CMU, Johns Hopkins and Aurora paper trains models with one likelihood loss across pretraining, SFT and RL

Large language models are usually built in a sequence of distinct phases: pretraining on web-scale text, sometimes a "midtraining" phase on curated data, supervised fine-tuning (SFT) on examples, then reinforcement learning (RL) from rewards. Each phase tends to have its own algorithm and its own infrastructure. A preprint posted to arXiv on October 29, titled "Training Foundation Models with a Single Algorithm", argues that this separation is a habit rather than a necessity.

The ten authors are from Carnegie Mellon University, Johns Hopkins University and Aurora Innovation; project lead Yuda Song and co-authors including Ruslan Salakhutdinov, Aarti Singh and Andrea Zanette are at CMU, while J. Andrew Bagnell is affiliated with both CMU and Aurora.

One loss for every stage

Pretraining, midtraining and SFT already share a goal: make the model assign high probability to a fixed set of text, a principle known as maximum likelihood. The paper's step is to place RL in the same family. Building on earlier work called MaxRL, the authors use the result that a policy-gradient update is equivalent to pushing the model towards its own outputs filtered for success. In practice, the model generates answers, the correct ones are kept and weighted more heavily when the task is hard for the model, and they are added to a shared training buffer.

With every stage expressed this way, the authors propose "unified training": a single loss applied to a mixture of static data and model-generated successes, with mixture weights that can shift over time. A conventional staged pipeline becomes a special case in which all the weight sits on one data source at a time.

What the experiments show

The tests are deliberately small. A 124-million-parameter model was trained from scratch on 20 billion tokens and then on reasoning puzzles; a 0.76-billion-parameter model was trained with math RL and evaluated on benchmarks such as MATH-500, GSM8K and HumanEval+. The authors report that the best unified setups beat the best staged ones on final accuracy, were more compute-efficient on accuracy at matched budgets, and kept general abilities closer to the pretrained model. Unified models also started from a stronger position when later trained on coding tasks. The gains were not uniform: in one setting, more mixing raised math coverage but lowered coding accuracy.

The paper has not been peer reviewed. Its authors acknowledge that the experiments do not establish scaling laws, that there is no systematic way yet to choose mixtures or schedules, and that most runs train on one RL task at a time.