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

TL;DR

This paper investigates how to optimally balance pretraining data and retrieval store size in retrieval-augmented language models, providing a scaling framework that guides resource allocation for improved performance.

Contribution

It introduces a three-dimensional scaling framework modeling performance as a function of model size, pretraining tokens, and retrieval store size, guiding data resource allocation.

Findings

Retrieval consistently improves performance across model scales.

A scaling manifold models performance based on size, data, and retrieval.

Optimal data allocation depends on model scale, task, and pretraining saturation.

Abstract

Retrieval-augmented generation (RAG) improves language model (LM) performance by providing relevant context at test time for knowledge-intensive situations. However, the relationship between parametric knowledge acquired during pretraining and non-parametric knowledge accessed via retrieval remains poorly understood, especially under fixed data budgets. In this work, we systematically study the trade-off between pretraining corpus size and retrieval store size across a wide range of model and data scales. We train OLMo-2-based LMs ranging from 30M to 3B parameters on up to 100B tokens of DCLM data, while varying both pretraining data scale (1-150x the number of parameters) and retrieval store size (1-20x), and evaluate performance across a diverse suite of benchmarks spanning reasoning, scientific QA, and open-domain QA. We find that retrieval consistently improves performance over parametric-only baselines across model scales and introduce a three-dimensional scaling framework that models performance as a function of model size, pretraining tokens, and retrieval corpus size. This scaling manifold enables us to estimate optimal allocations of a fixed data budget between pretraining and retrieval, revealing that the marginal utility of retrieval depends strongly on model scale, task type, and the degree of pretraining saturation. Our results provide a quantitative foundation for understanding when and how retrieval should complement pretraining, offering practical guidance for allocating data resources in the design of scalable language modeling systems.