arXiv:2402.07440 · cs.IR · 2024
Benchmarking and Building Long-Context Retrieval Models with LoCo and M2-BERT
Jon Saad-Falcon, Daniel Y. Fu, Simran Arora, Neel Guha, Christopher Ré
29 citations · 2 influential citations— Semantic Scholar
TL;DR · Semantic Scholar
This work introduces LoCoV1, a novel 12 task benchmark constructed to measure long-context retrieval where chunking is not possible or not effective, and presents the M2-BERT retrieval encoder, an 80M parameter state-space encoder model built from the Monarch Mixer architecture, capable of scaling to documents up to 32K tokens long.
Saved summary
AI-generated while a reader studied this paper. Check the original source for accuracy.
Problem: Existing retrieval pipelines perform poorly on long documents of 10\text{K} tokens or more because current benchmarks fail to evaluate situations requiring synthesis of information across the entire text. Furthermore, pretraining base models to represent both short query contexts and long document contexts is difficult, and standard contrastive finetuning methods like multiple negatives ranking loss (MNRL) require large batch sizes that are impossible to run under GPU memory constraints. Thus, developers face challenges in evaluating, pretraining, and finetuning models for long-context retrieval.
Method: To address these challenges, the authors introduce LoCoV1, a 12-task long-context retrieval benchmark spanning multiple domains where chunking is ineffective. They present the M2-BERT retrieval encoder, an 80\text{M} parameter state-space model built from the Monarch Mixer architecture that scales to 32\text{K} tokens. To pretrain this model, they employ a data mixture of short and long context sequences, and to finetune it, they utilize the batch-independent orthogonal projection loss (OPL) with single-sample batches.
Results: On the LoCoV1 benchmark, M2-BERT-32k outperforms the next best truncation baseline, E5-Mistral, by an average of 23.3 points and fine-tuned BGE-Large by 29.9 points. It also beats several popular API services, including OpenAI's text-embedding-ada-002, by 35.4 points on average despite containing upwards of 90 fewer parameters. Finally, M2-BERT-32K is 3 to 676 times more efficient at generating embeddings for long documents than E5-Mistral.
Limitations: While M2-BERT-128 performs comparably to SentenceBERT on short-context benchmarks, it averages 1.3 points lower on the BEIR benchmark. Additionally, the paper notes that while orthogonal projection loss (OPL) was effective, other loss functions with similar properties remain to be explored. Finally, pretraining models at a sequence length of 32\text{K} requires warm-starting from pretrained 8\text{K} checkpoints, as random initialization fails to converge within a reasonable time.
Abstract
Retrieval pipelines-an integral component of many machine learning systems-perform poorly in domains where documents are long (e.g., 10K tokens or more) and where identifying the relevant document requires synthesizing information across the entire text. Developing long-context retrieval encoders suitable for these domains raises three challenges: (1) how to evaluate long-context retrieval performance, (2) how to pretrain a base language model to represent both short contexts (corresponding to queries) and long contexts (corresponding to documents), and (3) how to fine-tune this model for retrieval under the batch size limitations imposed by GPU memory constraints. To address these challenges, we first introduce LoCoV1, a novel 12 task benchmark constructed to measure long-context retrieval where chunking is not possible or not effective. We next present the M2-BERT retrieval encoder, an 80M parameter state-space encoder model built from the Monarch Mixer architecture, capable of scaling to documents up to 32K tokens long. We describe a pretraining data mixture which allows this encoder to process both short and long context sequences, and a finetuning approach that adapts this base model to retrieval with only single-sample batches. Finally, we validate the M2-BERT retrieval encoder on LoCoV1, finding that it outperforms competitive Transformer-based models by at least 23.3 points, despite containing upwards of 90x fewer parameters.
Source: https://arxiv.org/abs/2402.07440 — Thank you to arXiv for use of its open access interoperability.
Papers citing this work · Semantic Scholar
- STAIR (STructure Aware Information Retriever): A novel dataset and LLM based retriever for document structure augmentation (2026) · 0 citations
- REIGN: Refurbished Embeddings with Integrated Guidance Networks for Efficient Context-Length Scaling (2026) · 0 citations
- CRAwLeR - Cross-Reference Aware Legal Retrieval (2026) · 0 citations
- EASE-TTT: Evidence-Aligned Selective Test-Time Training for Long-Context Question Answering (2026) · 0 citations
- RubiConv - Efficient Boundary-Respecting Convolutions (2026) · 0 citations
Read the original paper, explore key passages and ask questions in AewolFlow.