PaperlayerView exact revision on arXiv

Source-linked AI summary

Training Foundation Models with a Single Algorithm

Yuda Song, Alvin Zhang, Fahim Tajwar, Daman Arora, Jeff Schneider, Daniel Khashabi, Ruslan Salakhutdinov, J. Andrew Bagnell, Aarti Singh, Andrea Zanette

arXiv:2610.40564v1cs.LG

TL;DR

Foundation-model training is usually divided into stages with distinct recipes, while the effects of decisions made across training on final capabilities remain less understood. This paper puts pretraining, midtraining, SFT, and suitable RL under a shared likelihood-based framework and reports improved accuracy and coverage over staged training while preserving broader capabilities.

  • Problem

    Foundation models are commonly trained in separate stages, and less is known about how decisions throughout training jointly determine final-model capabilities.

    Source: §1 · p. 1

  • Method

    The paper derives likelihood targets for suitable RL objectives and combines static data with model generations under a common maximum-likelihood training rule.

    Source: §1 · p. 2

  • Results

    Unified training improves accuracy and coverage over staged training in controlled and real-world experiments while preserving broader capabilities.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The framework makes evolving data mixtures and their timing central training choices, providing a basis for studying the full training process.

    Source: §1 · p. 1, §7 · p. 13

    Experiments remain small-scale, identify no universally preferred mixture or schedule, and do not establish predictive scaling laws for compute allocation.

    Source: §7 · p. 13

Abstract

from arXiv · show

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.

1 Introduction

The paper argues that stage-specific recipes are not fundamental: pretraining, midtraining, SFT, and suitable RL can share a maximum-likelihood formulation. It reframes training around evolving data mixtures and reports empirical benefits over staged training.

1 Introduction

  • The unified view makes the evolving data distribution the central design choice: which sources to use, in what proportions, and when.Pretraining text, demonstrations, and model generations selected with verifiable rewards become sources for the same learning process.Source: §1 · p. 1
  • The paper identifies a gap in understanding how decisions made throughout training jointly determine final-model capabilities.A unified formulation allows decisions to be studied together by their effect on the final model rather than by stage-specific metrics.Source: §1 · pp. 1–2
  • Pretraining, midtraining, SFT, and suitable RL objectives share a maximum-likelihood formulation, with staged training as a special case.The formulation supports training on changing mixtures of static and generated data.Source: §1 · p. 2
  • Unified training improves accuracy and coverage over staged training in controlled and real-world experiments, while preserving broader capabilities.The framework figure reports higher accuracy and coverage on a reasoning task within the same total compute budget.Source: §1 · p. 2, Figure 1 · p. 2
  • By placing data composition, timing, and compute allocation in one framework, unified training offers a basis for studying principles of full training trajectories.The paper characterizes mixing proportions, schedules, and compute allocation across pretraining, midtraining, and RL.Source: §1 · p. 2

2 Preliminaries

The preliminaries define maximum-likelihood fitting for prompt-completion data and describe reinforcement learning as reward-based optimization of model-generated completions.

2.1 Maximum Likelihood Estimation

  • Maximum likelihood estimation fits prompt-completion examples by minimizing negative log-likelihood; pretraining, midtraining, and SFT apply it to different data sources.The fixed corpora and demonstrations used in these stages are called static data.Source: §2.1 · pp. 2–3
  • For a fixed training-data distribution, the loss and expected forward KL divergence differ only by a parameter-independent constant and have the same gradient.Source: §2.1 · p. 3

2.2 Reinforcement Learning

  • Reinforcement learning samples model completions for prompts and optimizes their reward; REINFORCE estimates expected reward, while PPO and GRPO add further update choices.The described reward takes values in {0, 1}.Source: §2.2 · p. 3

3 A Unified Likelihood View of Foundation-Model Training

Pretraining, midtraining, SFT, and suitable RL objectives can be expressed as likelihood fitting. Unified training keeps that learning rule fixed while the data sources and their mixture weights evolve.

3.1 MaxRL as Likelihood Fitting

  • MaxRL fits the current policy’s success-conditioned distribution as a likelihood target, placing RL in the same training form as pretraining, midtraining, and SFT.The target is held fixed during each update.Source: §3.1 · p. 3

3.2 Unified Training

  • Unified training represents each source with prompt and target distributions, then applies the same likelihood loss to their weighted mixture.Source: §3.2 · pp. 3–4
  • Both data sources and mixture weights may evolve, while conventional staged training is the special case of assigning all mixture weight to one source at a time.Source: §3.2 · p. 4
  • The framework also expresses other RL objectives through different target and effective prompt distributions, including REINFORCE and continuous-reward objectives such as TailRL.Source: §3.2 · p. 4, Table 2 · p. 4
  • The framework supports generations from a fixed policy other than the optimized model, though the resulting update is generally not its current-policy MaxRL gradient.Source: §3.2 · p. 4

4 Practical Implementation

The implementation uses a shared data representation, likelihood loss, and trainer for static data and RL generations. A shared buffer supports their mixture, while generation and optimization proceed asynchronously.

4 Practical Implementation

  • A shared system combines static examples and RL generations in one buffer, with asynchronous generation and optimization and periodic weight synchronization.Source: §4 · p. 4, Figure 2 · p. 5

4.1 Unified Loss and Data Format

  • Training examples use token sequences with per-token weights, allowing pretraining, midtraining, SFT, and RL to share a weighted likelihood format.Corpus tokens receive unit weights; SFT masks prompt tokens, while RL weights response tokens according to reward and sampled pass rate.Source: §4.1 · p. 5, Figure 2 · p. 5

4.2 Data Generation

  • For binary rewards, the generator samples completions from a recent model copy and weights correct responses by the inverse sampled pass rate.Incorrect completions receive zero weight.Source: §4.2 · p. 5
  • Finite-sample weighting preserves the success-conditioned target but down-weights prompts rarely solved by the generating policy; with one sample, it recovers REINFORCE.Source: §4.2 · p. 5

4.3 Data Mixing and Asynchronous Training

  • The trainer mixes static data and RL generations from a shared buffer and samples sources according to target token fractions adjusted for average example length.The trainer publishes weights every k updates, and generators load them between batches.Source: §4.3 · p. 5

5 Experiments

The experiments compare unified and staged training in controlled reasoning tasks and with heterogeneous real-world midtraining data. They examine performance, compute, retention, data mixing, and continual learning.

5.1.1 Unified Training can Obtain Better Task Performance and Generalization

  • Best unified configurations outperform staged training in Base and Extended task performance and on Extended generalization to larger puzzles unseen during midtraining and RL.The controlled runs use verifiable reasoning puzzles and compare final results after fixed RL-update budgets.Source: §5.1.1 · pp. 6–7, Figure 4 · p. 7

5.1.2 Unified Training can be Compute-optimal in More Settings

  • At matched compute, staged training is more accuracy-optimal at intermediate budgets and more coverage-optimal overall, while unified training leads at the smallest and larger accuracy budgets and almost all generalization budgets.The comparison includes compute from pretraining, midtraining, and RL.Source: §5.1.2 · p. 7, Figure 5 · p. 8

5.1.4 Additional Analyses

  • Pre-RL accuracy can mislead: staged initialization scores higher on Base, yet the best unified models finish ahead; separately normalized losses often help Extended but require tuning, and pretraining replay trades lower loss for reduced reasoning results in most comparisons.The shared loss’s relative source contributions follow the token mixture.Source: §5.1.4 · pp. 8–9

5.1.3 Unified Training can Retain General Capabilities Better

  • Unified mixtures maintain retention near pretraining-only levels, while the best staged configuration has substantially higher held-out loss and lower mean benchmark performance.Retention is assessed before RL using held-out pretraining loss and mean performance on five general benchmarks.Source: §5.1.3 · p. 8, Figure 6 · p. 9

5.2.1 Data Mixtures Tradeoff Accuracy, Coverage, and Retention

  • With heterogeneous midtraining data, larger nonzero mixtures improve retention and coverage, while Math accuracy changes little and Coding accuracy falls as mixing increases.Staged training raises held-out midtraining loss; SciQ accuracy generally improves with mixing.Source: §5.2.1 · p. 9, Figure 7 · p. 10

5.2.2 Data Mixing Schedules

  • Ending midtraining-data mixing halfway through RL improves Coding and SciQ accuracy over constant mixing while retaining less of the held-out corpus-loss reduction; Math accuracy and coverage remain similar across schedules.The schedules have an average configured midtraining proportion of 0.5%.Source: Figure 8 · p. 10, §5.2.2 · p. 10

5.2.3 Unified Training can Enable Better Continual Learning

  • After mathematical training, all unified models finish with higher coding accuracy on both benchmarks than staged training, and the 5% mixture leads on all four reported metrics.All models use the same subsequent coding RL without midtraining data.Source: Figure 9 · p. 11, §5.2.3 · p. 11

5.2.4 Additional Analyses

  • At fixed mixture, REINFORCE favors Math accuracy and retention while MaxRL favors Math coverage; tested KL penalties do not match unified training’s combination of outcomes, and corpus changes affect Coding accuracy differently by schedule.The corpus changes include removing designated code sources or replacing them with verified coding examples.Source: §5.2.4 · p. 11

6 Related Work

Related work studies foundation-model training stages, transfer, hybrid reinforcement learning, and asynchronous systems. This paper connects those threads with a common likelihood formulation and implementation for static and generated data.

6.1 Foundation-Model Training

  • Prior foundation-model systems combine pretraining with supervised learning, preference optimization, and RL in increasingly elaborate pipelines; this work instead studies their organization through a common formulation and implementation.The formulation allows data composition and timing to vary without a separate training algorithm for each stage.Source: §6.1 · pp. 11–12

6.2 Understanding Foundation-Model Training

  • Scaling, transfer, domain adaptation, and reasoning studies examine performance across model size, compute, training data, and stages; this work complements them with a common likelihood formulation for static and RL-generated data.The formulation supports studying composition, timing, and reuse across training.Source: §6.2 · p. 12

6.3 Hybrid Offline and Online Reinforcement Learning

  • Hybrid offline-online RL combines demonstrations or prior experience with interaction, while language-model methods often add a separately weighted supervised term; unified training instead sets both sources’ relative weight through one data mixture.Related approaches include InstructGPT’s added pretraining likelihood term and methods with tuned or scheduled coefficients.Source: §6.3 · p. 12

6.4 Connections between Reinforcement Learning and Supervised Learning

  • Work linking RL and supervised learning includes reward-guided maximum-likelihood training of generated responses; this formulation places static data and suitable RL targets in one likelihood framework, data format, loss, and trainer.The formulation builds on MaxRL and TailRL.Source: §6.4 · p. 12

6.5 Asynchronous Reinforcement Learning

  • Asynchronous RL systems separate actors from learners and manage policy lag in different ways; language-model systems likewise connect generation and training through buffers or streams.Examples discussed include bounding policy versions or updating generator weights during generation.Source: §6.5 · pp. 12–13

7 Discussion

Unified training makes the training rule common while allowing data sources and their proportions to change over time. The discussion highlights practical benefits but leaves systematic mixture selection and larger-scale evaluation open.

7 Discussion

  • Experiments show practical benefits, while a data source’s value depends on when it is introduced and how it interacts with later learning.The authors do not identify a universally preferred mixture or schedule, establish predictive scaling laws, or deeply study how optimization should change with mixture design.Source: §7 · p. 13
  • The experiments remain small-scale, mostly train one RL task at a time, and do not eliminate separate downstream adaptation stages.Longer task sequences are left for future study; preceding training provides initial evidence of improved adaptation.Source: §7 · p. 13

A.1 Data Mixing

  • Data mixtures specify how much data from each source enters a batch, using added examples in puzzle experiments and target midtraining-token fractions in real-world experiments.Example counts and token fractions can differ because example lengths vary; mean lengths are used to convert target fractions to example counts.Source: §A.1 · p. 23
  • During RL, unified training adds m ∈ {4, 8, 16, 32, 64} examples per update from the fixed corpus, while masking prompt tokens from the midtraining loss.Unlike during pretraining, only response and end-of-sequence tokens contribute to this RL-phase loss.Source: §A.1 · p. 23
  • Configured midtraining-token fractions use fixed pretraining-time length estimates, so the actual batch token fraction can shift as RL response lengths change.Each update retains 512 RL responses; rounding is applied to cumulative example counts.Source: §A.1 · p. 24
  • Schedules can allocate the same average configured proportion differently over training, and realized token proportions may still differ across runs.The described schedules include constant, full annealing, growth, and half annealing; response lengths can affect realized proportions.Source: §A.1 · p. 24
  • Unified loss weighting makes the midtraining contribution depend on its token count relative to RL response tokens, unlike a separate-loss objective with fixed coefficient α.The unified objective also has an overall scale of TR/(TR + TD).Source: §A.1 · p. 24

A.2 Compute Calculations in the Controlled Experiments

The controlled experiments distinguish data exposure from compute and estimate the latter using training and generation costs. Their reasoning puzzles use controlled reference compositions and verifiable solutions to study learning across training phases.

A.2 Compute Calculations in the Controlled Experiments

  • Matching RL update counts does not match total compute because unified training also processes midtraining data during pretraining and RL.Staged training instead processes 1, 2, or 4 midtraining-corpus passes before RL and uses no midtraining data during RL.Source: §A.2 · p. 25
  • Compute estimates count training and RL generation costs, but omit attention-length corrections, padding overhead, recomputation, evaluation, unused generations, and work discarded after restarts.The RL token coefficient accounts for generation followed by policy training; prompt tokens count even when masked from loss.Source: §A.2 · p. 25

B Reasoning Puzzles: Construction and Examples

  • The puzzle construction controls reasoning patterns and the number of people, using a symbolic oracle to create worked solutions before introducing additional compositions during RL.The puzzles are designed to study data choices throughout the training pipeline.Source: B Reasoning Puzzles: Construction and… · p. 27

B.1 Puzzle Construction

  • Each puzzle asks for a complete assignment of people to distinct occupied floors, with clues that can refer to people directly or through friendship and floor rank.A hidden assignment generates a clue chain; a symbolic solver verifies uniqueness and that constraint propagation can recover it.Source: §B.1 · p. 27

B.2 Reasoning Patterns and Midtraining Data

The reasoning-puzzle setup separates demonstrated reference operations from compositions introduced during RL. It uses oracle traces, examples, and distinct task settings to examine learning and generalization.

B.2 Reasoning Patterns and Midtraining Data

  • Midtraining teaches selected reference operations but omits their compositions; RL introduces those compositions, and held-out questions test the corresponding distributions.The omitted compositions appear in RL training rather than being reserved exclusively for evaluation.Source: §B.2 · p. 27, Table 4 · p. 28
  • Oracle traces use constraint propagation to narrow feasible floors, resolve indirect references, and produce a complete assignment without branching over alternatives.Up to four traces per question are generated by varying constraint order and verbalization, with duplicate traces removed.Source: §B.2 · p. 28

B.3 Examples

  • The examples show oracle training targets and evaluation prompts; gold assignments in evaluation examples are shown separately and are not model inputs.Figure 10 illustrates two omitted compositions using the known answer to explain the task, not as input to the model.Source: §B.3 · pp. 29–30, Figure 10 · p. 30

C Pretraining, Midtraining, and RL: Setup and Additional Results

  • The fixed midtraining corpus has 200,000 questions and, after duplicate-trace removal, 796,175 prompt–response examples totaling 509,942,713 tokens (about 510M); staged and unified training reuse this same corpus.Source: §C.1 · p. 32

C.1 Tasks and Midtraining Data

  • Base tests omitted compositions at the midtraining puzzle size, while Extended adds larger 10–11-person puzzles and evaluates generalization on 12-person puzzles unseen in training.The Extended test set with 12 people requires no additional training.Source: §C.1 · p. 32, Table 5 · p. 32

C.2 Model and Pretraining

  • Unified pretraining retains the web-text batch and adds 8, 16, 32, 64, or 128 packed midtraining sequences, increasing batch size and compute while keeping updates and web exposure fixed.The added sequences contribute approximately 133M to 2.14B midtraining tokens; all tokens, including prompts, receive language-model loss.Source: §C.2 · p. 33

C.3 Staged Initialization and RL Training

The appendix compares staged and unified initialization and RL configurations, then evaluates their reasoning performance, retention protocol, and compute tradeoffs. Unified training leads the best-observed accuracy–coverage frontier across most compute budgets, while staged training contributes at mid-level budgets for in-distribution evaluation.

C.3 Staged Initialization and RL Training

  • Staged runs use a separate midtraining phase before RL, while unified runs mix midtraining examples into RL batches from mixed-pretraining initializations.Both approaches use 512 RL responses per update; unified runs add 4–64 midtraining examples, with shared token-mean normalization.Source: §C.3 · p. 34

C.4 Evaluation Protocol

  • Accuracy is sampled pass@1 and coverage is pass@16, estimated from 64 responses per question; scores across two RL seeds are averaged without pooling generations.The held-out test sets each contain 1,024 questions.Source: §C.4 · p. 36

C.5 Reasoning Performance before RL

  • Before RL, larger unified midtraining mixtures improve reasoning, but staged initialization has higher Base pass@1 and the largest unified mixture reaches similar pass@16.On Extended and Extended generalization, even the largest mixture reaches only 2.34% and 0.75% pass@16, respectively; applying demonstrated patterns to new compositions remains difficult.Source: §C.5 · p. 37

C.6 Staged Performance after RL

  • The staged reference is selected by final RL pass@1, and the same two-epoch configuration at 2 × 10^-4 is selected for Base and Extended.More corpus exposure does not necessarily help: the best four-epoch configuration finishes below the best two-epoch configuration in both settings.Source: §C.6 · p. 39, Table 9 · p. 39

C.7 Additional Compute Comparisons

  • Unified training achieves higher best-observed accuracy and coverage at smaller and larger compute budgets across all three evaluations, while staged training contributes at mid-level in-distribution budgets.The comparison selects the best observed checkpoint within each budget, including checkpoints before a run ends.Source: Figure 12 · p. 40, §C.7 · p. 40

C.8 Retaining Pretraining Data during RL

The appendix tests whether pretraining text during RL preserves broader capabilities, compares loss aggregation choices, and reports retention before RL. Results vary with configuration: retained pretraining data usually improves retention at a reasoning cost, while no aggregation rule consistently wins matched comparisons.

C.8 Retaining Pretraining Data during RL

  • Adding pretraining text during RL improves held-out pretraining loss in all eight configurations and the general-benchmark average in seven, usually at the expense of reasoning accuracy and coverage.The benchmark gains are all smaller than 0.7 percentage points.Source: §C.8 · p. 41, Figure 13 · p. 41, Table 10 · p. 41

C.9 Combining Midtraining and RL Losses

  • Unified training has the highest observed final accuracy and coverage in Base and Extended generalization, but does not consistently beat separately normalized losses in matched comparisons.Relative performance depends on task, initialization, and the amount of midtraining data included during RL.Source: §C.9 · p. 42, Figure 14 · p. 43, Figure 15 · p. 44

C.10 Retention before RL

  • Before RL, unified mixtures keep mean pretraining loss near the reference and mean benchmark scores within 0.52 percentage points, whereas every tested staged variant worsens both aggregate measures.Staged midtraining improves TruthfulQA MC2 but lowers MMLU, ARC-Easy, HellaSwag, and WinoGrande performance.Source: §C.10 · p. 48, Table 11 · p. 48

D Details of Training with Real-world Midtraining Data

D.1 Experimental Setup

The real-world experiments vary the proportion, schedule, and composition of midtraining data mixed with mathematical RL, and compare RL objectives and training implementations. The effects differ across accuracy, coverage, retention, and task, so no single mixing choice dominates every metric.

D.1 Experimental Setup

  • The real-world setup uses a 240M-token additional midtraining corpus and Math12K prompts for mathematical RL; configured mixing percentages target a fraction of combined midtraining and RL-response tokens.The realized token fraction can vary as RL response lengths change.Source: §D.1 · p. 49

D.2 Effect of the Midtraining Data Proportion

  • At 5% mixing, Coding pass@32 is 45.17%, compared with 44.62% for the RL-only baseline; across proportions, coding accuracy remains below that baseline while coding coverage exceeds it at every nonzero proportion.Mathematical coverage is nonmonotonic, and increasing the midtraining proportion consistently lowers held-out corpus loss.Source: §D.2 · p. 53, Figure 22 · p. 54

D.3 Data Mixing Schedules

  • Ending midtraining exposure halfway through training lets Coding accuracy improve further and preserves SciQ gains, while held-out corpus loss increases; Math and MMLU results are similar across the four schedules.Math pass@1 spans 24.01–24.80%, and Math pass@32 spans 68.32–68.94%.Source: Figure 23 · p. 55, §D.3 · p. 55, Table 17 · p. 55

D.4 Effect of Code Data

  • Removing designated code sources raises Coding accuracy from 16.51% to 18.65% under constant mixing, but has little effect under full annealing; verified replacements do not improve accuracy.Replacing the sources improves Coding coverage under full annealing, so corpus effects depend on schedule and metric.Source: Figure 25 · p. 57, §D.4 · p. 57, Table 18 · p. 57

D.5 MaxRL versus REINFORCE

  • Across six mixing settings, MaxRL consistently achieves higher mathematical coverage, while REINFORCE yields lower held-out loss and higher SciQ and ARC-Challenge accuracy.Effects on mathematical accuracy and coding depend on mixing proportion and schedule.Source: §D.5 · p. 58, Figure 27 · p. 59

D.6 Unified Training with Zero Mixing versus Staged RL

  • At constant 0% mixing, unified training runs at 3.7 seconds per update versus 28.8 seconds for staged RL, while mathematical accuracy is similar and held-out loss increases more under unified training.The throughput comparison reflects the implemented systems and excludes launch and restart downtime; it does not isolate staleness or baseline subtraction effects.Source: §D.6 · p. 60

D.7 Coding RL after Mathematical Training

The section tests whether mathematical-training history affects subsequent coding RL and retention when models use the same coding procedure. Unified initializations retain mathematical accuracy and broader capabilities better than the staged model on several measures, though coverage and held-out loss do not improve uniformly.

Initialization

  • Models start coding RL from mathematical-training checkpoints and use the same coding configuration without midtraining data, isolating the effect of prior training history.The unified models use constant midtraining proportions of 0.5%, 1%, 2%, and 5%.Source: §D.7 · p. 61

Data and training

  • Coding RL uses 20,000 Python problems from KodCode-V1, with binary reward for passing all supplied unit tests.The problems comprise 10,000 labeled medium-difficulty and 10,000 labeled hard-difficulty examples; all models use shared RL settings.Source: §D.7 · p. 61, Table 21 · p. 61

Evaluation

  • Coding pass@1 and pass@16 are evaluated during training, with final coding scores reported after 4,000 updates.The evaluations use 16 responses for each of 164 HumanEval+ and 378 MBPP+ problems; mathematics, general knowledge, and held-out corpus loss are also assessed for retention.Source: §D.7 · p. 61, Table 22 · p. 61

Performance beyond coding

  • All unified models retain higher MMLU and SciQ accuracy than staged training after 4000 coding updates, while mathematical accuracy stays near its starting level.Coverage can increase or decrease depending on initialization; held-out corpus loss rises for every model, but rises less for unified models, whose final loss is lower.Source: §D.7 · pp. 61–62, Figure 28 · p. 62

D.8 WebInstruct RL after Mathematical Training

The section tests whether prior mathematical training with unified mixtures changes later WebInstruct RL. Both unified initializations finish with higher accuracy than staged training on validation and test data, while gains in pass@64 and coverage depend on answer type and mixture.

Initialization and training

  • The continuation compares staged training with unified models using 2% and 5% midtraining proportions, all trained with the same on-policy MaxRL procedure and no midtraining data.Each model starts from mathematical training, and WebInstruct training is a separate continuation.Source: §D.8 · p. 63

Data and reward

  • Training samples 20,000 nonmathematical WebInstruct questions and uses a binary reward requiring verifier acceptance and answer-format checks.Numerical answers must also satisfy a 1% relative-error check; the evaluation reserves disjoint validation and test sets.Source: §D.8 · p. 63

Evaluation

  • Validation accuracy is monitored every 100 updates, while test pass@1 and pass@64 are measured at updates 0, 1,000, 2,000, and 3,000.Test results are reported overall and separately for numerical and free-form answers; three prompts exceeding the response budget are excluded throughout.Source: §D.8 · p. 63

Results

  • Both unified models start below staged training in validation and test accuracy but finish above it after 3,000 updates.The 2% model also has higher test pass@64; the 5% model remains slightly below staged training on that metric.Source: §D.8 · p. 63, Figure 29 · p. 63
  • Both unified models achieve higher final free-form accuracy and coverage than staged training, while only the 2% model achieves higher numerical-answer coverage.Overall, the answer-type breakdown shows that benefits vary across answer types and metrics.Source: §D.8 · p. 63

E Comparison with KL Regularization

The section compares unified training with mathematical RL regularized toward a frozen base model, assessing both task performance and retention. Unified mixtures improve midtraining loss and retain stronger MMLU performance than the tested KL settings, but the approaches trade off on some task metrics.

Experimental setup

  • All comparisons start from the same base model and use Math12K prompts; unified runs mix midtraining data, whereas KL runs use no midtraining data during RL.The KL comparison varies β across 0, 0.001, 0.01, and 0.1; β = 0 is the unregularized control.Source: E Comparison with KL Regularization · p. 64

KL penalty

  • The KL penalty compares sampled-response token probabilities with a frozen copy of the base model and adds a β-scaled token penalty to the weighted loss.The clipped reference-minus-current log-probability difference is transformed as e^u − u − 1, and the penalty applies to both correct and incorrect responses.Source: E Comparison with KL Regularization · p. 64

Results

  • At final checkpoints, all five unified mixtures achieve higher MMLU accuracy than every tested KL-regularized run, and mixtures of 0.5% or more also achieve higher SciQ accuracy than those runs.The 0.2% mixture is slightly below the weakest KL penalty on SciQ; unified training lowers validation loss below the base value without the strongest penalty’s large reduction in mathematical performance.Source: E Comparison with KL Regularization · p. 65, Figure 31 · p. 65
  • Increasing a nonzero KL coefficient reduces deterioration in token-averaged midtraining validation loss, but the strongest penalty, β = 0.1, substantially reduces mathematical accuracy and coverage.Validation curves use a fixed set of 512 held-out windows; benchmark loss instead equally weights the 24 corpus components, so the two loss measures are distinct.Source: E Comparison with KL Regularization · pp. 64–65, Figure 30 · p. 64
  • None of the tested KL penalties matches unified training’s combination of mathematical performance, lower component-averaged corpus loss, and higher SciQ and MMLU accuracy.Weaker KL penalties yield higher coding accuracy, the strongest KL penalty yields the highest coding coverage, and ARC-Challenge shows no consistent advantage for either approach.Source: E Comparison with KL Regularization · p. 65

F Policy Gradient Objectives and Continuous Rewards

The section expresses policy-gradient updates as likelihood fitting to targets derived from a fixed generation policy. REINFORCE uses a reward-weighted completion target and prompt weighting, while TailRL represents continuous rewards through targets conditioned on reward thresholds.

F Policy Gradient Objectives and Continuous Rewards

  • For fixed generation policies, the framework constructs target distributions and prompt weights so likelihood updates recover REINFORCE and TailRL policy gradients when the generation policy is the current model.For a different generation policy, the update fits the target induced by that policy and need not match the current policy’s gradient.Source: F Policy Gradient Objectives and… · p. 66, §F.1 · p. 66, §F.2 · p. 67

F.1 REINFORCE

  • REINFORCE fits a reward-weighted target distribution and weights each prompt by its expected reward; for binary rewards, the target is the generation policy conditioned on correctness.The prompt weight cancels the target distribution’s normalization, and the expected reward can be absorbed into an effective prompt distribution.Source: §F.1 · p. 66
  • With on-policy sampling, fixed prompt-dependent baselines leave the expected gradient unchanged, and fixed positive prompt-dependent reward scales change relative prompt weights.A reward with a finite lower bound can be shifted to be nonnegative; the baseline invariance need not hold for uncorrected updates using an earlier policy.Source: §F.1 · p. 66
  • The likelihood identities concern expected on-policy gradients with fixed prompt-dependent baselines and scales; PPO and GRPO choices such as empirical advantages, clipping, and KL penalties need not preserve the fixed-target formulation.These identities therefore do not establish the same likelihood representation for those modified update procedures.Source: §F.1 · p. 67

F.2 Continuous Rewards with TailRL

  • TailRL represents continuous rewards by averaging success log-probabilities over thresholds and fitting a mixture of threshold-conditioned targets with uniform threshold weighting.When the generation policy is current, the likelihood update recovers the TailRL gradient; binary rewards reduce to the MaxRL objective and target.Source: §F.2 · p. 67
  • The TailRL derivation establishes population gradients, while continuous-reward finite-sample estimation must also account for different reward thresholds.The cited work derives the corresponding estimators and finite-order TailRL objectives.Source: §F.2 · p. 67

Related papers

Loading 2610.40564v1…