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

Summary (Overview)

• Core Problem: As training compute grows faster than the supply of natural language data, pretraining is entering a data-constrained, compute-rich regime where models train for multiple epochs over finite datasets, breaking assumptions of classical scaling laws.

• Key Contribution 1 - Regularization: Proposes Masked-Input Regularization (MIR), an auxiliary next-token prediction loss on randomly masked inputs, which improves validation loss on top of strong weight decay across 72M to 1.4B parameter models, with downstream gains including +10.2 points on BoolQ and +2.2 points on SciQ at 1.4B scale.

• Key Contribution 2 - Scaling Laws: Demonstrates that the additive Chinchilla scaling law is misspecified in the data-constrained regime (it predicts model-size-invariant loss gaps between data budgets, while empirically the gap grows with model size), and proposes SoftQ, a five-parameter coupled scaling law that substantially outperforms Chinchilla, Quanta, and Muennighoff-style alternatives.

• Key Finding: MIR's gain is estimated as equivalent to roughly 1.3× as much unique training data at 200M–400M token budgets.

• Key Insight: The previously reported autoregressive vs. masked diffusion language model gap is largely explained by insufficient AR regularization; when both are strongly regularized with large weight decay, their validation losses become comparable.

---

Introduction and Theoretical Foundation

Background: The Data-Constrained Regime

Classical scaling laws (Kaplan et al., 2020; Hoffmann et al., 2022) assume abundant data and a single pass over the corpus. However, training compute is growing faster than natural language data supply (Villalobos et al., 2024), creating a regime where:

• The unique dataset is fixed (bounded by $U$ tokens)
• Additional compute is spent on larger models and multiple epochs over the same corpus
• Overfitting becomes the central challenge

Problem Formulation

For a standard dense decoder-only transformer trained with next-token prediction, training compute is approximately:

$$C(N, D) \approx 6ND$$

where $N$ is the number of parameters and $D = U \cdot N_E$ is the total number of training tokens (with $N_E$ epochs over $U$ unique tokens).

The classical compute-optimal allocation problem is:

$$(N^{\star}(C), D^{\star}(C)) = \arg \min_{N, D} L_{\text{eval}}(N, D) \quad \text{s.t.} \quad C(N, D) = C$$

In the data-constrained regime, the goal shifts to modeling the optimized validation-loss envelope:

$$L^{\star}(N, U) = \inf_{h \in \mathcal{H}} L_{\text{eval}}(N, U; h)$$

where $h$ includes tunable hyperparameters (epochs, learning rate schedule, weight decay, regularization).

Theoretical Foundation: Masking as Regularization

The authors provide a theoretical intuition based on a context-specific noise model (formalized in Appendix C). The model decomposes each sequence into:

1. A context-specific component (enables memorization, acts as noise for generalization)
2. A generalizable component (contains predictive features)
3. An output token to be predicted

Theorem (Informal): Under the context-specific noise model in the data-constrained, compute-rich regime, standard autoregressive pretraining can minimize training loss by relying almost entirely on the context-specific component, thereby memorizing patterns that do not generalize. In contrast, MIR regularizes the model's dependence on context-specific components and encourages learning predictive patterns on generalizable components, strictly improving validation loss. Moreover, for fixed data size, this improvement increases as model capacity grows.

This motivates the hypothesis that larger overparameterized models benefit more from masking-based regularization.

---

Methodology

1. Masked-Input Regularization (MIR)

MIR is a minimal modification to standard autoregressive pretraining. For a clean sequence $x$ and a randomly masked version $\tilde{x}$:

$$\mathcal{L} = \mathcal{L}_{\text{NTP}}(x) + \lambda \mathcal{L}_{\text{NTP}}(\tilde{x})$$

Procedure:

• Sample a mask ratio $r \sim \text{Unif}(r_{\text{min}}, r_{\text{max}})$ for each input sequence
• At each position $t \in [0, T-1]$, a Bernoulli random variable with success probability $r$ determines whether to replace token $x_t$ with a special [MASK] token
• The model trains on both clean and masked inputs using standard next-token prediction

Key properties:

• No architectural changes required
• Preserves standard autoregressive decoding at inference
• Requires two forward passes per batch (increased training compute, but the regime is compute-rich)

2. Experimental Setup

Datasets:

• DCLM (Li et al., 2024) for natural language
• Stack-V2 (Lozhkov et al., 2024) for code-heavy text
• Both with 100M unique seed tokens, 10M reserved for validation

Model Scaling Ladder:

$$\text{ScalingLadder}(k) = (kW_1, kL_1, S_1, B_1)$$

where $W_1 = 1024$ (embedding dimension), $L_1 = 12$ (layers), $S_1 = 2048$ (sequence length), $B_1 = 128$ (batch size), and $k \in \{0.5, 0.75, 1, 1.5, 2\}$. Model sizes range from 72M to 1.4B parameters.

Hyperparameter tuning: Grid search over training steps, learning rate, and weight decay for each model (per-cell tuning following Kim et al., 2026a).

3. The SoftQ Scaling Law

Why Chinchilla is misspecified: The Chinchilla law decomposes loss as:

$$L_{\text{Ch}}(N, U) = E + \frac{A}{N^{\alpha}} + \frac{B}{U^{\beta}} \tag{1}$$

Its additive structure implies the loss gap between two data budgets is model-size-invariant:

$$L_{\text{Ch}}(N, U_1) - L_{\text{Ch}}(N, U_2) = \frac{B}{U_1^{\beta}} - \frac{B}{U_2^{\beta}}$$

This contradicts empirical observations where the gap grows with model size (Figure 1b).

SoftQ formulation (motivated by the skill-learning view of scaling laws):

$$L_{\text{SoftQ}}(N, U) = E + \left(\frac{A}{N^{\rho}} + \frac{B}{U^{\rho/(1+\alpha)}}\right)^{\alpha/\rho} \tag{4}$$

Properties:

• $\rho$ controls the sharpness of the transition between parameter-limited and data-limited regimes
• As $U \to \infty$: recovers parameter-scaling limit $L - E \propto N^{-\alpha}$
• As $N \to \infty$: recovers data-scaling limit $L - E \propto U^{-\alpha/(1+\alpha)}$
• Five fitted parameters $\{A, B, E, \alpha, \rho\}$ (same count as Chinchilla)
• When $\rho = 1$, SoftQ reduces to the Quanta law (strict nesting)

---

Empirical Validation / Results

Finding 1: Weight Decay Transfers Across AR and dLLM Pretraining

On DCLM with 100M unique tokens, four recipes were compared at 140M, 257M, and 664M:

Key finding: Large weight decay dramatically improves both AR and dLLM training. The previously reported AR–dLLM gap is largely explained by insufficient AR regularization.

Finding 2: MIR Improves Strongly Regularized AR Models

Validation loss improvements on DCLM 100M (averaged over 5 seeds):

• MIR improves validation loss at every model scale from 72M to 1.4B
• Average gain grows from ~0.006 loss at 72M to ~0.03 loss at 1.4B
• At 1.4B: mean validation loss reduced from 3.347 to 3.317

Transfer to code data (Stack-V2): MIR reduces validation loss at all five model sizes with gains from 0.008 to 0.020 loss (hyperparameters tuned only on DCLM).

Token-Level Analysis

• 6.61% of tokens satisfy $\Delta \ell_t \geq 1$ (MIR-better) vs. 5.41% with $\Delta \ell_t \leq -1$ (baseline-better)
• The gain comes from a broad set of hard validation tokens, not isolated outliers
• Top-gap examples: 62.6% word/subword continuations, 16.3% non-English/transliterated text, 11.6% punctuation
• The true token is often common (e.g., "and", "of", comma, closing parenthesis); difficulty comes from noisy prefix context (rare names, mixed scripts, broken word pieces)

Downstream Evaluations (1.4B models, DCLM U=100M)

MIR outperforms the baseline on six of eight zero-shot metrics.

Scaling Law Comparison

Grid: 5 model sizes × 4 data budgets = 20 points {72M, 140M, 257M, 664M, 1.4B} × {100M, 200M, 300M, 400M}.

Table 2: Scaling laws comparison (lower is better)

SoftQ wins on all three diagnostics: best in-sample fit, best held-out extrapolation to 400M, and best fit on the independent Kim et al. dataset.

MIR Data Efficiency Quantification

Using SoftQ as the baseline law, the regularized-baseline infinite-model curve is:

$$L_{\text{Reg},\infty}(U) = 0.306 + 2.249 U^{-0.125}$$

Solving $L_{\text{Reg},\infty}(U_{\text{eq}}) = E_U$ for each MIR asymptote $E_U$:

• At 200M–400M unique tokens, the regularized baseline needs about 1.28–1.34× as much unique data to match the MIR infinite-model asymptote
• SoftQ gives the most conservative data efficiency ratio at $U = 400$M among all scaling laws

---

Theoretical and Practical Implications

Theoretical Implications

1. Chinchilla's additive structure is fundamentally misspecified for data-constrained regimes. The coupling between model size and data size (captured by SoftQ's soft bottleneck) is essential for accurate loss prediction when data is the binding resource.

2. Masking provides complementary regularization to weight decay. The theoretical analysis (context-specific noise model) explains why: masking reduces the model's dependence on context-specific noise and encourages learning generalizable predictive features, with the benefit growing with model capacity.

3. The dLLM advantage over AR models is largely explained by regularization strength, not architectural differences. This reframes the comparison between autoregressive and diffusion language models in the data-constrained regime.

Practical Implications

1. MIR is a drop-in improvement for data-constrained pretraining: no architectural changes, no inference overhead, only an auxiliary loss during training. It provides consistent validation loss improvements and downstream gains.

2. Strong weight decay is critical and transfers across model families (AR and dLLM), confirming and extending prior findings (Kim et al., 2026a).

3. SoftQ enables better resource allocation in data-constrained settings, allowing practitioners to accurately predict the loss surface and estimate the value of additional unique data.

4. Data efficiency quantification: MIR is worth roughly 1.3× as much unique training data at 200M–400M token budgets, providing a concrete cost-benefit analysis for practitioners.

---

Conclusion

This paper addresses two linked questions in the data-constrained, compute-rich pretraining regime:

1. Regularization: The authors propose MIR (Masked-Input Regularization), which adds an auxiliary next-token prediction loss on randomly masked inputs to strongly regularized autoregressive training. MIR consistently improves validation loss across 72M to 1.4B parameters on both natural language and code data, with downstream gains including +10.2 points on BoolQ and +2.2 points on SciQ. The gain is estimated as equivalent to roughly 1.3× as much unique training data.

2. Scaling laws: The authors demonstrate that the additive Chinchilla law is misspecified in this regime and propose SoftQ, a five-parameter coupled scaling law that substantially outperforms Chinchilla, Quanta, and Muennighoff-style alternatives on in-sample fit, held-out extrapolation, and an independent dataset.

Limitations and future directions:

• Experiments span up to 1.4B parameters and 400M unique tokens—small relative to frontier-scale pretraining
• Model architecture and optimizer were held fixed; varying these could yield further gains
• The protocol relies on heavy per-cell hyperparameter search; a hyperparameter-transfer recipe for this regime is a natural next step

The code is released at https://github.com/yixinw-lab/dc_pretrain.