📌 Highlights
The central claim is that foundation-model training can be treated as one likelihood-based learning process driven by an evolving mixture of static examples and model-generated targets.
The strongest evidence comes from controlled reasoning experiments, where unified schedules outperform staged training on final task performance and larger-puzzle generalization, and from the compute-matched analysis showing unified training is preferable at many small and large budgets while staged training remains competitive at intermediate budgets.
In the heterogeneous mathematics-and-coding study, nonzero midtraining mixtures improve retention and later coding adaptation, but can trade coding accuracy for coverage and broader capability retention.
The main unresolved issue is scale and recipe dependence: the evidence comes from 124M- and 0.76B-parameter models, and the best mixture or schedule varies by task and objective.
- One trainer can express pretraining, midtraining, SFT, and suitable RL through token-level weights.
- Unified training changes the optimization problem from stage selection to data-source composition, timing, and compute allocation.
- MaxRL converts successful generations into a likelihood target; other RL objectives correspond to different target and prompt distributions.
- Results support better retention and continual learning, but not uniform superiority at every compute budget or task metric.
🎯 Introduction
Conventional foundation-model pipelines separate pretraining, midtraining, SFT, and RL into specialized stages with different data and optimization procedures. The paper argues that this separation is not fundamental because the first three stages already share maximum likelihood, while suitable RL objectives can also be written as likelihood fitting to reward-conditioned model outputs.
The paper-level objective is therefore to study the entire training trajectory through an evolving data distribution: which sources are sampled, in what proportions, and at what times. Staged training becomes an abrupt change in mixture weights, whereas unified training permits overlapping sources and gradual schedules. This framing is intended to connect capability development, data reuse, retention, and compute allocation in one analysis.
🔬 Methodology
For each source s at iteration t, the method defines a prompt distribution \(\rho_{s,t}\), a target completion distribution \(\mu_{s,t}\), and a mixture weight \(\lambda_{s,t}\). Static pretraining and midtraining use empty prompts, SFT uses instruction–demonstration pairs, and RL uses successful completions from the current or another generation policy. The unified distribution \(q_t\) is optimized with the same likelihood loss regardless of how its data were produced.
The practical system stores a token sequence and a nonnegative weight for every example. Pretraining and midtraining use unit token weights; SFT masks prompt tokens; RL masks prompts and gives successful response tokens a reward-derived weight. A generator appends evaluated responses to a shared buffer while the trainer samples static and generated data asynchronously, periodically synchronizing weights.
For binary rewards, the generator samples N completions per prompt, estimates the pass rate \(\hat p\), assigns \(1/\hat p\) to correct responses, and assigns zero to incorrect responses. Finite sampling leaves the success-conditioned target unchanged but down-weights prompts that are rarely solved; at \(N=1\) the weighting recovers REINFORCE-like prompt weighting, while larger N approaches the MaxRL prompt distribution.
Figure 2
Caption: Figure 2: Left: the training system. The generator produces RL data from a recent copy of the weights and appends it to the data buffer, which also holds static data; the trainer samples from the buffer and periodically publishes its weights, which the generator loads. Right: the data format. Each cell is a token and the number is its weight in Eq. (2); the stages differ only in the weights.
Why It Matters: It shows how one buffer, one trainer, and token-level weights operationalize the theoretical unification.
Figure 3
Caption: Figure 3: Data composition throughout training. The illustrated run from Appendix C.8 mixes midtraining tokens into pretraining, then combines RL generations, midtraining data, and pretraining data. Token counts include prompts and are rounded estimates; the axis break omits 10B pretraining tokens, and thin bands are enlarged for legibility.
Why It Matters: It illustrates that unified training is a changing data mixture rather than a fixed replacement of one stage by another.
Table 1
Caption: Table 1: Data sources for the shared likelihood objective. For binary rewards, ν+ is the success-conditioned distribution of the generation policy ν; ν = πθt recovers on-policy MaxRL.
| Source | Prompts \(\rho_{s,t}\) | Target \(\mu_{s,t}\) |
|---|---|---|
| Pretraining | empty | \(\mu_{\mathrm{pre}}\) |
| Midtraining | empty | \(\mu_{\mathrm{mid}}\) |
| SFT | \(\rho_{\mathrm{sft}}\) | \(\mu_{\mathrm{sft}}\) |
| RL | \(\rho_{\mathrm{rl}}\) | \(\pi^+_{\theta_t}\) or \(\nu^+\) |
Why It Matters: The table makes explicit that the stages differ in their data distributions, not necessarily in the optimizer or loss.
Table 2
Caption: Table 2: RL objectives within the unified likelihood framework. Different objectives correspond to different target and effective prompt distributions.
| Method | Reward | Target \(\mu_s\) | Effective prompts \(\rho_s\) |
|---|---|---|---|
| MaxRL | Binary | \(\pi^+_{\theta_t}\) | \(\rho(x)\) |
| REINFORCE | Binary | \(\pi^+_{\theta_t}\) | \(\propto \rho(x)p_{\theta_t}(x)\) |
| TailRL | \([0,1]\) | \(\int_0^1 \pi_{\theta_t}(\cdot\mid x,r>\tau)\,d\tau\) | \(\rho(x)\) |
| REINFORCE | \([0,1]\) | \(\frac{\pi_{\theta_t}(y\mid x)r(x,y)}{\mathbb{E}_{y'\sim\pi_{\theta_t}(\cdot\mid x)}[r(x,y')]}\) | \(\propto \rho(x)\mathbb{E}_{y\sim\pi_{\theta_t}(\cdot\mid x)}[r(x,y)]\) |
Why It Matters: It separates the target completion distribution from the effective prompt distribution, clarifying why different RL objectives can share the same trainer while producing different updates.
Unified mixture distribution
Meaning: This controls how source distributions are combined at each training iteration; staged training is the special case where one source has all mixture weight.
MaxRL likelihood identity
Meaning: This is the bridge that turns binary-reward MaxRL into maximum-likelihood fitting of successful completions.
Equation 2
Meaning: This is the implemented token-weighted loss; token masks and reward weights determine how each source contributes to the shared trainer.
📊 Experiments
The controlled study trains a 124M-parameter model from scratch on a fixed 20B-token pretraining corpus and a fixed 510M-token midtraining corpus. The task is a verifiable natural-language reasoning-puzzle suite in the style of Reasoning Gym. Base puzzles stay within the midtraining size range; Extended introduces larger puzzles during RL; Extended generalization evaluates on still larger puzzles excluded from both midtraining and RL. Both unified and staged systems use MaxRL, with 6000 RL updates for Base, 12000 for Extended, the same number of generations per update, and results averaged over two RL seeds. The reported metrics are held-out pass@1 accuracy and pass@16 coverage.
The best unified configurations outperform staged training on Base, Extended, and Extended generalization. More midtraining exposure during pretraining generally helps; mixing during RL is beneficial but less consistent, with an intermediate mixture best for Base at the largest pretraining mixture while Extended settings generally improve with more midtraining during RL. In compute-matched curves, staged training is strongest at some intermediate accuracy and coverage budgets, whereas unified training is stronger at the smallest and larger accuracy budgets and at most generalization budgets.
The retention analysis finds that unified initializations remain similar to pretraining-only references on held-out pretraining loss and mean performance across five general benchmarks, while the best staged retention configuration is worse on both measures. In the real-world experiment, a 0.76B-parameter model is trained with Math12K and a 240M-token heterogeneous corpus containing web text, mathematics, code, and reasoning traces. Math averages MATH-500 and GSM8K; Coding averages HumanEval+ and MBPP+. Constant midtraining mixtures from 0% to 5% keep Math accuracy broadly similar while larger mixtures improve coverage and retention, but coding accuracy declines as broader retention improves.
Schedule and continual-learning ablations expose the tradeoff structure. With an average configured mixture of 0.5%, ending the mixture halfway through training improves final coding and SciQ accuracy but preserves less held-out-loss reduction than constant mixing. After mathematical training, all tested unified mixtures from 0.5% to 5% achieve higher final coding accuracy on HumanEval+ and MBPP+ than the staged model; the 5% mixture is best on all four reported coding metrics, while smaller mixtures do not consistently improve HumanEval+ coverage. At constant 2% mixing, REINFORCE favors Math accuracy and retention, whereas MaxRL favors Math coverage.
Figure 4
Caption: Figure 4: Final performance and generalization. Mean pass@1 over two seeds after 6000 RL updates in Base and 12000 in Extended. Colors indicate midtraining tokens mixed during pretraining; the horizontal axis measures passes over the same fixed corpus during RL. Dashed lines show the best staged configuration on each setting’s primary test (Appendix C.6). Extended generalization evaluates the same Extended models on a larger puzzles excluded from RL training (Appendix C).
Why It Matters: It directly compares final unified configurations with staged training across task performance and larger-puzzle generalization.
Figure 5
Caption: Figure 5: Accuracy and coverage across training compute budgets. The top and bottom rows show pass@1 and pass@16 against cumulative compute, including pretraining, midtraining, and RL. Colored curves show unified configurations: darker colors indicate more midtraining compute before RL, while curves of the same color vary the midtraining mixture during RL. Gray dashed curves show all six staged configurations, spanning three midtraining budgets and two learning rates (Appendix C.6). Open circles mark the pre-RL evaluations in Appendix C.5. The black curve connects the best observed checkpoints for each metric, with colored or gray markers identifying unified or staged training. RL points av...
Why It Matters: It separates final quality from compute efficiency and shows where staged training remains competitive.
Table 1
Caption: Table 1: Data sources for the shared likelihood objective. For binary rewards, ν+ is the success-conditioned distribution of the generation policy ν; ν = πθt recovers on-policy MaxRL.
| Source | Prompts \(\rho_{s,t}\) | Target \(\mu_{s,t}\) |
|---|---|---|
| Pretraining | empty | \(\mu_{\mathrm{pre}}\) |
| Midtraining | empty | \(\mu_{\mathrm{mid}}\) |
| SFT | \(\rho_{\mathrm{sft}}\) | \(\mu_{\mathrm{sft}}\) |
| RL | \(\rho_{\mathrm{rl}}\) | \(\pi^+_{\theta_t}\) or \(\nu^+\) |
Why It Matters: The table makes explicit that the stages differ in their data distributions, not necessarily in the optimizer or loss.
Table 2
Caption: Table 2: RL objectives within the unified likelihood framework. Different objectives correspond to different target and effective prompt distributions.
| Method | Reward | Target \(\mu_s\) | Effective prompts \(\rho_s\) |
|---|---|---|---|
| MaxRL | Binary | \(\pi^+_{\theta_t}\) | \(\rho(x)\) |
| REINFORCE | Binary | \(\pi^+_{\theta_t}\) | \(\propto \rho(x)p_{\theta_t}(x)\) |
| TailRL | \([0,1]\) | \(\int_0^1 \pi_{\theta_t}(\cdot\mid x,r>\tau)\,d\tau\) | \(\rho(x)\) |
| REINFORCE | \([0,1]\) | \(\frac{\pi_{\theta_t}(y\mid x)r(x,y)}{\mathbb{E}_{y'\sim\pi_{\theta_t}(\cdot\mid x)}[r(x,y')]}\) | \(\propto \rho(x)\mathbb{E}_{y\sim\pi_{\theta_t}(\cdot\mid x)}[r(x,y)]\) |
Why It Matters: It separates the target completion distribution from the effective prompt distribution, clarifying why different RL objectives can share the same trainer while producing different updates.
🔮 Conclusion
The paper presents a coherent mathematical and systems-level case for treating foundation-model training as likelihood fitting over an evolving mixture of data sources. Its controlled and heterogeneous experiments suggest that overlapping data sources can improve task performance, generalization, retention, and later learning compared with abrupt stage transitions.
The practical conclusion is conditional rather than universal: unified training creates a more expressive design space, but the best result depends on the source mixture, schedule, RL objective, compute budget, and target capability. The reported gains are strongest when evaluation includes coverage, generalization, retention, or subsequent adaptation rather than only one immediate task score.
🛠️ Future Research Improvements
The paper's results motivate larger-scale studies that jointly sweep model size, source proportions, data reuse, synchronization lag, and RL duration under compute-matched conditions. Such studies should distinguish parameter count, training compute, inference-time sampling, and wall-clock throughput.
Further work should test more RL objectives and reward types, especially continuous, noisy, or imperfectly verifiable rewards, and should determine when unified mixing improves coverage without sacrificing exact-match accuracy. The schedule results also call for principled methods that choose when to introduce, decay, or remove domain data.
Reproduction should extend beyond reasoning puzzles and the reported mathematics-and-coding setting to broader instruction, tool-use, multilingual, and long-context workloads while measuring both immediate performance and continual-learning transfer.
🏭 Potential Industry Use Scenarios
A general language model could continue learning from verified task completions while retaining general text exposure, making the approach relevant to domains such as mathematical reasoning, code generation, structured extraction, and other workflows with automatic validators.
An asynchronous generator–trainer architecture could support production training systems in which evaluation or rollout workers generate reward-labeled data while optimization workers consume a shared buffer. The paper provides evidence for the architecture in research settings, but operational deployment would still require controls for data freshness, reward quality, privacy, and regression testing.
Mixture schedules could be used to manage specialization-versus-retention tradeoffs: constant mixing for stronger corpus retention, or decaying schedules when later-task accuracy is more important. These are plausible engineering applications of the paper's framework, not deployment results established by the experiments.
💬 Critical Analysis
The strongest aspect is the alignment between theory and implementation. The mixture formulation explains why static data and selected generations can share one objective, and the token-weighted data format makes that abstraction concrete. The experiments also evaluate more than immediate task accuracy: they examine coverage, larger-puzzle generalization, pretraining retention, schedule effects, and later coding RL.
The central claim should nevertheless be interpreted as a unification of objectives and infrastructure, not proof that stage boundaries are always harmful. The controlled experiments are small and highly structured, the real-world study remains below frontier scale, and several conclusions depend on tuned mixture proportions. Staged training wins at some intermediate compute budgets, and heterogeneous mixing can lower coding accuracy even when it improves coverage or retention.
For builders, the most important open question is whether unified training produces robust gains after scaling and after replacing clean verifiable rewards with the noisy feedback common in deployed systems. Reproduction should therefore preserve compute matching, report pass@1 separately from pass@16, compare multiple RL objectives, and measure broad-capability regressions rather than relying on a single aggregate score.