CMU's Unified LLM Training: One Maximum Likelihood Objective for All Stages

Large model training traditionally follows a pipeline: pre-training, mid-training, supervised fine-tuning, then reinforcement learning, each with different data, objectives, scripts, and systems. A new CMU paper introduces a more concise organization: the learning rule can always be maximum likelihood; the real design challenge is how different data sources appear, mix, and exit throughout the entire training lifecycle.

Stages Remain, but Boundaries Open

The conventional pipeline hands all weights to one data type at a time. Unified training allows overlap: web corpora, demonstration answers, and model-generated outputs all enter a single mixed data distribution. This shifts the core question from "which algorithm for the next stage?" to "what to learn now, at what proportion, and when to adjust?"

Figure 1: Static data and RL-generated data share weighted likelihood loss; right side compares performance under same total compute.

Key Step: Casting RL as Maximum Likelihood via MaxRL

Pre-training, mid-training, and supervised fine-tuning already optimize maximum likelihood. The difficulty lies in RL. Using MaxRL, the binary-reward update is reinterpreted as fitting the current policy's answer distribution conditioned on success. For each prompt, multiple responses are sampled; correct responses are weighted by the inverse of their pass rate, incorrect responses receive zero weight. These then join static corpora and demonstrations under the same per-token negative log-likelihood loss.

Table 1: Prompt and target distributions for four data sources under the same objective.

Figure 2: Generator and trainer run asynchronously, share a data buffer, and periodically sync model weights.

Engineering-wise, this eliminates the need for two completely separate training logics. Prompt tokens are masked; corpora, demonstrations, and RL responses differ only in weight, allowing them to feed into a single trainer.

Results Must Be Evaluated in Total Compute Budget

Controlled experiments use a 1.24B parameter model with fixed 200B pre-training tokens and 510M mid-training tokens, then learn verifiable reasoning puzzles. The best unified configuration outperforms the staged baseline on Base, Extended, and larger unseen puzzles. Notably, on unseen large puzzles, the unified approach achieves the highest accuracy and coverage across almost all compute budgets. The advantage is not uniform: staged training remains more compute-efficient at medium budgets, while unified training's accuracy gains appear mainly at higher budgets. Comparing training schemes requires accounting for both final score and the compute needed to reach it.

Figure 5: Horizontal axis includes all compute for pre-training, mid-training, and RL; top row pass@1, bottom row pass@16.

Real Data Changes the Capability Combination

The team further trains a 760M parameter model on math RL while mixing 240M tokens of heterogeneous data covering web, math, code, and reasoning trajectories. Non-zero mixing ratios reduce held-out corpus loss; larger ratios yield the highest math coverage and improve SciQ overall. Math accuracy drops only slightly, but code accuracy suffers, while code coverage exceeds both the base model and staged training. Unified training does not raise all metrics simultaneously; it provides an adjustable capability combination where mixing ratios and schedules become the new control knobs.

Figure 7: As data ratio increases from 0 to 5%, math, code, corpus losses and SciQ show different trends.

Benefits Extend to Subsequent Learning

After math training, all models undergo identical code RL without additional mid-training data. Despite some unified models starting with lower code accuracy, they all surpass staged models on HumanEval+ and MBPP+ after 4000 code RL updates, with the 5% mixing group achieving the best across four metrics. This indicates that a dataset's value should be judged not only by immediate scores but also by whether it leaves a better starting point for future learning.

Figure 9: Unified training initialization shows better continual learning results over 4000 subsequent code RL updates.

Current conclusions have clear boundaries: experiments are at small scale, most tasks optimize only one RL objective at a time, and no universal mixing recipe or predictive scaling law has emerged. Next steps require joint study of data scheduling and optimization strategies on larger models, multi-task settings, and longer task sequences.

Training Foundation Models with a Single Algorithm
https://arxiv.org/pdf/2610.40564
https://github.com/yudasong/unified-training