Training Foundation Models with a Single Algorithm

AI summary

Overview

Research area: Machine learning — foundation-model training pipelines, reinforcement learning, and training infrastructure.

Technical level: Intermediate (the paper is theoretically grounded but its central idea can be understood without advanced math).

Scope: The paper argues that pretraining, midtraining, supervised fine-tuning and reinforcement learning can all be run as one maximum-likelihood training algorithm over an evolving mixture of static data and model-generated data, and tests this claim across a controlled reasoning setting (124M parameters) and a real-world math/code setting (0.76B parameters).

What This Paper Is About

Foundation models are normally trained in separate stages — pretraining, midtraining, supervised fine-tuning (SFT), then reinforcement learning (RL) — and each stage typically uses its own data and its own learning algorithm. The authors show this algorithmic separation is not fundamental: by rewriting suitable RL objectives as maximum-likelihood fitting to a target distribution, every stage reduces to the same loss. The goal is a "unified training" system where the only design choices left are which data to train on, in what proportions, and when — with conventional staged training as a special case.

Key Contributions

1. A unified formulation of foundation-model training. The authors establish a common maximum-likelihood formulation covering pretraining, midtraining, SFT and suitable RL objectives, building on the likelihood interpretation of Maximum Likelihood Reinforcement Learning (MaxRL). Staged training becomes the special case where all mixture weight sits on one data source at a time.
2. Extensions beyond on-policy MaxRL. The framework handles generations produced by an arbitrary fixed policy ν (not only the model currently being optimized), and accommodates other objectives — REINFORCE, continuous-reward REINFORCE, and TailRL — through different choices of target distribution and effective prompt distribution.
3. A unified training infrastructure. A shared data format (token-level weights), a single weighted cross-entropy loss, and one trainer that samples static data and RL generations from a common buffer, with generation and optimization running asynchronously and periodic weight synchronization.
4. Empirical study of data mixing and compute allocation. Controlled and real-world experiments characterize how mixing proportions and schedules shape the tradeoffs among accuracy, coverage and retention, and how compute should be spread across pretraining, midtraining and RL.

Main Findings

• Better task performance and generalization (controlled setting). The best unified configurations outperform staged training in Base, in Extended, and in Extended generalization, which evaluates models on larger puzzles unseen during either midtraining or RL. The controlled study trains a 124M-parameter model from scratch on a fixed 20B-token pretraining corpus and a fixed 510M-token midtraining corpus, using Reasoning Gym-style puzzles, MaxRL with 6000 updates for Base and 12000 for Extended, pass@1 as accuracy and pass@16 as coverage, averaged over two training seeds.
• More data reuse is tolerated. Higher-intensity mixing of midtraining tokens during pretraining performs better; the two best-performing configurations perform 2 and 4 epochs over the midtraining corpus during pretraining. Staged training, by contrast, loses final task accuracy when midtraining is extended beyond its best-performing budget. The benefit of mixing midtraining data during RL is similar but less consistent — at the largest pretraining–midtraining mixture (2.14B tokens) Base peaks at an intermediate RL mixture, whereas Extended and Extended generalization improve with more midtraining data during RL.
• Compute-optimality depends on the metric and budget. Staged training is compute-optimal for accuracy at intermediate budgets, while unified training achieves higher accuracy at larger budgets and at the smallest budget in both Base and Extended. Staged training is more compute-optimal for coverage. On Extended generalization, unified training achieves the best accuracy and coverage for almost all compute budgets. Compute comparisons include pretraining, midtraining and RL.
• Better retention of general capabilities. Across all five unified pretraining–midtraining mixtures, unified training keeps retention close to the pretrain-only model even as midtraining data grows, while the best-retention staged configuration (one epoch, peak learning rate 5×10⁻⁵) has substantially higher held-out pretraining loss and lower mean score on five general benchmarks. The staged configuration with the best final RL accuracy does even worse on both retention measures.
• Evaluation timing matters. Before RL, the staged reference has higher Base accuracy than any unified initialization, yet the best unified models surpass it after RL — suggesting midtraining should be judged by how it supports later learning.
• Real-world mixture tradeoff. With a 0.76B-parameter model trained from scratch, Math12K, and a 240M-token corpus matching the midtraining composition (web text, mathematics, code, reasoning traces), sweeping the midtraining proportion from 0 to 5% leaves Math accuracy roughly similar with very small degradation, gives the highest coverage at the largest proportions, and lowers Coding accuracy even though Coding coverage exceeds both the base model and staged training. All nonzero mixtures reduce held-out midtraining loss below the base model (larger proportions give lower loss), while staged training increases it; SciQ accuracy generally improves.
• Schedules trade off differently. Comparing constant mixing, linear decay, linear growth, and decay to zero at the midpoint (average configured midtraining proportion 0.5%), ending mixing halfway improves Coding accuracy and preserves SciQ gains but retains less of the corpus-loss reduction; Math accuracy and coverage are similar across all four.
• Better continual-learning starting point. After the same coding RL (without midtraining data) for 4000 updates, all unified models achieve higher final accuracy on HumanEval+ and MBPP+ and higher MBPP+ coverage than staged training, despite lower initial coding accuracy; the 5% mixture is best on all four metrics, while HumanEval+ coverage at smaller mixtures is similar to or below staged. The finding is further verified with WebInstruct in a general reasoning domain.
• The RL objective still matters. At constant 2% mixing, REINFORCE gives higher Math accuracy, lower held-out loss, and higher SciQ and ARC-Challenge accuracy, while MaxRL gives higher Math coverage; these coverage and retention differences persist across proportions and schedules.
• KL regularization does not reproduce the benefit. Increasing the KL penalty toward the base model reduces deterioration in held-out loss, but the strongest penalty substantially reduces Math accuracy and coverage; no tested penalty matched unified training's combination of mathematical performance, lower corpus loss, and higher SciQ and MMLU accuracy.

Methodology in Plain English

The authors start from a simple observation: pretraining, midtraining and SFT are already the same algorithm — minimize negative log-likelihood on example text. RL looks different because it is written as reward maximization. Building on MaxRL, they show that an RL update can be rewritten as maximum-likelihood fitting to the current model's success-conditioned distribution — the distribution of completions that earned reward. That puts RL in exactly the same form as the other stages.

With every source expressed as a prompt distribution plus a target distribution, the whole pipeline becomes one loss over a weighted mixture of sources, where the weights can change over training. Conventional staged training is the case where the weights are one-hot and switch abruptly.

For implementation, each training example is a token sequence plus a per-token weight: corpus and demonstration tokens get unit weight, prompt tokens get zero, and RL completion tokens get a response weight. The generator samples N completions per prompt, estimates the pass rate p̂, and assigns weight 1/p̂ to correct completions and zero to incorrect ones — the finite-sample form of MaxRL. Generation and training run asynchronously through a shared buffer, with the trainer publishing weights every k updates.

They then evaluate in two settings: a controlled one with a 124M model and synthetic reasoning puzzles with verifiable answers and controllable difficulty (so they can measure Base, Extended, and generalization to larger puzzles), and a real-world one with a 0.76B model, heterogeneous midtraining data, and math and coding RL, where Math is the average of MATH-500 and GSM8K and Coding the average of HumanEval+ and MBPP+.

Why This Matters

Impact on research. The paper reframes foundation-model training as a question about training trajectories — which data, in what proportions, and when — rather than a sequence of separately studied stages. That makes it possible to judge midtraining, SFT and RL data by their effect on the final model instead of stage-local metrics, and it connects scaling-law research on individual stages to the question of how decisions across the whole pipeline interact.

Real-world applications:

• Training pipelines for reasoning models that must both learn new verifiable tasks and keep general knowledge intact.
• Continual learning and model refresh cycles, where an updated model should be a strong starting point for the next round of fine-tuning and RL.
• Data- and compute-constrained training runs where reusing a fixed corpus for multiple epochs is unavoidable and forgetting must be controlled.
• Mixture-schedule design for code and math post-training, where accuracy on one skill can degrade while coverage improves.

Industry relevance. The system is deliberately built as infrastructure — one data format, one loss, one trainer, asynchronous generation — which is directly relevant to teams that currently maintain separate codebases and recipes for SFT and RL. The results suggest that careful mixture design can trade off task accuracy, coverage and retention more flexibly than switching between entirely different training stages.

Future Directions

• Other RL objectives. The authors explicitly defer systematic study of RL objectives other than MaxRL to future work, presenting only a preliminary analysis at constant 2% mixing; how the objective interacts with the data mixture across proportions and schedules is left open.
• Schedule design principles. The preferred mixing schedule depends on how performance across tasks is balanced, and the authors state they leave a detailed analysis of these interactions to future work.
• A science of training trajectories. The paper positions unified training as a starting point for principles governing how data and compute should evolve over a model's lifetime, beyond the specific recipes tested here.
• Off-policy generation and staleness. The framework supports fitting the success-conditioned distribution of any generator, not just the current model; how far generation may drift from optimization (and the resulting staleness and throughput tradeoff) remains a design question rather than a settled result.

Target Audience

Machine-learning researchers and engineers working on post-training, RL for language models, or large-scale training infrastructure; practitioners designing data mixtures and compute allocation across pretraining, midtraining and RL; and students who want a compact theoretical bridge between likelihood-based training and policy-gradient methods. Readers seeking purely empirical benchmark rankings will find a controlled, small-scale study rather than frontier-scale results — the largest model reported is 0.76B parameters.

Authors’ abstract

Foundation models are typically trained in stages: pretraining, midtraining, supervised fine-tuning (SFT), and reinforcement learning (RL), with changes in both data and learning algorithms. We show that this algorithmic separation is not fundamental. By deriving suitable RL objectives as maximum-likelihood fitting to model-generated target distributions, we place all of these stages under a single likelihood-based training algorithm, unified training. This perspective turns the foundation-model training lifecycle into a problem of choosing which data to train on, in what proportions, and when, with conventional staged training as a special case. We instantiate this framework in a unified training system that mixes static data and RL generations within a single trainer. Across controlled and real-world experiments, unified training improves the accuracy-compute tradeoff and task coverage while preserving broader capabilities, and provides a stronger starting point for continual learning. This work lays the groundwork for a unified science of foundation-model training, connecting mathematical principles, training infrastructure, and the development of model capabilities.