🤖 machine learning

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

This paper addresses the data-constrained regime of language model pretraining by introducing Masked-Input Regularization (MIR) to improve validation and downstream performance, and proposing SoftQ, a new scaling law that accurately captures the interaction between model and data sizes under repeated training epochs where classical laws fail.

Original authors: Zhiwei Xu, Shihao Wu, Hanseul Cho, Wei Hu, Yixin Wang

Published 2026-06-08

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Zhiwei Xu, Shihao Wu, Hanseul Cho, Wei Hu, Yixin Wang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a student (the AI model) to write stories. In the past, the rule was: "Give the student a massive library of unique books, and let them read each book exactly once." This worked well because the library was huge, and the student had plenty of time to read everything once.

But now, the situation has changed. We have built super-fast computers (compute) that can read millions of books in a second, but the supply of new, unique books (natural language data) is running out. We are in a "data-constrained, compute-rich" world. The student now has to read the same small library over and over again, multiple times, to fill their brain.

The problem? If you just make the student read the same books over and over, they start memorizing the books instead of learning how to write. They become a parrot, repeating exact phrases they've seen, rather than understanding the rules of language. This is called "overfitting."

This paper tackles two main questions:

How do we stop the student from just memorizing? (Regularization)
What is the best way to balance the size of the student's brain against the size of the library? (Scaling Laws)

Here is the breakdown of their findings using simple analogies.

1. The "Blurry Glasses" Trick (Masked-Input Regularization)

The Problem:
When a student reads the same book for the 10th time, they might start memorizing the exact font or a weird typo on page 5, rather than learning the story. They need a way to force themselves to focus on the meaning rather than the exact text.

The Solution (MIR):
The authors tried a technique called Masked-Input Regularization (MIR). Imagine giving the student a pair of "blurry glasses" or a "scrambler" while they read.

Every time they read a sentence, the computer randomly hides (masks) a few words, turning them into [MASK].
The student has to guess the missing words based on the context, just like a "fill-in-the-blank" game.
They do this alongside their normal reading.

The Result:
Think of it like practicing a sport. If you only practice on a perfect, flat field, you might fail when the wind blows. But if you practice with a "wind machine" (the random masking) blowing at you, you learn to adapt.

The paper found that even when the student was already being disciplined (using "strong weight decay," which is like a strict coach telling them not to memorize), adding this "blurry glasses" trick made them even better.
It helped the model handle tricky situations, like weird names, mixed languages, or broken text, much better than the model that just read the clean text over and over.
The Payoff: They calculated that using this trick is mathematically equivalent to giving the student 1.3 times more unique books to read. It's a way to get more value out of the limited library we have.

2. The "Coupled" Growth Rule (SoftQ Scaling Law)

The Old Rule (Chinchilla):
Previously, experts had a formula (the Chinchilla law) to decide how big a student's brain should be based on how many books they have.

The Old Idea: "If you double the number of books, you should double the brain size." It treated the brain size and the book count as two separate things added together.
The Flaw: This formula assumed you could always get new books. But in our "running out of books" world, this formula breaks. It predicted that the penalty for having fewer books would be the same whether the student had a tiny brain or a giant one.

The New Rule (SoftQ):
The authors realized that in a world with limited books, the brain size and the book count are tightly linked.

The Analogy: Imagine a bucket (the brain) and a hose (the data).
- If the bucket is tiny, it doesn't matter if the hose is huge; the bucket fills up fast, and extra water just spills over (wasted data).
- If the bucket is huge but the hose is tiny, the bucket stays mostly empty (wasted brain power).
- The "penalty" for having a small hose gets much worse as the bucket gets bigger. A giant bucket with a tiny hose is a disaster.

The Result:
They proposed a new formula called SoftQ.

Unlike the old formula, SoftQ understands that as the model gets bigger, the need for unique data grows faster. It "couples" the two together.
When they tested it, SoftQ predicted the model's performance much more accurately than the old Chinchilla formula, especially when the data was limited and the model was reading the same data many times.

Summary of the "Takeaways"

Don't just repeat; scramble: When you are forced to train on a limited dataset, randomly hiding parts of the text (MIR) forces the AI to learn general rules instead of memorizing specific sentences. It's like studying for a test by doing practice quizzes with the answers covered, rather than just reading the textbook.
Size matters more when data is scarce: The old rules for building AI said "bigger model = better, regardless of data." The new rules (SoftQ) say "if you make the model bigger, you need disproportionately more unique data, or the model will fail."
Efficiency: By using the "scramble" trick (MIR) and the new "coupled" formula (SoftQ), we can build better AI models even when we are running out of new text to train on.

What the paper didn't say:
The paper does not claim this will immediately fix medical diagnosis, create perfect chatbots for customer service, or solve climate change. It strictly focuses on the math of how to train these models more efficiently when data is scarce. The results are based on experiments with models up to 1.4 billion parameters (which is small compared to the massive models used by big tech companies today), but the principles are intended to apply to larger scales.

Technical Summary: Data-Constrained Language Model Pretraining

Problem Statement

Classical scaling laws for Large Language Model (LLM) pretraining, such as Chinchilla, operate under the assumption of an "abundant data" regime where unique training data can scale freely with compute. These laws optimize the allocation of a fixed compute budget between model parameters and training tokens, typically assuming a single pass over the corpus.

However, the paper argues that the field is shifting toward a data-constrained, compute-rich regime. In this setting, the supply of high-quality natural language data is fixed or bounded, while training compute continues to grow. Consequently, training strategies must adapt to utilize larger models and perform multiple epochs (repeated passes) over a finite dataset. The authors identify two critical gaps in the current understanding of this regime:

Regularization: How can models avoid overfitting when compute increases but unique data remains fixed? While strong weight decay is known to help, it is unclear if other mechanisms, such as the random masking used in Diffusion Language Models (dLLMs), offer complementary benefits.
Scaling Laws: Do classical additive scaling laws (which decouple model size and data size) accurately describe the loss surface in this regime? The authors hypothesize that these laws are misspecified because they fail to capture the interaction between model capacity and limited unique data.

Methodology

1. Regularization: Masked-Input Regularization (MIR)

To investigate regularization, the authors introduce Masked-Input Regularization (MIR). This is a minimal modification to standard autoregressive (AR) pretraining that does not alter the model architecture or inference procedure.

Mechanism: For a clean input sequence $x$ , a randomly masked version $\hat{x}$ is generated by replacing tokens with a [MASK] token based on a sampled mask ratio.
Objective: The training loss is augmented to include an auxiliary next-token prediction loss on the masked sequence:
$L = L_{NTP}(x) + \lambda L_{NTP}(\hat{x})$
Rationale: The authors hypothesize that random masking acts as a regularizer, forcing the model to learn generalizable predictive features rather than memorizing context-specific noise. This is tested against a baseline of "strongly regularized" AR models (using high weight decay) and compared to Diffusion Language Models (dLLMs).

2. Scaling: The SoftQ Law

To address the scaling question, the authors construct a grid of experiments varying model size ( $N \in \{72M, \dots, 1.4B\}$ ) and unique data budget ( $U \in \{100M, \dots, 400M\}$ ). They evaluate several scaling laws:

Chinchilla: An additive form $L = E + A N^{-\alpha} + B U^{-\beta}$ , which assumes the loss gap between data budgets is independent of model size.
Muennighoff-style: Uses effective resources that saturate with repetition.
Quanta: A coupled law derived from skill-learning views.
SoftQ (Proposed): A five-parameter scaling law that couples model size and data size through a "soft bottleneck":
$L_{SoftQ}(N, U) = E + \left( \frac{A}{N^\rho} + \frac{B}{U^{\rho/(1+\alpha)}} \right)^{\alpha/\rho}$
Here, $\rho$ controls the sharpness of the transition between parameter-limited and data-limited regimes. As $U \to \infty$ , it recovers parameter scaling; as $N \to \infty$ , it recovers data scaling.

Key Contributions and Results

1. Regularization Findings

Weight Decay Transfer: The authors demonstrate that strong weight decay (optimized for AR) substantially improves dLLM performance, reducing the validation loss gap between AR and dLLM models. When both are strongly regularized, their performances become comparable.
MIR Efficacy: MIR provides complementary regularization on top of strong weight decay. Across model sizes from 72M to 1.4B trained on the DCLM dataset (100M unique tokens), MIR consistently lowers validation loss compared to the strong-weight-decay-only baseline.
- The gain increases with model size (from ~0.006 loss reduction at 72M to ~0.03 at 1.4B), aligning with the theoretical intuition that overparameterized models benefit more from masking-based regularization.
- Downstream Gains: At 1.4B parameters, MIR yields significant downstream improvements, including +10.2 points on BoolQ and +2.2 points on SciQ.
- Generalization: The benefits transfer to code-heavy data (Stack-V2) even when hyperparameters are tuned solely on natural language data.
Mechanism Analysis: Token-level analysis reveals that MIR helps most on "hard" tokens where the context is noisy, mixed-script, or contains rare names, suggesting MIR improves robustness to local prefix context rather than just frequency-based prediction.

2. Scaling Law Findings

Chinchilla is Misspecified: Empirical results show that the validation loss gap between different unique data budgets grows with model size. The additive Chinchilla law, which predicts a model-size-invariant gap, fails to capture this "fan-out" behavior.
SoftQ Superiority: The proposed SoftQ law, which explicitly couples model and data size, fits the data-constrained experiments significantly better than Chinchilla, Quanta, and Muennighoff-style laws. It achieves lower RMSE and MAE on both in-sample fits and held-out extrapolation (predicting 400M budget performance from smaller budgets).
Data Efficiency Quantification: Using SoftQ as the baseline, the authors estimate that MIR's performance gain is equivalent to roughly 1.3 times the amount of unique training data (specifically in the 200M–400M token budget range).

Significance and Claims

The paper claims to provide a foundational understanding of the data-constrained, compute-rich regime, which is becoming increasingly relevant as compute scales faster than data availability.

Practical Recipe: The authors propose a practical, low-overhead recipe (MIR) that improves data efficiency for autoregressive pretraining without architectural changes or inference overhead.
Theoretical Correction: They demonstrate that classical additive scaling laws are insufficient for the data-constrained regime and propose SoftQ as a more accurate model for predicting loss surfaces when unique data is the binding resource.
Interplay of Regularization and Scaling: The work highlights that regularization strategies (like weight decay and masking) and scaling laws are deeply intertwined in this regime; effective scaling requires models that are properly regularized to prevent overfitting on repeated data.

The authors remain modest, noting that their experiments are limited to 1.4B parameters and 400M unique tokens, which are small relative to frontier-scale pretraining. They suggest that future work should explore hyperparameter transfer recipes and varying model architectures.