Data Augmentations for Data-Constrained Language Model Pretraining

2026-06-15

AI Summary

This work addresses the severe overfitting that autoregressive language models exhibit under data-constrained yet compute-rich pretraining regimes, where repeated training epochs on a fixed corpus degrade generalization. To mitigate this, the authors propose data augmentation as a regularization mechanism enabling efficient pretraining for hundreds of epochs on static datasets. Three orthogonal augmentation strategies are introduced: token-level noise (e.g., random token replacement), sequence reordering (e.g., right-to-left prediction and infilling), and target-shifted prediction (e.g., forecasting future tokens). Empirical results demonstrate that each strategy effectively reduces validation loss, with random token replacement yielding the strongest individual gains. Combining these augmentations further lowers validation loss, substantially delaying overfitting and enhancing training efficiency.

Abstract

As AI labs approach a data ceiling where compute capacity outpaces the rate of new high-quality text generation, language model pretraining is shifting toward a data-constrained, compute-abundant regime that demands productive multi-epoch training on fixed corpora. Standard autoregressive (AR) pretraining overfits severely in this setting, reaching its optimum early and then continuously deteriorating. We investigate data augmentation as a regularizer to mitigate this overfitting and enable productive training for hundreds of epochs on the same data. We introduce three orthogonal categories of augmentation for AR pretraining: token-level noise (masking, random replacement), sequence permutations (right-to-left prediction, Fill-in-the-Middle), and target offset prediction ($x_{t+i}$ for $i > 1$). Through systematic ablations, we find that individual augmentations delay overfitting and lower validation loss relative to the baseline, with random token replacement achieving the best minimum loss among individual methods. Combining augmentation categories further lowers the minimum validation loss. Our experiments demonstrate that data augmentations mitigate AR pretraining's data inefficiency and offer a promising solution to the data-constrained regime. All code and data are available at https://github.com/michaelchen-lab/data-augmentations-for-pretraining

Problem

• data-constrained pretraining
• overfitting
• language models
• data augmentation
• autoregressive pretraining

Innovation

• data augmentation
• autoregressive pretraining
• overfitting mitigation
• multi-epoch training
• data-constrained regime