🤖 AI Summary
As pretraining enters a new regime characterized by limited data availability and abundant computational resources, conventional scaling laws and regularization techniques become ineffective. This work proposes Masked Input Regularization (MIR), which adapts the stochastic masking mechanism from diffusion models into autoregressive pretraining to enhance generalization. Furthermore, it introduces SoftQ, the first coupled scaling law tailored for multi-epoch training under finite data constraints. Evaluated across model sizes ranging from 72M to 1.4B parameters, MIR consistently reduces validation loss and yields measurable gains on downstream tasks. SoftQ demonstrates superior fitting accuracy compared to classical scaling laws in data-limited settings, effectively amplifying the usable training data by approximately 1.3×.
📝 Abstract
Classical scaling laws for language model 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 language models 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.