Technical Summary: Repetition Mismatch in Data Mixture Experiments

Problem Statement

In the pre-training of language models (LMs), determining the optimal composition of data mixtures from multiple sources is critical for downstream performance. A common strategy involves running small-scale experiments and extrapolating the results to the target training budget. However, practitioners frequently observe that mixtures optimized at small scales fail to transfer effectively to larger regimes.

This paper identifies repetition mismatch as the primary cause of this extrapolation failure. When high-quality data is scarce relative to the training budget, it must be repeated multiple times during training. Crucially, the number of repetitions required for a fixed dataset changes as the total training budget grows. Standard scaling-based mixture selection methods assume that the relationship between mixture proportions and performance is invariant to scale. The authors argue this assumption is flawed because small-scale proxy experiments impose a fundamentally different repetition regime on high-quality data compared to the target run. This discrepancy distorts the loss landscape, shifting the apparent optimal mixture in ways that small-scale experiments cannot anticipate.

Methodology

The authors propose a repetition-aware subsampling procedure to isolate repetition as a variable in mixture prediction. This approach, building on work by Li et al. (2025), involves downsampling all data sources such that the high-quality data undergoes the exact same number of repetitions in the small-scale experiment as it would in the full-scale target run, while using only a fraction of the total tokens.

Experimental Setup:

• Datasets: The study combines limited high-quality datasets (WikiText or PubMed) with a large-scale web crawl corpus (FineWeb).
• Models: Experiments span four model sizes (30M, 124M, 345M, and 757M parameters) using a modified NanoGPT architecture with Muon optimizer and Rotary Positional Embeddings.
• Comparison: The proposed repetition-controlled approach is compared against a standard scaling-law baseline that extrapolates optimal mixture ratios from shorter training runs without matching repetition rates.
• Metrics: Performance is measured by the distance of the predicted mixture from the true optimal mixture (measured in mixing ratio error) and validation loss on held-out high-quality domains.

Key Contributions and Results

1. Repetition Mismatch is a Dominant Confounder

The paper demonstrates that matching the repetition rate of the target run is sufficient to recover accurate mixture predictions from small-scale experiments.

• Two-Source Results: In a setting with one high-quality source and FineWeb, a single repetition-controlled experiment using only 1/16 of the target tokens recovers a mixture within 0.05 of the optimum for a 757M parameter model. In contrast, the standard scaling-law approach without repetition control yields an error of 0.75.
• Cost Efficiency: Achieving comparable accuracy without repetition control requires three to four training horizons, consuming 44% to 94% of the target token budget. The repetition-controlled method achieves this with a single horizon.
• Model Capacity Interaction: The benefit of repetition control scales with model size. The error reduction is most pronounced for larger models (757M), while smaller models (30M) show less benefit or even outperform the repetition-controlled method with standard scaling laws, suggesting a lower bound on model scale where this method is most effective.

2. Effectiveness in Multi-Source Settings

The authors extend the findings to a three-source setting (WikiText, PubMed, and FineWeb).

• While the larger mixture space requires more experiments to constrain, repetition control remains highly effective.
• At the 757M scale, two repetition-controlled horizons are sufficient to recover the target optimal mixture, outperforming baselines that require full two-source experiments to construct.
• At the 124M scale, multi-horizon repetition-controlled predictions (up to four horizons) converge to the optimum, effectively matching the true loss (2.91950 vs. 2.91820) while significantly reducing the computational cost compared to non-controlled baselines.

3. Repetition as a First-Class Variable

The results indicate that data repetition dynamics, not just scale, determine whether small-scale mixture experiments generalize. The paper argues that repetition should be treated as an explicit variable in mixture optimization rather than an incidental side effect of limited data.

Significance and Claims

The paper claims that controlling for repetition mismatch resolves the extrapolation problem in data-constrained regimes. By preserving the repetition dynamics of the target run through subsampling, researchers can predict optimal data mixtures with high accuracy using a fraction of the computational budget.

The authors position this as a simple, orthogonal intervention that requires no parametric modeling or proxy training runs, making it applicable to any existing mixture prediction method. They assert that as practitioners increasingly rely on smaller ablation runs to inform mixture decisions for billion-parameter models, accounting for repetition mismatch will become critical to avoiding systematic prediction errors that grow with model capacity and the gap between proxy and target scales.

The study is modest in its scope, noting that experiments were conducted up to 757M parameters and ~3.8B tokens, smaller than modern LLM pre-training. However, the consistent trend of repetition control becoming more effective with model size suggests these findings will likely carry over to larger scales. The authors also note that their evaluation focuses on validation loss over high-quality domains, and confirming transfer to downstream benchmarks remains a direction for future work.