MLKV：基于磁盘键值存储高效扩展大型嵌入模型训练
机器学习
2025-04-03 v1
摘要
许多现代机器学习（ML）方法依赖嵌入模型来学习实体（嵌入表）的向量表示。随着日益多样化的 ML 应用利用嵌入模型且嵌入表的规模和数量不断增长，针对特定任务训练大型嵌入模型的开发呈现出 ad-hoc 增长。尽管不同嵌入模型训练任务中出现的可扩展性问题相似，但这些框架各自独立地重新发明并定制针对特定任务的存储组件，导致在开发和部署方面存在大量重复的工程工作。本文提出了 MLKV，这是一个高效、可扩展且可重用的数据存储框架，旨在解决嵌入模型训练中的可扩展性挑战，特别是数据停滞和数据陈旧问题。MLKV 通过赋能基于磁盘的键值存储，使先前专属各个特定框架的优化变得民主化，并为嵌入模型训练任务提供易用的接口。针对开源工作负载以及 eBay 支付交易风险检测和卖家支付风险检测等应用进行了广泛实验，表明 MLKV 在基于工业级键值存储的卸载策略之上性能提升 1.6-12.6 倍。MLKV 已在 https://github.com/llm-db/MLKV 上开源。
引用
@article{arxiv.2504.01506,
title = {MLKV: Efficiently Scaling up Large Embedding Model Training with Disk-based Key-Value Storage},
author = {Yongjun He and Roger Waleffe and Zhichao Han and Johnu George and Binhang Yuan and Zitao Zhang and Yinan Shan and Yang Zhao and Debojyoti Dutta and Theodoros Rekatsinas and Ce Zhang},
journal= {arXiv preprint arXiv:2504.01506},
year = {2025}
}
备注
To appear in ICDE 2025