One likelihood loss for pretraining, SFT and RL: what the unified training paper assumes

Every post-training pipeline I have seen is a relay race. Pretraining hands a checkpoint to midtraining, midtraining hands it to supervised fine-tuning, SFT hands it to reinforcement learning, and each leg has its own data loader, loss, infrastructure and failure modes. Training Foundation Models with a Single Algorithm, posted on October 29 by Yuda Song, Andrea Zanette and eight co-authors including Ruslan Salakhutdinov and J. Andrew Bagnell, argues the race is unnecessary. Their claim is that the right RL objective is just maximum likelihood on a target distribution the model generates, so every stage becomes one likelihood loss over a data mixture. I read the 67 pages to find out whether that is a real simplification or a change of notation.

The core identity

The derivation assumes binary, verifiable rewards: an answer is right or wrong, and a checker can tell. Picture a math prompt the model solves one time in four, so its success probability p is 0.25. The paper shows that the gradient of log p equals the expected gradient of the model's log-likelihood over its own successful samples. Concretely, sample a batch, discard the wrong answers, roughly three in four, and run ordinary supervised fine-tuning on the survivors with that set held fixed during the update. That is exactly one on-policy step of their objective, which they call MaxRL. If you already run a rejection-sampling SFT loop, you are one weighting term away from it.

In finite samples this looks like rejection sampling with a weight. Each correct completion gets weight one over the estimated pass rate, incorrect ones get zero. The difference from REINFORCE is the prompt weighting. REINFORCE maximizes expected reward, which implicitly weights each prompt by how often the model already solves it. MaxRL maximizes log pass rate, so hard prompts count as much as easy ones. With N samples per prompt, the effective weight is one minus (one minus p) to the power N, which reduces to REINFORCE at N equal to one and approaches uniform as N grows.

The unification is exact at the level of the gradient. Everything that makes RL expensive is still there, just moved into the data generator.

That is my honest reading of the "is it only notation" question. The identity is mathematically simple, and once you see it, RL with a 0 or 1 reward really is likelihood training on filtered self-generated data. What it does not remove is generation, verification and staleness. The paper's own system uses a generator that writes completions into a shared buffer and a trainer that publishes weights every k updates, so the samples come from a recent copy of the model and the update is "generally not the current-policy MaxRL gradient." Continuous rewards need extra machinery: a TailRL target that averages over reward thresholds, or a reward-weighted variant.

What the experiments show

The practical payoff of one loss is one batch. Because every stage is now a token-weighted negative log-likelihood, the trainer can fill a single batch with pretraining text, midtraining traces and RL generations, mixed by target token fraction, instead of handing a checkpoint from one pipeline to the next. The paper tests this only on models trained from scratch, at two small scales. The first is 124M parameters on a 20B-token corpus, with Reasoning Gym-style puzzles as the RL tasks. The second is 0.76B with RL on Math12K, followed by three questions: did math improve (MATH-500, GSM8K), did code improve (HumanEval+, MBPP+), and how much general knowledge survived (SciQ, ARC-Challenge, MMLU). During RL in that second setting, midtraining data stayed between 0.5% and 5% of tokens, under constant and annealed schedules, and that share is the knob the results below turn on.

The best unified configurations beat staged training on the puzzle tasks. The clearest win is retention: mixing even a small fraction of static data into RL kept held-out loss below the base model, while staged RL pushed it up. The mixed results are just as informative. Coding accuracy fell as the share of midtraining data grew, even while coding coverage rose. At 2% mixing, plain REINFORCE beat MaxRL on math accuracy and on SciQ and ARC-Challenge, while MaxRL won on math coverage. Staged training was more compute-optimal for coverage. And no KL penalty they tried reproduced the retention gain, which is a useful negative result for anyone tuning a KL coefficient today.

The limits the authors state

The authors are direct about scope. "The scope of the experiments also remains small-scale." There is no universally preferred mixture or schedule, no scaling law with predictive power, most runs optimize a single RL task, and downstream adaptation "will inevitably require a separate training stage." Below one billion parameters trained from scratch is a long way from post-training a 400B open-weights checkpoint, and the verifiable-reward assumption leaves preference tuning with learned reward models outside the clean part of the theory. That is the setting where the rubric becomes the reward and failure modes look very different.

For cost, the honest summary is that the paper saves engineering, not FLOPs. One trainer, one loss and one data pipeline are cheaper to maintain and easier to reproduce than four. The generator still dominates RL compute, and choosing the mixture and schedule becomes a new hyperparameter search.

What I would try

If you already run RL with verifiable rewards on an open model, two changes are cheap enough to test this month. Keep 1% to 5% of your SFT or general text data in every RL batch as plain likelihood loss, and track held-out loss on general text alongside your task metric, the same way I watch for the usual fine-tuning overfitting signs. Then compare REINFORCE-style weighting with one-over-pass-rate weighting on correct samples at your normal N. Both are a few lines in an existing trainer, and the retention curve will tell you within one run whether the unified view pays off at your scale.