To Memorize or to Retrieve: Scaling Laws for RAG-Considerate Pretraining

1. Problem Statement

Current scaling laws for Large Language Models (LLMs) treat the training corpus as a monolithic resource, assuming all data is consumed parametrically (compressed into model weights). However, Retrieval-Augmented Generation (RAG) introduces a new degree of freedom: a portion of the data can be held out as an external datastore and accessed at inference time.

The core problem addressed is the resource allocation trade-off under a fixed data budget. Given a corpus of $N$ tokens, how should one optimally split the data between:

1. Pretraining Data ($D$): To build parametric knowledge (internal weights).
2. Retrieval Store ($R$): To build non-parametric knowledge (external index).

Prior work has studied scaling laws for pretraining and the benefits of RAG separately, but no systematic analysis exists on how these two mechanisms interact or compete for the same data budget across different model scales and task types.

2. Methodology

Experimental Setup

• Models: The authors trained OLMo-2 based LMs ranging from 30M to 3B parameters.
• Data: Used 100B tokens of the DCLM (DataComp-LM) corpus.
• Variable Scales:
  • Pretraining Data ($D$): Varied from 1x to 150x the number of model parameters.
  • Retrieval Store ($R$): Varied from 1x to 20x the number of parameters, constructed from held-out slices of the same DCLM corpus.
• Index Construction: Used FAISS with IVPQ indexing. Embeddings were generated using Qwen3-Embedding-8B. The retrieval store was built from disjoint subsets of the corpus to ensure controlled scaling comparisons (smaller stores are strict subsets of larger ones).
• Evaluation: Evaluated across a diverse suite of benchmarks including reasoning (PIQA, StrategyQA), scientific QA (SciQ, OpenBookQA), and open-domain QA (Natural Questions, SimpleQA).
• Metrics: Primary metric is Perplexity (PPL) rather than accuracy. PPL provides a continuous, length-normalized measure that captures incremental improvements better than discrete accuracy, which often exhibits step-like behavior.

Modeling Approach

The authors propose a 3D Scaling Law that models loss $L$ as a function of Model Size ($N$), Pretraining Tokens ($D$), and Retrieval Tokens ($R$).

They extend the standard power-law scaling (Kaplan et al., 2020; Hoffmann et al., 2022) by adding a logarithmic gain term for retrieval:
$$L(N, D, R) = A \left(\frac{N}{10^9}\right)^{-\alpha} + B \left(\frac{D}{10^9}\right)^{-\beta} - C \log\left(1 + \eta \frac{R}{10^9}\right) + L_0$$

Where:

• $\alpha, \beta$: Scaling exponents for model size and data.
• $C, \eta$: Parameters governing the magnitude and saturation rate of retrieval gains.
• $L_0$: Irreducible loss floor.

3. Key Contributions

1. Unified Scaling Framework: Introduced a 3D scaling manifold that models performance as a joint function of model size, pretraining data, and retrieval corpus size.
2. Quantification of Substitutability: Defined a metric ($\sigma$) representing the number of pretraining tokens saved per retrieval token. This allows for calculating the "equivalent pretraining budget" a retrieval store provides.
3. Identification of Crossover Points: Discovered a scale-dependent threshold where retrieval becomes an efficient substitute for pretraining.
4. Task-Dependent Analysis: Demonstrated that the utility of retrieval is not uniform; it varies significantly based on task type (factual vs. reasoning) and model saturation.

4. Key Results

A. Retrieval as a Substitute for Pretraining

• Crossover Threshold: Retrieval is ineffective at replacing pretraining in low-data regimes. However, beyond a threshold of approximately 4.14 pretraining tokens per parameter ($D/N \approx 4.14$), retrieval becomes a highly efficient substitute.
• Log-Linear Gains: In the saturated regime, each retrieval token can replace multiple pretraining tokens, with gains growing log-linearly.
• Diminishing Returns: While retrieval remains an efficient substitute at larger scales, the absolute marginal benefit (loss reduction per unit of retrieval data) decreases as models grow larger and become more saturated.

B. Scale and Task Dependence

• Small Models Benefit Most: Smaller models (e.g., 30M) derive the largest marginal benefit from retrieval. As model size increases (up to 3B), the marginal utility of retrieval diminishes.
• Task Sensitivity:
  • Knowledge-Intensive Tasks: Tasks like CommonsenseQA and SimpleQA show high substitutability ($\sigma$), meaning retrieval effectively replaces the need for memorizing facts.
  • Reasoning-Heavy Tasks: Tasks like PIQA and StrategyQA show weak or negative gains from retrieval. The authors suggest these tasks rely on internal reasoning capabilities rather than external fact retrieval.

C. Retrieval Quality

• The study tested query formulations (Question-only vs. Question + Gold Answer). While better retrieval (oracle-style) improves performance, it does not fundamentally alter the scaling trends. This indicates that the observed trade-offs are driven by the allocation of data resources rather than just retrieval precision.

5. Significance and Implications

• Rethinking Data Allocation: The paper argues against the assumption that all available data should be compressed into weights. Instead, it suggests a joint optimization of pretraining and retrieval under a fixed compute/data budget.
• Design Guidelines:
  • For smaller or undertrained models, allocating a significant portion of the data budget to a retrieval store is highly efficient and can substitute for massive pretraining.
  • For large, saturated models, the marginal return on retrieval decreases, suggesting a shift back toward parametric learning for reasoning-heavy tasks.
• Cognitive Alignment: The findings mirror human cognition, where abstract reasoning is internalized (parametric) while factual recall is offloaded to external memory (non-parametric).
• Future Systems: This work provides a quantitative foundation for designing "RAG-considerate" pretraining pipelines, where corpora are explicitly partitioned into data for internalization vs. data for external access, optimizing for both training cost and inference efficiency.

In summary, the paper establishes that retrieval and pretraining are not independent design choices but competing mechanisms for knowledge storage. The optimal allocation depends on the model's scale, the specific task requirements, and the degree of pretraining saturation.