Training Foundation Models with a Single Algorithm
Yuda Song, Alvin Zhang, Fahim Tajwar, Ruslan Salakhutdinov, Aarti Singh, Andrea Zanette and colleagues (Carnegie Mellon and Johns Hopkins, with Aurora Innovation) place pretraining, midtraining, SFT and suitable RL objectives under one likelihood-based algorithm, turning training into a choice of data, proportions and timing.
Ask this paper
Unified formulation. Suitable RL objectives are derived as maximum-likelihood fitting to model-generated targets, so RL generations and static data share one loss and one trainer. Staged training becomes a special case.
Controlled study. On a 124M model trained from scratch with a 20B-token pretraining corpus and a 510M-token midtraining corpus, the best unified runs outperform staged training on RL tasks and on larger unseen puzzles. Those runs reuse the midtraining corpus 2 and 4 times.
Compute budgets. Unified training has the best accuracy at the smallest and larger budgets, while staged training leads at intermediate budgets. Generalization favors unified training at almost all budgets.
Real-world tradeoffs. On a 0.76B model with math RL and 0 to 5% midtraining data, math accuracy stays similar and retention improves. Coding accuracy falls as the proportion grows, and ending the mix halfway recovers some of it.
Continual learning. After math training, all unified models reach higher final coding accuracy than staged training despite starting lower, and the 5% mixture performs best. The experiments are small-scale and identify no universally preferred mixture.
Abstract
Foundation models are typically trained in stages: pretraining, midtraining, supervised fine-tuning (SFT), and reinforcement learning (RL), with changes in both data and learning algorithms. We show that this algorithmic separation is not fundamental. By deriving suitable RL objectives as maximum-likelihood fitting to model-generated target distributions, we place all of these stages under a single likelihood-based training algorithm, unified training. This perspective turns the foundation-model training lifecycle into a problem of choosing which data to train on, in what proportions, and when, with conventional staged training as a special case. We instantiate this framework in a unified training system that mixes static data and RL generations within a single trainer. Across controlled and real-world experiments, unified training improves the accuracy-compute tradeoff and task coverage while preserving broader capabilities, and provides a stronger starting point for continual learning. This work lays the groundwork for a unified science of foundation-model training, connecting mathematical principles, training infrastructure, and the development of model capabilities.