Research map / Architectures and efficiency / Attention and transformers
Benchmarking and Building Long-Context Retrieval Models with LoCo and M2-BERT
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Introduces LoCoV1, a 12-task long-context retrieval benchmark, and M2-BERT, an 80M-parameter Monarch Mixer state-space retrieval encoder scaling to 32K tokens, finetuned with orthogonal projection loss at batch size 1.
Problem
Retrieval encoders perform poorly on long documents (10K+ tokens) where relevance requires synthesizing information across the full text; existing benchmarks can be solved by truncation or chunking, and contrastive MNRL finetuning needs large batches that long-context GPU memory cannot afford.
Approach
Combine a Monarch Mixer (subquadratic state-space) BERT-like encoder, a mixed short/long pretraining corpus with warm-started long-context stages, and orthogonal projection loss, which enables single-sample (B=1) finetuning for retrieval, plus a new long-context benchmark (LoCoV1) where chunking fails.
Key results
- M2-BERT-32k (80M) scores 94.7 average nDCG@10 on LoCoV1 versus 71.4 for E5-Mistral (7.11B), a 23.3-point gain under truncation.
- OPL with batch size 1 reaches 94.7 on LoCoV1 versus 70.4 for MNRL at batch size 2 when finetuning M2-BERT-32k.
- M2-BERT-32k embeds 32K-token documents 676x faster than E5-Mistral on A100, while M2-BERT-128 is within 1.3 nDCG@10 of SentenceBERT on BEIR.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Increase maximum sequence length from 128 to 32768 tokens | M2-BERT · LoCoV1 | nDCG@10: 70.3 -> 94.7 | Helped |
| Pretrain on mixed short and long examples rather than short examples only (ablation) | M2-BERT · LoCoV1 | nDCG@10: 37.2 -> 55.4 | Helped |
| Pretrain on mixed short and long examples rather than long examples only (ablation) | M2-BERT · LoCoV1 | nDCG@10: 44.9 -> 55.4 | Helped |
| Use an 8k pretrained checkpoint to warm-start 32k pretraining rather than random initialization (ablation) | M2-BERT-32k · M2-BERT-32k pretraining | MLM accuracy: 4.8 -> 33.9 | Helped |
| Use orthogonal projection loss instead of MNRL | M2-BERT-32k · LoCoV1 | nDCG@10: 70.4 -> 94.7 | Helped |
| Use orthogonal projection loss instead of prototype loss (ablation) | M2-BERT-32k · LoCoV1 | nDCG@10: 63.2 -> 94.7 | Helped |
| Use chunked embeddings instead of truncation for Jina Embeddings (ablation) | Jina Embeddings · LoCoV1 | nDCG@10: 67.2 -> 19.2 | Hurt |
| Use chunked embeddings instead of truncation for BGE-Large Fine-tuned (ablation) | BGE-Large Fine-tuned · LoCoV1 | nDCG@10: 64.8 -> 61.6 | Hurt |
| Use chunked embeddings instead of truncation for E5-Mistral (ablation) | E5-Mistral · LoCoV1 | nDCG@10: 71.4 -> 70.3 | Hurt |
| Use chunked embeddings instead of truncation for OpenAI Ada (ablation) | OpenAI Ada · LoCoV1 | nDCG@10: 63.2 -> 63.4 | No clear effect |
| Increase M2-BERT maximum sequence length from 128 to 32768 tokens | M2-BERT · BEIR | nDCG@10: 38.7 vs. 40.0 | No clear effect |
| Use M2-BERT-128 rather than SentenceBERT | M2-BERT-128 · MTEB | average score: 63.1 -> 63.3 | No clear effect |
Surprising findings
- Chunking with embedding averaging barely helps and often hurts baseline retrievers (e.g., Jina Embeddings drops from 67.2 to 19.2 average on LoCoV1 with chunks).
- Prototype loss from a short-context teacher underperforms MNRL for long-context finetuning (63.2 vs 70.4).
Limitations
- Baselines with max length below document length are evaluated with truncation or chunking, which are not optimal; comparisons are sensitive to these choices.
- M2-BERT is weaker on some short-context BEIR datasets (e.g., SciFact 39.9 vs 56.9 for SentenceBERT).
- Pretraining data is relatively small and does not match large-scale retriever pretraining used by baselines.
Experimental setup
Datasets: C4 Wikipedia BookCorpus SummScreenFD Gov. Report QMSUM QASPER MultiFieldQA 2WikiMQA Passage Retrieval CourtListener Australian Legal Case Reports
Benchmarks: LoCoV1 BEIR MTEB Needle-in-the-Haystack (NQ-based synthetic)
Metrics: nDCG@10 throughput / embedding time
Baselines: BGE-Large E5-Mistral Jina Embeddings OpenAI text-embedding-ada-002 ColBERTv2 BM25 SentenceBERT Voyage-001 Cohere embed-english-v3.0
Builds on: Monarch Mixer multiple negatives ranking loss orthogonal projection loss BEIR
Closest papers by meaning
- Retrieval meets Long Context Large Language Models (ICLR 2024 · similarity 0.90)
- Can Language Models Actually Retrieve In-Context? Drowning in Documents at Million Token Scale (NeurIPS 2026 · similarity 0.89)
- BTR: Binary Token Representations for Efficient Retrieval Augmented Language Models (ICLR 2024 · similarity 0.88)
- Long-Context LLMs Meet RAG: Overcoming Challenges for Long Inputs in RAG (ICLR 2025 · similarity 0.87)
- Contextual Document Embeddings (ICLR 2025 · similarity 0.87)
- Training with “Paraphrasing the Original Text” Teaches LLM to Better Retrieve in Long-Context Tasks (AAAI 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.