Dual-objective Language Models: Training Efficiency Without Overfitting

AI summary

Overview

Research area: Natural Language Processing — language model pretraining objectives (autoregressive vs. masked-diffusion).

Technical level: Advanced. The paper assumes familiarity with transformer architectures, attention masking, loss functions, diffusion processes, and Gaussian process regression.

Scope: A systematic empirical study of training a single transformer on a weighted combination of autoregressive and masked-diffusion objectives, spanning 50 trained models and varying degrees of training-data repetition.

What This Paper Is About

Autoregressive language models learn quickly but overfit badly when training data must be repeated, while masked-diffusion language models resist overfitting but learn slowly. The authors ask whether a single model can be trained on both objectives simultaneously — using autoregression for fast early learning and diffusion as a regularizer — and, if so, what the optimal mixing weight should be at different levels of data repetition.

Key Contributions

1. A dual-objective training method that combines autoregressive and masked-diffusion losses in one transformer with no architectural modification, enabling one model to handle both unidirectional and bidirectional tasks.
2. A large-scale empirical map of the relationship between data repetition, the objective ratio α, and downstream performance, based on 50 trained models evaluated on nine zero-shot tasks.
3. The finding that dual-objective training beats single-objective training in every evaluated setting — including in regular data settings, where dual-objective models outperform pure masked-diffusion models even though only a small fraction of training goes to the diffusion objective.
4. Two practical recommendations for choosing the objective ratio in regular and data-constrained regimes, plus a third recommendation on induced prefix language modeling.

Main Findings

• Autoregressive models overfit, diffusion models do not but converge slowly. In the extreme setting of 128 repetitions of the training corpus, the autoregressive objective converges fastest and then "catastrophically overfits," while the masked-diffusion objective converges slowly without being negatively affected by the repetitions. The combined model shows both fast convergence and robustness to overfitting.

• Combining objectives always helps. It was optimal to combine both objectives under all evaluated settings, and the optimal balance is similar whether targeting autoregressive or masked-diffusion downstream performance.

• Dual-objective models beat pure masked-diffusion models even without data constraints. Switching to bidirectional evaluation, the single-objective α = 1 model performs poorly, while models with α between 255/256 and 15/16 perform similarly and substantially outperform pure masked-diffusion training.

• Two regimes emerge. A regular-data region — roughly 16 repetitions of training data and less — where a purely autoregressive model does not overfit, and a data-constrained region — roughly 32 repetitions and more — where overfitting is a serious consideration.

• Recommended ratio for regular data. With 16 repetitions or less, train with a small amount of masked-diffusion objective (α ≈ 63/64) to get stronger bidirectional performance than pure masked-diffusion training without losing autoregressive performance.

• Recommended ratio for data-constrained settings. With more than 32 repetitions, choose α so that the autoregressive objective sees roughly 16 repetitions of the training data. The authors note that more than 32 autoregressive repetitions lead to overfitting, while fewer than 8 autoregressive repetitions lead to underfitting.

• Detailed task results favor the dual objective. On normalized autoregressive scores across the nine tasks, the average was 26.9 for the dual model (α = 63/64) versus 26.1 for the autoregressive-only model at 1 repetition; 23.9 for the dual model (α = 3/4) versus 22.0 at 32 repetitions; and 19.1 for the dual model (α = 1/8) versus 13.3 at 128 repetitions.

• Dual objectives learn from very little data. The detailed scores show nontrivial performance even when the model is exposed to just 256M tokens of training data under 128 repetitions.

• Prefix language modeling emerges for free. Processing the conditional part of a prompt fully bidirectionally improves over fully unidirectional evaluation by reliably more than one percentage point across most configurations that combine both training objectives, without any additional training.

• Consistent with prior work on masked diffusion being suboptimal alone. The authors note their results confirm the conclusion of Xue et al. (2025) that masked diffusion alone is a suboptimal objective for language, but show that simply combining both objectives captures the benefits of diffusion without losing performance.

Methodology in Plain English

The authors train one transformer that is used in two modes, differing only in its input and attention mask. In the autoregressive mode, it sees the full unchanged text with a causal (unidirectional) attention mask; in the masked-diffusion mode, it sees partially masked text with a fully bidirectional attention mask. Both modes are unified by using masked next-token prediction (MNTP), where the model always uses the hidden state at position i to predict the token at position i+1. This means the exact same architecture and parameters serve both objectives; the paper proves this parameterization is as expressive as the standard approach in an appendix.

Training minimizes a weighted sum of the two losses, α·L_AR + (1−α)·L_MD. To keep computation efficient, each GPU device is assigned a single objective rather than mixing objectives within a batch. Training is distributed across 256 devices, which permits 256+1 possible values of α.

The experiments train each 470-million-parameter model (360M non-embedding weights) on 32 billion tokens. A repetition factor R means a unique subset of 32B/R tokens is sampled and repeated R times. This budget is more than 4× past the Chinchilla compute-optimal point, chosen because it reflects how modern language models are actually trained.

Models have 24 layers, hidden size 1,024, 16 attention heads, feed-forward intermediate size 3,554, and a vocabulary of 51,200 tokens. They use pre-normalization with RMSNorm, rotational positional embeddings, and Swish-gated linear units. Optimization uses the Muon optimizer (specifically the variant from Liu et al., 2025) with learning rate 0.007, a warmup-stable-decay schedule without warmup steps and 2,048 steps of linear decay, 8,192 total steps, 4M tokens per global batch, and sequence length 2,048. Regularization comes from weight decay (strength 10⁻¹) and an auxiliary z-loss term (strength 10⁻⁴).

Data comes from the HPLT v2 corpus — English documents sampled to 32B tokens, combining extracted webpages from the Internet Archive and CommonCrawl — with a smaller disjoint subset held out for validation. A byte-level BPE tokenizer with 51,200 subwords is trained directly on the full training data.

Evaluation covers nine zero-shot tasks, following OLMES guidelines with two changes: only zero-shot evaluation, and only the "cloze" formulation of each task. Autoregressive scores use Equation 2; bidirectional scores use the semi-autoregressive variation of pseudo log-likelihood (PLL) proposed by Samuel (2025), which the authors report is more than ten times faster than Monte-Carlo estimation and more accurate. Scores are normalized so the random baseline is 0 and the maximum is 1, then averaged across tasks. To handle noisy data, the authors fit a Gaussian process regression model with a composite kernel — a constant kernel multiplied by an anisotropic Matérn kernel (ν = 1.5), combined additively with a white-noise kernel — optimized with L-BFGS-B via SciPy, achieving R² over 0.99 in all cases.

Evaluation tasks and sizes:

BLiMP and MMLU contain 67 and 57 sub-tasks respectively, reported as macro-averages.

Why This Matters

Impact on research. The paper challenges the assumption that autoregressive and masked-diffusion modeling must be treated as competing paradigms. It shows the two are complementary and that combining them is never worse than either alone in the evaluated settings, while also extending the effective repetition limit of autoregressive training — the paper notes that prior work showed autoregressive models cannot meaningfully learn from more than 16 data repetitions, and demonstrates that this value is at least an order of magnitude larger with the dual objective.

Real-world applications:

• Training on limited or proprietary text. Organizations with restricted corpora (legal, medical, or specialized-domain text) can repeat data more times before performance degrades.
• Low-resource language modeling. Languages with small web footprints may benefit from the finding that dual-objective models retain nontrivial performance after only 256M tokens of exposure.
• Dropping into existing pipelines. Because the method requires no architectural changes and no inference overhead, models can be deployed as standard autoregressive transformers.
• Prompt-processing improvements at inference. Bidirectional processing of the conditioning part of a prompt improves downstream scores with no extra training.

Industry relevance. The paper situates its work against the "data wall" — the expected exhaustion of available training data relative to established scaling laws. As compute budgets grow exponentially while high-quality text does not, the ability to extract more signal from repeated data becomes an economic concern for anyone training at scale. The authors release training and evaluation code at https://github.com/ltgoslo/dual-language-models and models under the Apache 2.0 license at https://huggingface.co/ltg/dual-lm-470m.

Future Directions

1. Verifying at larger scale. The authors believe their recommendations hold for larger models because overfitting behavior does not depend on model size according to prior work, and because the relative burden of representing two modes decreases with model size — but they state that reliably answering this would require expensive experimentation.
2. Testing much higher repetition factors. The study covers up to 128 repetitions, leaving open how the optimal α behaves at the hundreds of repetitions the authors expect to become relevant for frontier development.
3. Understanding why small amounts of diffusion help. The authors hypothesize that the prevalence of the autoregressive objective drives fast convergence while a small amount of masked-diffusion induces useful modeling priors, but this mechanism is not resolved.
4. Extending the induced prefix capability. The finding that dual training induces prefix language modeling without any additional training raises the question of what other attention patterns or inference-time behaviors might emerge from the same training recipe.

Target Audience

Researchers and engineers working on language model pretraining — particularly those interested in training objectives, data-constrained or repeated-data regimes, and scaling behavior. It is most useful to readers with a working understanding of transformer attention masking and diffusion-based language modeling, since the method is described in terms of loss functions and attention patterns rather than at an introductory level. Practitioners designing pretraining recipes for smaller or low-resource corpora will find the two α recommendations the most directly actionable part of the paper.

Authors’ abstract

This paper combines autoregressive and masked-diffusion training objectives without any architectural modifications, resulting in flexible language models that outperform single-objective models. Autoregressive modeling has been a popular approach, partly because of its training efficiency; however, that comes at the cost of sensitivity to overfitting. On the other hand, masked-diffusion models are less efficient to train while being more resilient to overfitting. In this work, we demonstrate that dual-objective training achieves the best of both worlds. To derive the optimal balance between both objectives, we train and evaluate 50 language models under varying levels of data repetition. We show that it is optimal to combine both objectives under all evaluated settings and that the optimal balance is similar whether targeting autoregressive or masked-diffusion downstream performance.