Four Training Stages May Just Be a Special Case of One Algorithm

A paper from CMU, Johns Hopkins, and Aurora Innovation puts the four common stages in foundation model training into a single algorithm. The paper title is straightforward: Training Foundation Models with a Single Algorithm.

For readers without a training background, the four stages can be roughly understood as follows: pretraining teaches the model to predict the next word, mid-training continues using specific data to shape capabilities, SFT uses instruction data to teach the model to answer questions, and RL uses reward signals to adjust outputs. Usually they are split into different stages, and the data, loss functions, and trainers may all change.

The paper's entry point is to rewrite a suitable RL objective as maximum likelihood fitting to the target distribution generated by the model. In this way, RL generation and static data share one loss function and one trainer. Staged training becomes a special case of unified training. The training design question then changes: which data to select, in what proportions, and when to feed it in.

The experiments are in several parts. The authors first train a 124M model from scratch, with a 20B-token pretraining corpus and a 510M-token mid-training corpus. The best unified training runs surpass staged training on RL tasks and also do better on larger puzzles the model has not seen. Those runs reused the mid-training corpus 2 and 4 times. Under a fixed compute budget, unified training has the best accuracy at the smallest and larger budgets, while staged training leads at intermediate budgets. Generalization results almost all favor unified training.

More realistic tests use a 0.76B model for math RL, mixing in 0% to 5% mid-training data. Math accuracy remains similar, and retained capabilities improve. Coding accuracy decreases as the mixing ratio increases, and ending in the middle of mixing recovers some of it. Continual-learning results are clearer: after math training, all unified models ultimately have higher coding accuracy than staged training, despite starting lower; 5% mixing is best.

Some commenters said that if stages really can be merged, many training pipeline diagrams would get shorter. Others focused on the intermediate budget: staged training leads at intermediate budgets, suggesting that unified training's search sometimes lands on a worse point, and that point is already included in the search space. The paper abstract does not explain what happens in the intermediate range; this question is more interesting than average scores. Another commenter asked whether the RL part still needs a reward model. The paper converts the RL objective into maximum likelihood fitting; the specific reward setup still requires looking at the original implementation.

The authors themselves acknowledge that the experiments are small in scale and that they did not find a universally optimal mixture or schedule across scenarios. This sentence is key. Unified training moves the engineering problem from splitting stages to selecting data, setting proportions, and deciding timing. The direction is interesting, but it is still far from redrawing large-model training pipelines.

Paper: https://arxiv.org/abs/2610.40564

DAIR.AI paper page: https://academy.dair.ai/papers/training-foundation-models-with-a-single-algorithm-2610.40564