Research map / Large language models / Language, safety and interpretability

Scaling Retrieval-Based Language Models with a Trillion-Token Datastore

NeurIPS 2024 · Large language models › Language, safety and interpretability · approach: LLM pruning and scaling data · 8 citations (OpenAlex)

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.

Builds MassiveDS, a 1.4-trillion-token open datastore, and an efficient subsampling pipeline to study how inference-time datastore scale improves retrieval-in-context LMs on language modeling and QA/reasoning tasks.

Problem

Scaling laws usually consider only training data and parameters; it is unknown how the size of an inference-time retrieval datastore affects retrieval-based LMs across broad tasks, since prior trillion-token datastores were proprietary or evaluated narrowly.

Approach

Treat inference-time datastore size as a scaling dimension: construct an open, multi-domain 1.4T-token datastore and reorder the pipeline (index and retrieve top-K once, then subsample, dedup, decontaminate and rerank per variant) so that datastore scaling studies are ~10x cheaper while equivalent to the naive pipeline with high probability.

Key results

• Datastore scaling monotonically reduces perplexity without saturation, and Llama-2 7B with retrieval from MassiveDS outperforms the LM-only Llama-2 13B on language modeling.
• On knowledge-intensive QA, a Llama-2 7B retrieving from under 100B tokens outperforms the LM-only Llama-3 8B on TriviaQA and NQ.
• Under equal training compute, retrieval-based Pythia and OLMo models achieve better compute-optimal performance than LM-only models; the pipeline reduces datastore construction compute by more than an order of magnitude.

Measured effects reported in the paper

Surprising findings

• Retrieval-based Pythia models benefit on TriviaQA and NQ but show no meaningful gain on MMLU and MedQA, which stay near random even at 12B.
• Llama-3 8B underperforms Llama-2 7B on RedPajama perplexity.

Limitations

• Compute-optimal analysis relies on intermediate checkpoints of Pythia and OLMo as proxies for models trained on different token counts.
• Only a single retriever (Contriever-MSMARCO) was used for the full scaling study, since changing it requires re-indexing.
• MassiveDS lacks high-quality reasoning-heavy and in-domain data, so gains on MMLU and MedQA are limited; evaluation covers mostly QA tasks.

Experimental setup

Datasets: MassiveDS (RedPajama Books, peS2o, RPJ ArXiv, Wikipedia, StackExchange, GitHub, OpenWebMath, NaturalProofs, PubMed, Comm RedPajama S2ORC TriviaQA Natural Questions MMLU MedQA

Benchmarks: TriviaQA Natural Questions MMLU MedQA RedPajama perplexity S2ORC perplexity

Metrics: accuracy perplexity

Baselines: LM-only models RETRO Sphere single-domain datastores (Wikipedia, PubMed, peS2o, RedPajama subsets) cross-encoder reranker (MiniLM-L12)

Builds on: RETRO RALM (retrieval-in-context) Contriever Pythia scaling suite OLMo

Closest papers by meaning

• Data Engineering for Scaling Language Models to 128K Context (ICML 2024 · similarity 0.89)
• Scaling Laws Meet Model Architecture: Toward Inference-Efficient LLMs (ICLR 2026 · similarity 0.89)
• LLMs on the Line: Data Determines Loss-to-Loss Scaling Laws (ICML 2025 · similarity 0.89)
• Scaling Laws with Vocabulary: Larger Models Deserve Larger Vocabularies (NeurIPS 2024 · similarity 0.88)
• Inference Scaling for Long-Context Retrieval Augmented Generation (ICLR 2025 · similarity 0.88)
• Scaling Inference-Efficient Language Models (ICML 2025 · similarity 0.88)