Data-Constrained Language Model Pretraining: Improved Regularization and Scaling Laws

Abstract: Classical scaling laws for LLM pretraining balance model size against training dataset size under a fixed compute budget, assuming abundant data and a single pass over the corpus. As training compute grows faster than the supply of natural language data, pretraining is likely to enter a data-constrained, compute-rich regime where models train for multiple epochs over a finite dataset. We study data-constrained pretraining along two axes, regularization and scaling. For regularization, we study masked-input regularization (MIR), an auxiliary next-token prediction loss on randomly masked inputs. MIR tests whether the random masking central to diffusion LLMs can benefit autoregressive pretraining without architectural changes or inference overhead. Across 72M to 1.4B parameter models, we find that MIR added on top of strong weight decay improves validation loss over autoregressive strong-weight-decay-only models, with downstream gains at 1.4B. For scaling, we propose SoftQ, a scaling law that couples model size and data size to capture their interaction under repeated data. Classical alternatives such as the Chinchilla law use an additive form that decouples these terms, making them misspecified in the data-constrained regime. We find that SoftQ fits data-constrained experiments substantially better than these alternatives, and estimates MIR's gains as equivalent to roughly 1.3 times as much unique training data. We release our code at https://github.com/yixinw-lab/dc_pretrain.

• The paper demonstrates that applying Masked-Input Regularization (MIR) atop strong weight decay markedly reduces validation loss, especially as model size increases.
• The study reveals that traditional additive scaling laws fall short in data-limited regimes, leading to the proposal of the coupled SoftQ scaling law for improved predictive fidelity.
• Empirical findings show MIR enhances data efficiency by up to 1.3× and yields better downstream task performance, presenting a practical approach for finite corpus training.

Expert Summary of "Data-Constrained LLM Pretraining: Improved Regularization and Scaling Laws" (2606.06888)

Introduction and Motivation

The paper conducts a systematic study of LLM (LM) pretraining under data-constrained, compute-rich conditions, addressing central questions of regularization and scaling law formulation in this emerging practical regime. Classical scaling laws (e.g., Chinchilla [hoffmann2022training], Kaplan [kaplan2020scaling]) historically assume that available training data is abundant relative to computation, justifying a single-pass training paradigm where only compute budget constrains optimal tradeoffs between model size and data exposure. Recent evidence shows compute availability is outpacing new data acquisition [villalobos2022], necessitating study of multi-epoch training over finite corpora, where overfitting and memorization risks increase.

This work targets two orthogonal axes:

1. Regularization: How can neural LMs be robust to overfitting in repeated epochs over finite data? Is strong weight decay optimal or can random masking (as in masked diffusion LMs) further regularize autoregressive (AR) training, even atop aggressive weight decay?
2. Scaling Laws: Are additive loss decompositions (as in Chinchilla) still adequate, or does the interaction between model- and data-limited regimes call for new coupled formulations?

Regularization in Data-Constrained Pretraining

The authors revisit empirical claims that discrete masked diffusion LLMs (dLLMs) outperform AR LMs under data-constrained, high-epoch training. They demonstrate that with sufficiently strong weight decay, the purported dLLM-AR gap nearly disappears across scales, i.e., prior dLLM advantages are largely ablated when AR regularization is optimized (Figure 1).

Figure 1: MIR improves validation loss over strongly regularized AR baselines, especially on larger models.

To probe for additive effects, the authors introduce Masked-Input Regularization (MIR): an auxiliary NTP objective applied to versions of training sequences with randomly masked tokens, added atop the standard clean AR NTP loss (with weighting coefficient $\lambda$). MIR is architecturally minimal, requiring only one [MASK] token addition and double forward passes at training, with no changes at inference.

Empirically (across scales 72M–1.4B and both natural language and code data), MIR yields consistent validation loss reductions on top of strong weight decay, with gains increasing as model size grows—from 0.006 at 72M to ~0.03 at 1.4B. Importantly, loss improvements are not restricted to rare or outlier tokens: token-level analysis at 1.4B finds the MIR advantage is broad, localizing especially to contexts with rare prefixes, noisy text, or mixed scripts (Figure 2).

Figure 2: Left, MIR produces a heavier positive tail in absolute loss-gap analysis; right, MIR boosts next-token probability particularly in contexts with noisy or unusual prefixes.

Theoretical underpinnings are developed via a controlled toy model, formalizing that in the compute-rich, data-limited regime, memorization on context-specific components leads to overfitting in the AR objective, while MIR regularizes toward generalizable (context-invariant) patterns. The prove that MIR consistently improves generalization as model size grows, with gain growing in model capacity. This is corroborated by empirical results: downstream, MIR-trained models score up to +10.2 on BoolQ versus the strong AR baseline, demonstrating that improved validation loss directionally transfers to challenging tasks.

Scaling Laws: Additive Decoupling Is Misspecified

The study exposes a central limitation in the Chinchilla loss model, which uses an additive separation:

$$L_{\mathrm{Ch}}(N, U) = E + AN^{-\alpha} + B U^{-\beta}$$

In this form, the loss gap between two data budgets (e.g., $U_1$ vs. $U_2$) is independent of model size. However, empirical data reveal this gap grows with $N$: for small LMs, additional data provides marginal utility; as $N$ increases and model becomes less capacity-limited, the same absolute increment in data increasingly impacts valid loss (Figure 3).

Figure 3: Regularized baseline: loss difference between unique data budgets grows with model size, contradicting Chinchilla's prediction of constancy.

To address this, the authors propose the SoftQ scaling law, which smoothly couples both axes using a five-parameter soft-bottleneck formulation inspired by the "quanta hypothesis" [quanta-scalinglaw]:

$$L_{\mathrm{SoftQ}}(N, U) = E + \left(A N^{-\rho} + B U^{- \rho / (1+\alpha)} \right)^{\alpha / \rho}$$

This law interpolates between parameter- and data-limited regimes, with the sharpness governed by $\rho$. Fitting is conducted on a dense baseline grid (5 model sizes × 4 data budgets), compared also to Muennighoff's effective-resource law and the Quanta law. On full-data fit, held-out extrapolation, and external baselines, SoftQ uniformly minimizes RMSE, MAE, and AIC (Figure 4, 7).

Figure 4: Fitted curves for Chinchilla (left) and SoftQ (right); only SoftQ correctly captures the growing gap across data budgets as a function of model size.

Figure 5: MIR and SoftQ: across model sizes and budgets, MIR shifts the scaling curve, conferring the same loss as a $1.3\times$ increase in unique data.

Quantitative Data Efficiency of MIR

Using the baseline SoftQ law (fitted on AR models with aggressive weight decay) as a reference, the loss asymptotes of MIR-trained models are compared to the baseline's infinite-$N$ curve. The unique data efficiency of MIR is then measured as the factor by which the baseline's data budget would need to be increased to match MIR's minimum attainable loss. Across language and code tasks, the MIR approach yields $$L_{\mathrm{Ch}}(N, U) = E + AN^{-\alpha} + B U^{-\beta}$$0 greater data efficiency at 200M–400M budgets, a nontrivial margin in practical high-$$L_{\mathrm{Ch}}(N, U) = E + AN^{-\alpha} + B U^{-\beta}$$1 pretraining.

Practical and Theoretical Implications

• Regularization: MIR is an efficient, general-purpose regularizer for the data-limited regime, with negligible inference overhead and easy adoption in AR pipelines. Its increased effectiveness with scale makes it highly relevant as practitioners push toward ever-larger LMs on fixed corpora.
• Scaling Law Modeling: The consistent empirical failure of additive decoupling, and SoftQ's robust fit, suggests future research into scaling law design should prioritize coupled forms, especially as diminishing returns from data repetition and large models interact.
• Benchmarking: Since compute is increasingly "free" relative to new data, research should prioritize data efficiency (loss at fixed unique data) as a primary metric, superseding compute efficiency for this regime.

Limitations and Future Directions

Experiments are constrained to 1.4B–1.5B parameter models and $$L_{\mathrm{Ch}}(N, U) = E + AN^{-\alpha} + B U^{-\beta}$$2M tokens, orders of magnitude below commercial or frontier-scale LMs. All training assumes standard Llama-style architectures and AdamW; exploring other optimizers, layer norms, attention variants, or curriculum scheduling under MIR may yield additive benefits. While MIR is shown here as an auxiliary loss, hybrid objectives or data-conditional masking could further refine regularization schedules.

The heavy per-cell hyperparameter search is likely impractical at full scale; transfer or automated hyperparameter routines warrant development to operationalize MIR and SoftQ scaling in the broader ecosystem.

Conclusion

This work establishes that in the data-constrained, compute-rich regime, (1) random masking (MIR) offers robust additional regularization atop large weight decay, with effects intensified at large model scales and (2) classical additive scaling laws are misspecified, with the SoftQ law providing superior descriptive and predictive fidelity for validation loss surfaces. MIR confers significant data efficiency improvements—up to $$L_{\mathrm{Ch}}(N, U) = E + AN^{-\alpha} + B U^{-\beta}$$3 at large budgets—highlighting an actionable path to better utilization of finite corpora in next-generation LMs.