One loss for pretraining, midtraining and RL: unified training lifts accuracy and coverage at matched compute
Synopsis
Foundation models are usually trained in stages, with both the data and the learning algorithm changing at each one. Unified training writes pretraining, midtraining, supervised fine-tuning and reinforcement learning as maximum-likelihood fitting to target distributions, so static corpora and RL generations mix inside a single trainer, and staged training becomes one special case. In controlled reasoning tasks and a real-world midtraining experiment, it improved the accuracy-compute tradeoff and task coverage while preserving broader capabilities.
Interpretation
The whole foundation-model training pipeline can run under one fixed likelihood learning rule, with only the distribution of training data changing. Earlier pipelines treated pretraining, midtraining, supervised fine-tuning and reinforcement learning as separate learning problems, each with its own recipe and infrastructure. The equivalence rests on rewriting the MaxRL gradient as maximum-likelihood fitting to the current policy's success-conditioned distribution, which puts RL in the same loss form as pretraining, midtraining and SFT.
In the implementation, static data and RL generations share one weighted example format and are sampled together by a single trainer. Staged training is the case where all mixture weight sits on one source at a time; unified training lets sources overlap and their proportions vary over training. Generation and optimization run asynchronously, with generators appending completed batches to a shared buffer and the trainer publishing weights every few updates; the controlled study uses a 124M-parameter model with a 20B-token pretraining corpus and a 510M-token midtraining corpus, and the real-data study a 0.76B-parameter model with a 240M-token heterogeneous midtraining corpus.
Perspective
The demonstrated scope is two settings: a controlled study with a 124M-parameter model on a fixed 20B-token pretraining corpus and a 510M-token reasoning-puzzle midtraining corpus, and a real-world study with a 0.76B-parameter model mixing a 240M-token heterogeneous midtraining corpus into mathematical RL; compute comparisons cover pretraining, midtraining and RL.
The authors state that the experiments reveal tradeoffs among accuracy, coverage and retention but do not identify a universally preferred mixture or schedule; the compute study characterizes the evaluated configurations without establishing a predictive scaling law; most runs optimize a single RL task at a time, and downstream adaptation still inevitably requires a separate training stage.