One learning rule lets training data change across stages
Training Foundation Models with a Single Algorithm · arXiv:2610.40564
In plain words
Foundation models are often trained in separate stages, with different data and learning recipes. This paper introduces unified training, a framework and system that uses one learning rule while changing which examples the model learns from.
It can learn from ordinary examples and from model answers that receive rewards. In controlled reasoning puzzles, the strongest unified setups beat a staged reference in final accuracy, including on a larger puzzle size absent from training.
The research
Foundation models commonly use separate training stages, with different data and learning algorithms. Could one learning rule cover the whole process, so that training instead depends on which data the model sees and when? The paper investigates this question for foundation-model training.
Unified training uses one likelihood-based rule for static examples and suitable rewarded model answers, while letting their proportions change over training.
In controlled puzzle experiments, researchers compared unified setups with a staged reference after training both on the same pretraining stream and reasoning corpus, but exposing the reasoning data at different times and amounts.
The strongest unified setups had higher final accuracy on the main tests and on 12-person puzzles, whose size was absent from training.
A compute-based comparison qualified the result: the leader depended on the accuracy or coverage measure and the budget. The evidence supports flexible data timing in these tests, not one best schedule.
The experiments are small-scale and mostly train on one reinforcement-learning task at a time. They do not establish a universally best data mixture or schedule, or show how results scale to larger models and longer sequences of tasks.
What the evidence shows
- One rule across stages
Unified training is the paper’s framework and training system for applying one likelihood-based learning rule across foundation-model training stages. The question is whether those stages need different learning algorithms, or whether the same rule can learn from changing sources of data. The framework treats staged training as one possible schedule for changing the data mixture.
- How the data is combined
The system combines static examples with suitable reinforcement-learning targets. For binary rewards, it can learn from generated answers that succeeded; the generator may be the model being trained or another model.
The rule stays the same as the data sources and their proportions change, while the framework’s likelihood interpretation applies only to suitable reinforcement-learning objectives.
- Puzzle results and scope
In controlled reasoning puzzles, researchers compared unified configurations with a selected staged reference. The best unified configurations had higher final accuracy on the main tests and on 12-person puzzles, which were larger than puzzles used in either midtraining or reinforcement learning. This supports a benefit in the tested puzzle setting, including for that larger puzzle size.
- Compute qualifies the result
The result changes when the comparison accounts for training compute. The best observed approach depended on the accuracy or coverage measure and the compute budget; on the 12-person test, unified training led on both measures at almost all tested budgets.
These observed checkpoints qualify the final-score result, rather than establishing a general scaling rule.
- Tradeoffs and later learning
A separate mathematical training study found tradeoffs as the share of a heterogeneous data corpus changed: math accuracy stayed similar with small degradation at higher proportions, while math coverage and retention measures improved.
Coding accuracy declined as the mixture grew. Yet after a shared coding-training procedure, the tested unified starting models had higher final accuracy than the staged model on two coding tests.
Deep
- Original paper authors
- Yuda Song, Alvin Zhang, Fahim Tajwar, Daman Arora, Jeff Schneider
- Submitted to arXiv
- License
- CC BY 4.0