홈 › Papers

MLKV: Efficiently Scaling up Large Embedding Model Training with Disk-based Key-Value Storage

2025-04-02 · Yongjun He, Roger Waleffe, Zhichao Han, Johnu George, Binhang Yuan, Zitao Zhang, Yinan Shan, Yang Zhao, Debojyoti Dutta, Theodoros Rekatsinas, Ce Zhang

Many modern machine learning (ML) methods rely on embedding models to learn vector representations (embeddings) for a set of entities (embedding tables). As increasingly diverse ML applications utilize embedding models and embedding tables continue to grow in size and number, there has been a surge in the ad-hoc development of specialized frameworks targeted to train large embedding models for specific tasks. Although the scalability issues that arise in different embedding model training tasks are similar, each of these frameworks independently reinvents and customizes storage components for specific tasks, leading to substantial duplicated engineering efforts in both development and deployment. This paper presents MLKV, an efficient, extensible, and reusable data storage framework designed to address the scalability challenges in embedding model training, specifically data stall and staleness. MLKV augments disk-based key-value storage by democratizing optimizations that were previously exclusive to individual specialized frameworks and provides easy-to-use interfaces for embedding model training tasks. Extensive experiments on open-source workloads, as well as applications in eBay's payment transaction risk detection and seller payment risk detection, show that MLKV outperforms offloading strategies built on top of industrial-strength key-value stores by 1.6-12.6x. MLKV is open-source at https://github.com/llm-db/MLKV.

📄 PDF Abstract BibTeX arXiv:2504.01506

Code (1)

llm-db/mlkv 공식 구현

Methods 이 논문이 사용한 방법론

SET Dynamic Sparse Training method where weight mask is updated randomly periodically

MLKV: Multi-Layer Key-Value Heads for Memory Efficient Transformer Decoding

2024-06-13 · Zayd Muhammad Kawakibi Zuhri, Muhammad Farid Adilazuarda, Ayu Purwarianti, Alham Fikri Aji

Auto-regressive inference of transformers benefit greatly from Key-Value (KV) caching, but can lead to major memory bottlenecks as model size, batch size, and sequence length grow at scale. We introduce Multi-Layer Key-V…

MoME: Mixture-of-Memory Embeddings for Context-Aware Sparse Lookup

2026-09-14 · Muchen Li, Leonid Sigal, Renjie Liao hf

Scaling large language models efficiently has motivated sparse capacity mechanisms such as Mixture-of-Experts and, more recently, conditional memory: token-indexed embedding tables that augment the backbone with cheap pa…

Mitigating Position Bias in Transformers via Layer-Specific Positional Embedding Scaling

2026-06-26 · Changze Lv, Zhenghua Wang, Yiran Ding, Yixin Wu 외 arxiv

Large Language Models (LLMs) still struggle with the ``lost-in-the-middle'' problem, where critical information located in the middle of long-context inputs is often underrepresented or lost. While existing methods attem…

Little Giants: Synthesizing High-Quality Embedding Data at Scale

2024-10-24 · Haonan Chen, Liang Wang, Nan Yang, Yutao Zhu 외

Synthetic data generation has become an increasingly popular way of training models without the need for large, manually labeled datasets. For tasks like text embedding, synthetic data offers diverse and scalable trainin…

Synthetic Data Generation

Efficiently Enhancing Zero-Shot Performance of Instruction Following Model via Retrieval of Soft Prompt

2022-10-06 · Seonghyeon Ye, Joel Jang, Doyoung Kim, Yongrae Jo 외

Enhancing the zero-shot performance of instruction-following models requires heavy computation, either by scaling the total number of training datasets or the model size. In this work, we explore how retrieval of soft pr…

Instruction Following Retrieval