Training Foundation Models with a Single Algorithm
Organizations: Carnegie Mellon University · Johns Hopkins University · Aurora Innovation
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.
Figures & tables
| Source | Prompts | Target |
|---|---|---|
| Pretraining | empty | |
| Midtraining | empty | |
| SFT | ||
| RL | or |
| Method | Reward | Target | Effective prompts |
|---|---|---|---|
| MaxRL | Binary | ||
| REINFORCE | Binary | ||
| TailRL | |||
| REINFORCE |
Appendix figures & tables42 assets
Supplementary material from the paper’s appendix.
Appendix
| Midtraining sequences per update, | Midtraining token fraction |
|---|---|
| 8 | 3.03% |
| 16 | 5.88% |
| 32 | 11.11% |
| 64 | 20.00% |
| 128 | 33.33% |
| Availability | Example reference |
|---|---|
| Included in midtraining | Ava’s best friend |
| The best friend of Ava’s best friend | |
| The person directly above Ava | |
| The person on the second-highest floor | |
| The best friend of the person on floor 19 | |
| Introduced during RL | The best friend of the person on the highest floor |
| People | Omitted steps | Base | Extended | |
| RL prompts: demonstrated patterns | 8–9 | 0 | 10000 | 10000 |
| RL prompts: omitted compositions | 8–9 | 2–4 | 18570 | 11666 |
| RL prompts: larger puzzles | 10–11 | 2–4 | – | 10000 |
| Primary test questions | 8–9 | 2–4 | 1024 | – |
| Primary test questions | 10–11 | 2–4 | – | 1024 |
| Size-generalization test questions | 12 | 2–4 | – | 1024 |
| Setting | Value |
|---|---|
| Parameters / layers | 123587328 / 12 |
| Hidden / feedforward size | 768 / 2048 |
| Attention heads | 12 |
| Embeddings and normalization | Tied embeddings, rotary position embeddings, RMSNorm |
| Tokenizer / context length | GPT-2 / 2048 tokens |
| Global batch size | 256 packed sequences of 2048 tokens |
| Setting | Staged midtraining | RL (both approaches) |
|---|---|---|
| Optimizer | AdamW | AdamW |
| Adam betas | ||
| Learning rate | peak | |
| Schedule | 5% warmup, then cosine decay | Constant, no warmup |
| Weight decay | 0.01 | 0.01 |
| Gradient clip | 1.0 | 1.0 |
| Demonstrated patterns | Base | Extended | Extended generalization | |||||
| Midtraining | pass@1 | pass@16 | pass@1 | pass@16 | pass@1 | pass@16 | pass@1 | pass@16 |
| None | 0.00 | 0.00 | 0.00 | 0.00 | 0.00 | 0.00 | 0.00 | 0.00 |
| Unified: mixed during pretraining | ||||||||
| 133M tokens | 0.92 | 12.11 | 0.12 | 1.76 | 0.00 | 0.00 | 0.00 | 0.00 |
| 267M tokens | 7.23 | 55.98 | 0.34 | 4.67 | 0.15 | 0.00 | 0.00 | |
| 534M tokens | 16.74 | 79.67 | 0.55 | 7.01 | 0.05 | 0.81 | 0.02 | |
| Base | Extended | Extended generalization | |||||
|---|---|---|---|---|---|---|---|
| Epochs | Peak LR | pass@1 | pass@16 | pass@1 | pass@16 | pass@1 | pass@16 |
| 1 | 21.79 | 72.40 | 16.14 | 63.51 | 6.58 | 41.78 | |
| 1 | 12.24 | 54.54 | 9.88 | 51.30 | 3.18 | 26.31 | |
| 2 | 24.54 | 73.51 | 16.91 | 64.80 | 5.58 | 35.90 | |
| 2 | 12.74 | 55.25 | 7.25 | 41.13 | 2.17 | 18.86 | |
| 4 | 15.15 | 57.49 | 11.64 | 53.20 | 4.10 | 28.53 | |
| Midtraining | Pretraining loss | Benchmark mean (%) | ||||
|---|---|---|---|---|---|---|
| Setting | Tokens during pretraining | Examples per RL update | Without | With | Without | With |
| Base | 534M | 8 | 2.570 | 2.542 | 41.66 | 42.13 |
| Base | 534M | 32 | 2.575 | 2.535 | 41.65 | 41.98 |
| Base | 2.14B | 8 | 2.572 | 2.545 | 41.93 | 42.08 |
| Base | 2.14B | 32 | 2.571 | 2.537 | 42.01 | 41.85 |
| Extended | 534M | 8 | 2.613 | 2.546 | 41.60 | 42.25 |
| Held-out loss | Benchmark score (%) | ||||||||
| Midtraining | ClimbMix | Dolmino | Mean | MMLU | ARC-E | HellaSwag | TQA | WinoG | Mean |
| None | 2.784 | 1.794 | 2.537 | 28.65 | 55.09 | 35.39 | 39.15 | 53.59 | 42.37 |
| Unified: mixed during pretraining | |||||||||
| 133M tokens | 2.784 | 1.795 | 2.537 | 28.62 | 54.67 | 35.55 | 39.09 | 53.20 | 42.23 |
| 267M tokens | 2.786 | 1.796 | 2.538 | 28.79 | 54.97 | 35.17 | 39.17 | 53.99 | 42.42 |
| 534M tokens | 2.785 | 1.797 | 2.538 | 28.22 | 54.67 | 35.67 | 39.57 | 51.14 | 41.85 |
| Setting | Pretraining | Midtraining |
|---|---|---|
| Optimizer | AdamW | |
| Adam ; | ; | |
| Weight decay; gradient clipping norm | 0.1; 1 | |
| Global batch size (sequences) | 2048 | |
| Sequence length (tokens) | 2048 | |
| Peak / minimum learning rate | / | |
| Setting | Unified | Staged RL |
|---|---|---|
| Optimizer | AdamW | AdamW |
| Adam | ||
| Adam | ||
| Weight decay | 0 | 0 |
| Gradient clipping norm | 1 | 1 |
| Learning rate |
| Benchmark | Pass@1 | Pass@32 |
|---|---|---|
| MATH-500 | 6.54 | 43.76 |
| GSM8K | 17.84 | 71.63 |
| Math (average) | 12.19 | 57.69 |
| HumanEval+ | 11.97 | 32.32 |
| MBPP+ | 21.05 | 57.14 |
| Coding (average) | 16.51 | 44.73 |
| Mixing | MATH pass@1 | GSM pass@1 | MATH pass@64 | GSM pass@64 | MATH pass@256 | GSM pass@256 |
|---|---|---|---|---|---|---|
| 0% | 16.83 | 32.68 | 66.73 | 83.88 | 78.20 | 91.81 |
| 0.1% | 16.90 | 32.14 | 65.75 | 85.56 | 77.00 | 93.03 |
| 0.2% | 16.56 | 32.35 | 65.03 | 84.74 | 76.20 | 92.49 |
| 0.5% | 16.83 | 31.71 | 65.60 | 84.06 | 77.20 | 91.81 |
| 1% | 16.11 | 31.61 | 65.85 | 83.71 | 76.00 | 91.74 |
| 2% | 16.68 | 31.30 | 67.51 | 84.90 | 78.40 | 93.03 |
| Mixing | loss | MMLU | SciQ | ARC-C | MBPP+ | HumanEval+ |
|---|---|---|---|---|---|---|
| 0% | 37.12 | 50.20 | 27.90 | 25.41 | 15.17 | |
| 0.1% | 39.30 | 54.60 | 29.69 | 24.82 | 14.44 | |
| 0.2% | 40.04 | 54.30 | 31.06 | 23.71 | 13.76 | |
| 0.5% | 39.99 | 62.00 | 32.34 | 21.05 | 11.97 | |
| 1% | 40.54 | 66.10 | 33.53 | 19.31 | 11.66 | |
| 2% | 40.66 | 67.00 | 34.22 | 19.23 | 11.55 |
| Schedule | Math pass@1 | Math pass@32 | Coding pass@1 | Coding pass@32 | loss | MMLU |
|---|---|---|---|---|---|---|
| Constant 0.5% | 24.27 | 68.32 | 16.51 | 46.38 | 39.99 | |
| Full anneal 1% | 24.70 | 68.94 | 17.76 | 46.16 | 39.96 | |
| Ramp to 1% | 24.01 | 68.94 | 16.33 | 46.55 | 40.17 | |
| Half anneal 2% | 24.80 | 68.49 | 18.85 | 46.93 | 40.14 |
| Setting | MATH | GSM | loss | SciQ | MBPP+ | HumanEval+ |
|---|---|---|---|---|---|---|
| constant: original | 16.83 | 31.71 | 62.00 | 21.05 | 11.97 | |
| constant: code removed | 16.05 | 32.15 | 61.90 | 23.66 | 13.64 | |
| constant: code replaced | 16.54 | 31.22 | 64.80 | 20.80 | 12.12 | |
| anneal: original | 15.50 | 32.55 | 59.00 | 23.69 | 14.41 | |
| anneal: code removed | 16.40 | 32.07 | 56.80 | 24.01 | 14.33 | |
| anneal: code replaced | 17.00 | 31.81 | 61.70 | 23.38 | 13.11 |
| Method | Math pass@1 | Math pass@32 | Coding pass@1 | Coding pass@32 | loss | MMLU | SciQ | ARC-C |
|---|---|---|---|---|---|---|---|---|
| MaxRL | 23.99 | 69.61 | 15.39 | 45.51 | 40.66 | 67.00 | 34.22 | |
| REINFORCE | 24.81 | 67.70 | 15.80 | 45.21 | 40.69 | 71.70 | 36.43 |
| Mixing setting | Method | Math pass@1 | Math pass@32 | Coding pass@1 | Coding pass@32 | loss | MMLU | SciQ | ARC-C |
|---|---|---|---|---|---|---|---|---|---|
| Constant 0.5% | MaxRL | 24.27 | 68.32 | 16.51 | 46.38 | 39.99 | 62.00 | 32.34 | |
| REINFORCE | 22.98 | 65.29 | 14.45 | 46.00 | 41.13 | 73.90 | 38.82 | ||
| Constant 1% | MaxRL | 23.86 | 68.25 | 15.49 | 44.89 | 40.54 | 66.10 | 33.53 | |
| REINFORCE | 24.51 | 66.41 | 15.60 | 44.24 | 40.82 | 70.30 | 35.58 | ||
| Constant 2% | MaxRL | 23.99 | 69.61 | 15.39 | 45.51 | 40.66 | 67.00 | 34.22 | |
| REINFORCE | 24.81 | 67.70 | 15.80 | 45.21 | 40.69 | 71.70 | 36.43 |
| Setting | Value |
|---|---|
| Algorithm | On-policy MaxRL with reward-baseline subtraction |
| Optimizer | Fresh AdamW |
| Adam ; | ; |
| Weight decay; gradient clipping norm | 0; 1 |
| Learning rate / schedule | / constant, no warmup |
| Prompts per update | 32 |
| HumanEval+ | MBPP+ | |||
|---|---|---|---|---|
| Mathematical training | Pass@1 | Pass@16 | Pass@1 | Pass@16 |
| Staged | 22.94 | 45.12 | 29.81 | 52.12 |
| Unified 0.5% | 23.67 | 43.90 | 30.31 | 52.91 |
| Unified 1% | 23.70 | 45.12 | 29.94 | 54.76 |
| Unified 2% | 24.09 | 44.51 | 31.76 | 56.88 |
| Unified 5% | 24.96 | 46.95 | 32.69 | 58.47 |
LLM Training
Large language model training studies the data, objectives, optimizers, and systems used to build and adapt language models.
See all 353 papers