PaperlayerView exact revision on arXiv

Source-linked AI summary

Training-Free On-Policy Distillation

Bangji Yang, Jiajun Fan, Zhaoyu Zhu, Hongbo Ma, Ge Liu

arXiv:2610.22236v1cs.LGcs.AIcs.CL

TL;DR

Domain-specific rewards produce specialists that excel in their own domains but lag in others, making it difficult to combine their capabilities in one model. For specialists sharing a base model, the paper replaces training-based on-policy distillation with exact layer-wise fitting from rollout statistics; on Olmo 3 7B, the fused model keeps most of each specialist’s gains across seven benchmarks and is more accurate on average than either specialist alone.

  • Problem

    Domain-specific rewards produce specialists that excel in their own domains but lag in others, so a user seeking multiple domains must build a model from the specialists.

    Source: §1 · p. 1

  • Method

    For specialists sharing a base model, the method solves each layer’s least-squares distillation objective using second moments gathered as teachers read their own rollouts, without backpropagation.

    Source: §1 · p. 2, §3.1 · p. 4

  • Results

    On Olmo 3 7B, the fused model keeps most of each specialist’s gain on seven benchmarks and is, on average, more accurate than either specialist alone.

    Source: §1 · pp. 2–3

  • Takeaways & Limitations

    For shared-base specialists, adding a specialist costs its rollouts and one more term in each layer’s sums; nothing is retrained.

    Source: §4.6 · p. 10

    The layer statistics must come from teachers processing their own generated answers; reference solutions cannot provide that distribution.

    Source: §3.2 · p. 5

Abstract

from arXiv · show

On-policy distillation consolidates specialist models into one student at the price of rollouts, teacher passes and gradient steps for every teacher. We show that the consolidation does not need training: with forward passes only, the capabilities of separately tuned specialists can be aggregated into one model that is, on average, more accurate than either specialist alone. When specialists share a base model, distillation has an exact layer-wise solution: each linear layer of the fused model must reproduce every specialist's output on the inputs that specialist receives while reading its own answers, a least-squares problem whose solution needs only the second moments of those inputs. One forward pass over a few hundred rollouts provides them. We call this training-free on-policy distillation (\method). What the rollouts contain matters more than the solver: answers truncated in repetition loops are few among a mathematics specialist's rollouts but carry a large share of their tokens, and removing them lifts the fused model on AIME25 from task arithmetic's level to the specialist's at matched precision. On Olmo~3 7B, a few GPU-hours of forward computation per specialist yield a generalist that keeps most of each specialist's gain on each of seven mathematics and code benchmarks.

1 Introduction

The paper proposes training-free on-policy distillation to combine specialists that share a base model, using their own rollouts to compute exact layer-wise least-squares solutions. On Olmo 3 7B, the fused model retains most specialist gains across seven benchmarks and has higher mean accuracy than either specialist, while rollout filtering matters on hard problems.

1 Introduction

  • The paper frames specialist consolidation as a gap between training or ignoring rollouts: on-policy distillation uses rollouts but requires teacher passes and gradient steps, while weight arithmetic ignores inputs.RegMean uses input statistics from task training corpora rather than a reasoning model’s generated text.Source: §1 · p. 1
  • TF-OPD uses specialist-generated rollouts to gather layer-input second moments, then solves a least-squares problem for each fused layer without backpropagation.Each specialist reads its kept answers once; a linear solve produces each layer.Source: §1 · pp. 2–3, Figure 2 · p. 3
  • On Olmo 3 7B, the fused model retains most of each specialist’s gains across five mathematics and two code benchmarks and has higher mean accuracy than either specialist.The method uses forward computation only, with a few GPU-hours per specialist.Source: §1 · pp. 2–3, Figure 1 · p. 2
  • Answers truncated in repetition loops are few among mathematics rollouts but account for a large share of tokens, and removing them lets the fused model match its mathematics specialist on AIME25 in the same format.The rollout filter separates a model that falls short of the specialist from one that matches it.Source: §1 · p. 3
  • A close token-level fit does not ensure success on hard problems; the paper reports that student rollouts, domain reweighting, and mined hard problems do not improve on selection, while numerical precision does.The paper attributes the outcome on hard problems to which inputs the statistics weigh.Source: §1 · pp. 2–3

2 Related Work

The related work places TF-OPD between distillation and model merging: it retains teacher matching and generated-text statistics while replacing optimization with layer-wise closed-form solutions. It also distinguishes its rollout-based statistics from prior methods that use weights alone or training-input data.

Knowledge distillation

  • On-policy distillation supervises model-generated text and has been applied to reasoning and knowledge fusion, but the cited approaches use optimization; TF-OPD solves the matching objective layer by layer instead.Prior rollout-selection approaches select hard problems, weight by student pass rate, or balance student-generated states across teachers.Source: §2 · p. 3

Model merging

  • Weight averaging, task arithmetic, TIES-Merging, DARE, and spectral methods merge models using weights, while Fisher merging and RegMean bring in data through parameter importance or input second moments.TF-OPD uses RegMean’s closed form but gathers statistics on specialists’ own generations rather than supervised-task training inputs.Source: §2 · pp. 3–4

Other training and evaluation settings

  • Other cited work studies multi-domain fine-tuning, distributional Pareto-optimal reinforcement learning, interactive guidance, imperfect demonstrations, supervision choices, and distinct evaluation or application settings.The passage groups these lines as other training and evaluation settings.Source: §2 · p. 4

3 Method

The method solves a layerwise distillation objective using input statistics from specialists’ own rollouts, without back-propagation. Its result depends on which rollouts are retained and on storing the fused model at sufficient precision.

3.1 Distillation of one layer is least squares

  • The layerwise objective matches each teacher’s outputs on its own inputs, and its solution depends on weighted input second moments.Shrinking off-diagonal moments improves conditioning and is equivalent to a per-coordinate penalty pulling the fused weight toward each teacher.Source: §3.1 · p. 4
  • When teacher input moments match, the solution reduces to weight averaging; with orthogonal input ranges, each teacher update is reproduced on its own range.The passage identifies interference as arising only where input distributions overlap.Source: §3.1 · pp. 4–5

3.2 On-policy statistics

  • On-policy statistics come from layer inputs as each teacher processes its own answers, gathered in one teacher-forced forward pass over selected rollouts.Gram matrices average over token positions so longer-answer domains do not dominate by length.Source: §3.2 · p. 5

3.3 Solving

  • The method solves the layer objective once per member layer in double precision, while combining embeddings, normalization parameters, and biases by task arithmetic.The solve also reports, for each layer and teacher, the share of the update left unreproduced.Source: §3.3 · p. 5

3.4 Which rollouts to use

  • Repetition-loop answers are few but contribute many tokens, concentrating Gram-matrix statistics on degenerate inputs; selection excludes answers that do not finish or are flagged as loops.When automatically gradable, mathematics answers are additionally filtered for correctness and code answers are run against tests.Source: §3.4 · p. 5
  • The fused model has the highest mean accuracy across seven benchmarks, ahead of both specialists on average.Source: Table 2 · p. 6

3.5 Precision

  • The fused model is stored in float32 because rounding its update to bfloat16 measurably distorts it, and is evaluated in float16 for single-GPU memory fit.The paper states that format changes shift measured accuracy, so scores are reported with their evaluation format.Source: §3.5 · p. 6

4 Experiments

Experiments on Olmo 3 7B find that the fused model is more accurate on average than either specialist while staying close to each specialist on its domain. The findings also show that rollout selection matters on hard mathematics problems, while the approach uses forward passes and a layer-wise solve rather than training.

4.1 Setup

  • The setup uses shared-base Olmo 3 7B math and code specialists, selected training rollouts, and seven mathematics and code benchmarks.Each specialist answers 256 training problems; evaluation uses separate problems, with four sampled answers per problem on AIME24, AIME25, and HMMT25 and one on the other benchmarks.Source: §4.1 · p. 6

4.2 Main results

  • The fused model has the highest mean accuracy of the four models, stays close to each domain specialist, and keeps most specialist gains across benchmarks.It equals the mathematics specialist on AIME24, is slightly ahead on OlympiadBench, and is nearly level with the domain specialist on MATH-500 and both code benchmarks.Source: §4.2 · pp. 6–7

4.3 Case study

  • On HMMT25, TF-OPD solves a geometry problem where all four mathematics-specialist samples fail; on a code task, the fused model implements the described rule while the coding specialist fails in every seed.The geometry solution uses the fact that a square’s diagonals share a midpoint.Source: Figure 3 · p. 7, §4.3 · pp. 7–8

4.4 Qualitative analysis

  • Looping answers are few by count but contribute several times more tokens, so their inclusion weights the statistics toward loop behavior and hurts hard mathematics, while code benchmarks remain matched.The all-answer model can remain close to the specialist in next-token predictions on correct solutions even while matching task arithmetic on AIME25.Source: §4.4 · pp. 8–9

4.5 Ablation study

  • Removing looping answers brings the fused model to the mathematics specialist’s AIME25 level; stricter correctness filtering adds no further gain, and code-test verification changes code accuracy by no measurable amount.Task arithmetic and the all-answer fused model fall short of the mathematics specialist in the matched evaluation format.Source: §4.5 · p. 9
  • Three-to-one mathematics weighting recovers most of the loss at a coding cost, mined hard problems give no clear gain, and easier benchmarks do not distinguish merge rules.The share of the coding specialist’s update left unreproduced nearly doubles under reweighting; task arithmetic, its low-rank variant, and the closed form are indistinguishable on MATH-500 and both code benchmarks.Source: §4.5 · p. 9

4.6 Discussion

  • For specialists sharing a base, the layer-wise closed form yields a model close to each teacher without retraining; training remains useful when teachers differ in architecture or size from the student.Adding a specialist requires its rollouts and another term in each layer’s sums.Source: §4.6 · p. 10

5 Conclusion

TF-OPD aggregates separately tuned specialists without training when they share a base model, using each teacher’s input second moments to solve layer-wise least squares. On hard mathematics problems, removing looping answers lets the fused model match its mathematics specialist in the same format.

Conclusion

  • With a shared base model, TF-OPD solves each linear layer’s least-squares distillation objective exactly from second moments of the teachers’ inputs.The teachers obtain those moments by reading a few hundred of their own answers; this requires forward passes, not training.Source: §5 · p. 10
  • When separately tuned specialists share a base model, their capabilities can be aggregated without training-based on-policy distillation, yielding a model more accurate on average than any source.For a linear layer, the least-squares solution is exact given teachers’ input second moments, which are obtained by having teachers read a few hundred of their own answers at the cost of forward passes.Source: §5 · p. 10
  • Removing looping answers, which are few in number but carry a large share of tokens, lets the fused model match its mathematics specialist on hard problems in the same format.Source: §5 · p. 10

A Derivation and properties

The derivation expresses layer-wise distillation as a quadratic least-squares objective and gives a closed-form update. Its properties describe how input statistics govern averaging, teacher recovery, and interference.

Closed form

  • The layer-wise objective is convex and quadratic in the fused weights, and its solution is the same whether solving for the weight or its update.The closed form uses the teachers’ input second moments; shrinkage modifies those statistics before the solve.Source: A Derivation and properties · p. 15

Shrinkage

  • Shrinkage adds a penalty weighted by each input coordinate’s second moment, and makes the summed regularized statistics invertible when all coordinates are active and α < 1.Source: A Derivation and properties · p. 15

Identical statistics

  • When teachers have identical input statistics, TF-OPD reduces to averaging their updates rather than summing them as task arithmetic does.This is the equal-weight case when teachers cannot be distinguished by their inputs.Source: A Derivation and properties · p. 15

Dominant teacher

Orthogonal inputs

  • When teachers’ input-statistic ranges are mutually orthogonal, the solution reproduces each teacher’s outputs exactly on its own inputs, with zero objective.Interference is confined to input directions the teachers share.Source: A Derivation and properties · p. 15

Comparison with task arithmetic

  • Task arithmetic applies every teacher’s update to every input, whereas TF-OPD minimizes the resulting input-weighted layer error.The diagnostic ε_d is zero when a teacher is reproduced on its inputs and one when the fused layer does no better than the base layer.Source: A Derivation and properties · pp. 15–16

B Implementation and evaluation protocol

The implementation gathers statistics from selected specialist rollouts, solves each layer’s linear system, and combines other tensors by task arithmetic. The evaluation uses held-out mathematics and code benchmarks, with no gradients or optimizer state.

Algorithm

  • The procedure generates specialist answers, keeps those that are correct, finished, and not loops, accumulates layer statistics, and solves the systems.Embeddings and other non-linear-map tensors are combined by task arithmetic.Source: B Implementation and evaluation protocol · p. 16

Models

Rollouts

  • Each specialist generates one answer for 256 training problems, with separate mathematics and code datasets and generation limits matched to evaluation.No evaluation problem is used for rollout collection.Source: B Implementation and evaluation protocol · p. 16

Selection

  • Selection keeps 175 correct, finished, loop-free mathematics answers; for code, it keeps 92 answers that pass every test among 241 problems read by standard input.All 19 truncated mathematics answers are repetition loops, including two graded correct before looping; no code answer is truncated or a loop.Source: B Implementation and evaluation protocol · p. 16

Statistics

  • Teacher-forced forward hooks accumulate normalized X⊤X statistics over prompt and answer tokens, grouping linear layers that receive identical tensors.In Olmo 3, each block has four such groups, and final hidden states provide output-head statistics.Source: B Implementation and evaluation protocol · pp. 16–17

Solve

  • The solve uses α = 0.9 and unit teacher weights, solves each system in 64-bit floats, and stores the result in 32-bit floats.Embeddings, normalization gains, and biases use task arithmetic; the fused-update cosine after bfloat16 rounding is about 0.96.Source: B Implementation and evaluation protocol · p. 17

Evaluation

  • Mathematics benchmarks sample at temperature 0.7 and code benchmarks at 1.0; released checkpoints use bfloat16 while fused models use float16.The evaluation includes AIME24, AIME25, HMMT25, MATH-500, OlympiadBench, and two code benchmarks.Source: B Implementation and evaluation protocol · p. 17

Compute

C Ablations and diagnostics

The all-answers fused model loses to the mathematics specialist on AIME25, while outperforming it on MATH-500 and matching the coding specialist on code benchmarks. Diagnostics show that rollout selection matters on hard mathematics problems, while other tested variants provide no clear improvement.

Weight-only rules

  • The weight-only comparison includes task arithmetic and a low-rank variant that sums low-rank update components while averaging the rest.The low-rank portion is selected using a singular-value threshold of 1.05 times the Marchenko–Pastur edge.Source: C Ablations and diagnostics · p. 17

Amount of rollouts

  • The rollout comparison uses 128 problems per domain with a 16k-token limit versus 256 problems at evaluation limits; the usual closed-form model uses the larger set.The latter statistics include 2.6 million mathematics answer tokens before selection and 1.2 million afterward.Source: C Ablations and diagnostics · p. 17

Student rollouts

  • The tested student-rollout variants use student activations or teacher-forced statistics on student answers, but the reported agreement is the same for the student-rollout and all-answers models.Source: C Ablations and diagnostics · pp. 18–19

Reweighting and mined problems

  • Weighting mathematics three to one recovers most of the loss at the coding specialist’s expense, while mined hard-problem statistics give no clear gain.The mined-problem setup uses problems selected by a 1.5B reasoning model and leaves code statistics unchanged.Source: C Ablations and diagnostics · p. 18

Results

  • On AIME25, the all-answers model loses to the mathematics specialist; on MATH-500 it leads by a margin whose interval excludes zero, and on code benchmarks it is level with the coding specialist.With thirty AIME25 problems, the intervals are wide; the final model uses selected mathematics statistics and verified code statistics, not the all-answers statistics.Source: C Ablations and diagnostics · p. 18, Table 6 · p. 18, Table 7 · p. 18

Selected statistics

  • The final model’s selected mathematics statistics and verified code statistics differ from the all-answers model’s statistics.The paired AIME25 comparison reports a loss for the all-answers model to the mathematics specialist, with thirty problems and wide intervals.Source: C Ablations and diagnostics · pp. 17–18

Teacher-forced agreement

  • Teacher-forced agreement with the mathematics specialist does not degrade along its correct solutions, and the student-rollout model has the same agreement as the all-answers model.The comparison covers 51 AIME25 and 64 MATH-500 solutions.Source: C Ablations and diagnostics · p. 19, Table 8 · p. 19

Unreproduced share

  • Keeping only correct, loop-free mathematics answers shifts the unreproduced update share slightly toward mathematics; weighting mathematics three to one halves its share but nearly doubles code’s.Source: C Ablations and diagnostics · p. 19, Table 9 · p. 19

D Case studies

The case studies show the fused model solving examples that the domain specialist misses in every sample, across mathematics and code. Its mathematics solution uses a square’s shared diagonal midpoint, while the code examples highlight correct handling of task rules and edge cases.

Case selection

  • Case studies are problems where the fused model succeeds while its domain specialist fails in every sample; specialists were sampled four times per model on competition benchmarks, versus one fused-model sample against four coding-specialist seeds.Examples include HMMT25 problems 21 and 26 and AIME25 problem 25, plus seven listed code problems.Source: D Case studies · p. 19

HMMT25, problem 21

  • The fused model gives HMMT25 problem 21’s answer, 63, in three of four samples, while the mathematics specialist is wrong in all four.The fused model equates the diagonal midpoints to obtain H − 26 = 37.Source: D Case studies · p. 19, Figure 6 · p. 20

HumanEval/126

  • On HumanEval/126, the fused model passes base and additional tests, while all four coding-specialist seeds mishandle repeated numbers and fail the base tests.The task requires checking ascending order and returning false if a number occurs more than twice.Source: D Case studies · p. 20

Mbpp/92

  • On Mbpp/92, the fused model passes all tests, while the coding specialist fails in all four seeds, including three that miss additional edge cases.The task asks whether different digits alternate; the fused model checks the first pair and then compares each digit with the one two places earlier.Source: D Case studies · p. 20

Related papers

Loading 2610.22236v1…