Demystifying Training-Time Augmentation for Data-Constrained Language Model Pretraining

Michael K. Chen; Xikun Zhang; Fan Bai; Zhengding Hu; Zhen Wang

Source-linked AI summary

Demystifying Training-Time Augmentation for Data-Constrained Language Model Pretraining

Michael K. Chen, Xikun Zhang, Fan Bai, Zhengding Hu, Zhen Wang

arXiv:2606.16246v2Jun 15, 2026cs.LGcs.AIcs.CL

TL;DR

Fixed-corpus autoregressive pretraining overfits in the data-constrained regime, limiting productive use of abundant compute. This paper evaluates training-time augmentation families and finds that their strongest combination lowers minimum validation loss from 4.015 to 3.805.

Problem
The paper asks whether training-time augmentations can reduce severe overfitting during many-epoch autoregressive pretraining on fixed corpora.
Method
The authors systematically ablate token noise, sequence permutations, and target-offset prediction as augmentation families for autoregressive pretraining.
Results
The strongest combination lowers minimum validation loss from 4.015 to 3.805, below every individual method and naive stack.
Takeaways & Limitations
Augmentation can mitigate autoregressive pretraining’s data inefficiency, with random replacement and right-to-left prediction performing better than masking and Fill-in-the-Middle.
Takeaways & Limitations
Validation loss does not perfectly predict downstream accuracy, which the authors treat as corroborating rather than primary evidence.

Abstract

from arXiv · show

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 training-time 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~\footnote{All code and data are available at https://github.com/ michaelchen-lab/ data-augmentations-for-pretraining.

1 Introduction

As high-quality text becomes supply-limited, language-model pretraining must extract more reusable signal from fixed corpora, but repeated autoregressive training shifts from generalization to memorization and rising held-out loss. This study evaluates training-time augmentations as regularizers for many-epoch, fixed-corpus pretraining and finds that carefully composed, evaluation-aligned views are most effective.

Introduction: Compute is growing faster than the supply of high-quality human text, shifting pretraining toward a compute-abundant, data-constrained regime.The binding constraint becomes how much generalizable signal a model can extract from each token.
Introduction: Repeated autoregressive next-token training on a fixed corpus causes memorization, rising held-out loss, and rapidly diminishing value from additional epochs.Past a few epochs, the value of repeated data falls toward zero, while deeper training drives held-out loss upward.
Introduction: The paper studies token corruption, sequence permutations, and target-offset prediction as diffusion-style view generation for regularizing many-epoch fixed-corpus AR pretraining.Experiments use a 150M-parameter Llama-based model trained on 75M filtered-web-text tokens, with held-out validation loss as the primary metric.
Introduction: Augmented views help when they remain close to left-to-right evaluation: random replacement outperforms masking and right-to-left helps, while Fill-in-the-Middle provides no benefit.Target diversity likewise works only when anchored near next-token prediction; a wide uniform horizon erases signal, whereas an exponentially weighted horizon acts as an implicit curriculum.
Introduction: 4.015 to 3.805: low-rate random replacement plus right-to-left prediction and offset prediction achieves the strongest minimum validation loss.Token noise and offset prediction interfere when noise corrupts local context, while right-to-left and offset prediction reinforce one another; the combination beats every individual method and naive stack.

2 Method

The method introduces three orthogonal training-time augmentation categories for autoregressive pretraining: token-level noise, sequence permutations, and target-offset prediction. Augmentations modify only training input-label pairs, can be composed, and are disabled at evaluation under standard left-to-right next-token prediction.

Overview: Three augmentation categories target autoregressive pretraining: token-level noise, sequence permutations, and target-offset prediction.The categories are designed as orthogonal transformations of training examples.
Token-level noise: Token noise masks or randomly replaces selected content tokens while preserving original labels and protecting control tokens.Masking removes lexical signal, whereas random replacement supplies a plausible but incorrect token.
Sequence permutations: Sequence permutations route samples to right-to-left prediction or Fill-in-the-Middle arrangements, with labels matching the resulting prediction order.FIM rearranges prefix, suffix, and middle segments so the model predicts the middle after seeing both surrounding segments.
Target-offset prediction: Target-offset prediction replaces immediate-next-token prediction with xt+i, using a per-sample offset token and masking positions without a valid target.At evaluation, the offset is fixed to i = 1, restoring standard next-token prediction.
Composition and evaluation: Augmentations compose sequentially: token noise precedes FIM, then direction and offset controls are added before labels are constructed.At evaluation, noise and FIM are disabled, direction is left-to-right, and i = 1.

3 Experimental Setup

The experiments use a 150M-parameter Llama-based causal language model trained with a Warmup-Stable-Decay schedule, while the baseline reaches its minimum validation loss early and then deteriorates.

Training setup: The model has a maximum context length of 2048 tokens and tied input/output embeddings to keep the parameter count tractable.Training uses the Warmup-Stable-Decay learning-rate schedule, which separates constant stable training from a short final decay.
Baseline behavior: The baseline autoregressive model reaches its lowest validation loss at epoch 16 and deteriorates continuously afterward.The extended training budget therefore does not improve minimum loss after the optimum is reached.

4 Experiments

Across experiments, augmentation delays the early collapse of standard autoregressive pretraining, but its effectiveness depends on preserving the evaluation signal and combining compatible objectives. The best configuration, random 5% noise with R2L and exponentially weighted offsets up to 5, achieves a minimum loss of 3.805 at epoch 68, while stable-phase rankings largely persist after decay.

Baseline: The baseline reaches its minimum loss of 4.015 at epoch 16, after which validation loss increases monotonically as training collapses into memorization.Standard AR pretraining becomes counterproductive within the first 20% of the 100-epoch run.
Individual augmentations: Random token replacement regularizes better than masking because plausible-but-wrong tokens require judging correctness from context rather than detecting an obvious gap.The corruption type therefore changes the strength of regularization, not merely the presence of noise.
Sequence permutations: R2L at 50% achieves a minimum loss of 3.910 at epoch 32, outperforming R2L at 25%, which reaches 3.942 at epoch 24.The balanced direction split provides more effective exposure to the reversed objective than the lopsided split.
Sequence permutations: FIM reaches a minimum loss of 3.947 at epoch 16, then climbs steeply and surpasses the baseline by epoch 40.The results attribute this failure to mismatch between FIM’s rearranged sequences and the standard left-to-right evaluation setting.
Target offsets: Exponential weighting over offsets i≤5 achieves a minimum loss of 3.870 at epoch 60, whereas uniform sampling over i≤5 produces a baseline-like minimum loss of 4.016 at epoch 32.Offset diversity helps when prediction usually remains anchored to the next token rather than being spread uniformly across a wide horizon.
Combinations and validation: Random 5% + R2L 50% + i≤5 exp. achieves a minimum loss of 3.805 at epoch 68, improving on the best individual method by 0.036 and beating all previously tested configurations.By contrast, token noise and offset prediction interfere strongly, while R2L and offset prediction combine synergistically; stable-phase rankings are largely preserved after decay, with a best decay minimum of 3.792.

5 Related Work

Related work identifies data scarcity and multi-epoch overfitting as central pretraining challenges, while exploring diffusion models and training-time objectives or augmentations as alternative regularization strategies.

Data-constrained pretraining: Data-constrained pretraining studies find that repeated data produces diminishing returns, memorization, degradation, and overfitting despite abundant compute.These works characterize failure modes in the multi-epoch regime and motivate remedies acting primarily at the data or objective level.
Diffusion language models: Diffusion language models resist multi-epoch overfitting better than autoregressive models, potentially because varied corruption levels and factorization orders regularize training.Adopting diffusion models wholesale would require leaving the autoregressive stack supporting current training and inference infrastructure.
Training-time objectives and augmentation: Prior training objectives and augmentations generate alternative sequence views through masking, permuted factorization orders, denoising mixtures, infilling, and instance-level augmentation.Examples include masked language modeling, UL2’s Mixture-of-Denoisers, Fill-in-the-Middle, Meet-in-the-Middle, and augmentation methods from vision.

6 Conclusion and Discussion

The study finds that standard autoregressive pretraining is highly data-inefficient in data-constrained settings, while augmentation mitigates overfitting and merits consideration as a core pretraining technique. Random token replacement is the strongest token-level noise method, outperforming conventional masking for identifiable mechanistic reasons.

Data inefficiency: At epoch 16, the baseline reaches minimum validation loss, after which over 80% of the training budget becomes counterproductive.Validation loss degrades continuously thereafter.
Data inefficiency: Data augmentation directly addresses the baseline’s inefficiency in the data-constrained regime.The experiments position augmentation as a practical response to continued post-optimum degradation.
Practical implication: Augmentation deserves serious attention as a first-class technique for pretraining in data-constrained settings.This is presented as a broader implication of the study’s results.
Token-level noise: Random token replacement consistently outperforms masking among token-level noise methods.The authors attribute this to the difficulty of disambiguating plausible-but-wrong tokens and the greater distribution shift between training and validation data.

A Limitations · B Future Work

The study is limited to one small model and data scale, with incomplete coverage of augmentation combinations. Future work should test scaling, data-to-parameter regimes, interference patterns, and dataset sensitivity.

A Limitations: Experiments use a single 150M-parameter Llama-based model trained on 75M tokens.This training set is 40× below the Chinchilla-optimal data budget.
A Limitations: Compute constraints prevent evaluating whether augmentation rankings generalize to larger models.The question remains open beyond the tested model and data scale.
A Limitations: The experiments do not establish whether rankings hold near the Chinchilla-optimal data budget.The tested corpus is 40× below that reference regime.
A Limitations: Hyperparameter and augmentation-combination coverage is not exhaustive.Only a subset of possible two- and three-category combinations is explored.
B Future Work: Future ablations should span multiple model sizes and data-to-parameter ratios.This would clarify whether observed rankings are universal or regime-specific.
B Future Work: Future work should test whether scale preserves interference patterns between augmentation categories.One example is token noise disrupting offset prediction.
B Future Work: Dataset sensitivity should be evaluated across corpora with different domain distributions.The current experiments use a single web-text corpus without a model-based quality filter; C4 is one suggested example.

C Extended Related Work

Prior work establishes autoregressive next-token prediction as the dominant pretraining paradigm while highlighting a data wall and severe multi-epoch inefficiency. Alternatives and augmentation methods address overfitting, but existing studies differ from this work in motivation, scope, or systematicity.

Autoregressive language model pretraining: Autoregressive pretraining predominantly uses causal left-to-right next-token prediction, with model quality scaling predictably with parameter count and dataset size.The paradigm was established by the GPT line and carried forward by modern large language models.
The data wall and AR inefficiency: High-quality public internet text may be exhausted within a few years as GPU compute grows faster than data availability, intensifying autoregressive data inefficiency.The related work frames this mismatch as a hard limit for conventional scaling.
Diffusion language models as an alternative: Diffusion language models are reported to be more robust to high-epoch overfitting than autoregressive models, potentially because denoising across corruption levels and factorization orders regularizes training.These studies present diffusion as an alternative response to data-constrained pretraining.
Training-time augmentation objectives: Earlier training-time objectives, including masked language modeling and limited token-noise comparisons, were developed mainly for compute-constrained single-epoch optimization rather than systematic multi-epoch regularization.The cited data-constrained ablations use only one or two augmentation types and do not systematically vary hyperparameters.
Data augmentation in computer vision: Computer vision has long relied on augmentation methods such as random cropping, flipping, color jitter, CutMix, MixUp, and RandAugment in state-of-the-art classifiers and self-supervised models.The passage presents augmentation as foundational to deep learning success in computer vision.

D Training Details

The 150M decoder-only model follows the Llama architecture, with hyperparameters extrapolated from the DCLM scaling recipe. Training uses AdamW, a Qwen2 tokenizer with tied embeddings, and augmentation-specific protected control tokens.

Model architecture: The decoder-only causal model uses Llama-style pre-normalization with RMSNorm, SwiGLU blocks, RoPE, and tied input/output embeddings.RMSNorm uses ε = 10^-6; tied embeddings keep the parameter count manageable for the large vocabulary.
Model architecture: The 150M model is scaled to 20 layers, width 512, 4 attention heads, and intermediate size 1,536 while preserving dhead = 128.These settings are linearly extrapolated downward from the DCLM recipe and use an intermediate size approximately 3 × 512.
Optimization: Optimization uses AdamW with peak learning rate 6 × 10^-4, weight decay 0.033, 100 warmup steps, gradient clipping at 1.0, and batch size 512 × 2,048 tokens.The optimization settings follow the DCLM recipe for learning rate and weight decay.
Tokenization: The Qwen2 tokenizer has vocabulary size 151,646, while tied embeddings reduce the embedding contribution from approximately 77M parameters by half.At hidden size 512, the embedding matrix contributes 151,646 × 512 ≈ 77M parameters before tying.
Augmentation implementation: Augmentation-specific direction, offset, mask, and FIM control tokens are added as needed and protected from token-noise corruption.The tokens include direction pairs, per-offset prediction tokens, one mask token, and three FIM control tokens.

E Held-Out Validation Details

Held-out validation loss is measured on a fixed dataset from a different DCLM-RefinedWeb shard than training. Both splits use identical preprocessing, while validation uses standard left-to-right next-token prediction.

Validation dataset: The primary metric is held-out validation loss on a fixed validation dataset from a different DCLM-RefinedWeb shard than the training split.The passage identifies held-out validation loss as the primary evaluation metric.
Preprocessing: Both training and validation splits tokenize and pack documents into contiguous 2,048-token blocks, discarding remainder tokens.Because packing is contiguous, individual examples may span document boundaries.
Evaluation protocol: Validation is always evaluated with standard left-to-right next-token prediction using i = 1.The passage specifies the evaluation direction and target offset.

F Decay-Phase Training Details

The decay phase resumes each configuration from its best stable-phase validation-loss checkpoint, then applies a 1−√· learning-rate decay over approximately 20% of the stable-phase training steps. Other hyperparameters remain unchanged, with no additional warmup at resumption, and final loss is the minimum validation loss across checkpoints.

Checkpoint selection: Each of the eight configurations resumes decay from the checkpoint with the lowest held-out validation loss during stable training.The resume checkpoint is specified by its global training step and corresponding epoch.
Decay schedule: The decay phase uses the 1−√· learning-rate schedule, decreasing the multiplier from 1 at its start to 0 at the final training step.The schedule is parameterized by the current step, final step, and number of decay steps.
Decay schedule: Ndecay is set to approximately 20% of the stable-phase training steps, with exact per-run values listed in Table 7.Table 7 reports the resume checkpoint and total decay steps for all eight runs.
Decay schedule: Batch size, weight decay, AdamW β values, and context length stay identical to the stable phase, and no additional warmup is applied at resumption.The final converged loss is the minimum validation loss observed at any checkpoint during decay.

G Downstream Evaluation Details

Downstream evaluation uses zero-shot, left-to-right next-token prediction without task-specific fine-tuning or prompting, retaining five benchmarks that provide informative variation at 150M parameters. Other tasks were excluded because they saturated, collapsed, lacked discriminative variation, or were near-random.

Evaluation uses lm-evaluation-harness in zero-shot mode with standard left-to-right next-token prediction, without task-specific fine-tuning or prompting.
Five benchmarks were retained: HellaSwag, PIQA, ARC-Challenge, WinoGrande, and COPA.All tasks are scored by accuracy; HellaSwag, PIQA, ARC-Challenge, and WinoGrande use length-normalised accuracy, while COPA uses standard accuracy.
Exclusion rationale: At 150M parameters, LAM-BADA produced 0% accuracy universally because its long-range dependencies exceeded the small model’s capacity.
Exclusion rationale: BoolQ, RTE, and CommonsenseQA collapsed to majority-class predictions regardless of augmentation, indicating no task-relevant representations were learned.
Five retained tasks varied by at least 2 pp across configurations and showed plausible ordering relative to validation loss, making them the most informative available signal.ARC-Easy and OpenBookQA had near-identical scores across configurations with variance <0.5 pp, while SciQ was near-random at ≈5%.