PaperlayerView exact revision on arXiv

Source-linked AI summary

Trust-Region Behavior Blending for On-Policy Distillation

Daniil Plyusov, Alexey Gorbatovski, Alexey Malakhov, Nikita Balagansky, Boris Shaposhnikov, Daria Korotyshova, Daniil Gavrilov

arXiv:2605.31159v1cs.LGcs.AI

TL;DR

Early on-policy distillation can expose teacher supervision to weak student-generated prefixes. TRB instead uses a teacher-guided behavior policy within a student-centered KL trust region during warmup, and achieves the strongest average across two math-reasoning settings.

  • Problem

    Early student rollouts can generate low-quality prefixes, limiting whether student-visited trajectories carry usable teacher signal for on-policy distillation.

  • Method

    TRB optimizes a teacher-guided behavior policy within a student-centered KL trust region during warmup while keeping the per-prefix reverse-KL objective fixed and annealing blending away.

  • Results

    Across two math-reasoning distillation settings, TRB achieves the strongest average among the compared methods.

  • Takeaways & Limitations

    TRB’s results support using teacher-guided off-policy behavior as a warmup rather than maintaining it throughout training.

  • Takeaways & Limitations

    The study covers only two Qwen3-Base math-reasoning settings, and TRB adds temporary warmup cost from online teacher decoding and student–teacher co-residency.

Abstract

from arXiv · show

On-policy distillation (OPD) trains a student on prefixes sampled from its own policy while matching a stronger teacher. This addresses the prefix mismatch of offline distillation, but early student rollouts can still be poor, placing teacher supervision on weak or low-quality prefixes. We propose Trust-Region behavior Blending (TRB), a warmup method that replaces the early rollout policy with the closest-to-teacher behavior policy inside a student-centered KL trust region, while keeping the per-prefix reverse-KL OPD loss unchanged. The KL budget is annealed to zero, so training returns to pure student rollouts after warmup. Across two math-reasoning distillation settings, TRB attains the strongest average among the compared methods.

1 Introduction

On-policy distillation avoids fixed-prefix mismatch by rolling out the student, but early weak prefixes can make teacher supervision brittle. Trust-Region behavior Blending addresses this early regime by guiding rollout behavior within a student-centered KL trust region, then annealing the intervention away; across two math-reasoning settings, it attains the strongest average.

  • Motivation: On-policy distillation rolls out the current student, avoiding the prefix distribution mismatch caused by fixed teacher-forced or teacher-generated prefixes.Knowledge distillation matches teacher predictions, while OPD applies teacher supervision to student-generated prefixes.
  • Motivation: Early OPD is brittle because weak students can generate low-quality prefixes, and its effectiveness depends on whether visited trajectories carry usable teacher signal.Pure student rollouts preserve the target distribution, whereas stronger teacher intervention improves local prefix quality by moving collection off-policy.
  • Method: TRB selects the feasible behavior policy closest to the teacher subject to D_KL(µ ∥π_S) ≤ε, keeping teacher-guided behavior close to the student.At each prefix, the student policy defines the KL trust region and TRB chooses the closest feasible teacher-oriented behavior policy.
  • Method: TRB controls behavior during early rollout collection without changing the per-prefix distillation objective, then anneals the intervention away after warmup.The method is used only in the early regime, allowing training to return to pure student rollouts.
  • Evaluation: Across two math-reasoning distillation settings, TRB attains the strongest average against vanilla OPD and alternative teacher-guidance and warmup methods.Comparisons include target-side reformulation, direct token replacement, persistent blending, and simpler warmup heuristics.

2 Background

On-policy distillation samples prefixes from the current student policy and uses a reverse-KL objective. TRB leaves this per-prefix loss unchanged while modifying only the prefix-generating behavior policy within a student-centered trust region.

  • 2 Background: OPD samples prefixes from the current student policy rather than a fixed offline dataset.The prefix distribution is induced by student rollouts.
  • 2 Background: TRB keeps the per-prefix reverse-KL loss fixed and changes only the behavior policy used to generate prefixes.The behavior policy is denoted µ.
  • 2 Background: The constrained objective measures closeness to the teacher with DKL(µ∥πT) and defines the student-centered trust region with DKL(µ∥πS).The trust-region formulation uses µ as the rollout behavior policy.
  • 2 Background: The reverse-KL term is estimated on a truncated student top-k support, with this approximation fixed across rollout-side variants.This implementation follows top-k OPD.

3 Related Work

Related work moves from fixed-distribution distillation toward student-generated trajectories, while addressing instability in teacher supervision and gaps between student and teacher behavior. TRB differs by its control surface, rather than changing the target distribution, divergence, or supervised positions.

  • From offline KD to OPD: Classical knowledge distillation uses fixed data and creates exposure bias, whereas GKD and OPD place teacher supervision on student-generated trajectories.MiniLLM further argues that…
  • Stabilizing teacher supervision: Veto changes the target distribution, Entropy-Aware OPD changes the divergence, and TIP changes where supervision is concentrated.These methods address stabilization through different intervention points.
  • Bridging the student–teacher gap: TRB differs in its control surface from related methods that modify targets, divergences, supervision locations, sampling, or distillation intermediates.The passage identifies this distinction as central to the method’s framing.
  • Bridging the student–teacher gap: SKD, MiCoTA, and Li et al. address student–teacher gaps through accepted teacher samples, intermediate assistants and traces, or trajectory-level signal compatibility.Li et al. argue that OPD succeeds only when student-visited states carry compatible and transferable teacher signal.

4 Trust-Region behavior Blending

TRB collects rollout prefixes with a teacher-guided behavior policy that stays within a KL trust region centered on the current student, then applies the unchanged reverse-KL OPD update. Its KL budget anneals to zero during warmup, transitioning from teacher guidance to pure student sampling.

  • Behavior policy: TRB moves the sampling policy toward the teacher within an explicit KL trust region around the current student before reverse-KL OPD updates.The resulting prefixes are used in the reverse-KL OPD update from Section 2.
  • Behavior policy: The behavior policy selects the most teacher-like distribution subject to a student-centered local KL-deviation constraint.The objective pulls sampling toward teacher-supported tokens while bounding off-policy deviation from the current student.
  • Trust-region solution: The closed-form behavior policy is µ∗(· | h) = µβ∗(h)(· | h), where β∗(h) is the largest feasible teacher-mixing coefficient.β controls movement toward the teacher, and Zβ(h) normalizes the distribution.
  • Trust-region solution: At ε = 0, TRB uses the student; when DKL(πT ∥πS) ≤ε, it uses the teacher, otherwise it finds β∗(h) by binary search.The binary search is justified because DKL(µβ∥πS) is monotone in β.
  • Warmup schedule: During warmup, the KL budget shrinks linearly to zero, shifting rollout collection from greater teacher guidance back to pure student sampling.TRB uses initial KL budget ε0 and warmup horizon K as its two method hyperparameters.

5 Experiments & Results

Experiments test whether limited early behavior-side guidance improves OPD across two model-pair settings. TRB achieves the strongest average results, while trajectory and prefix analyses indicate that its main behavioral differences arise during warmup.

  • Experimental setup: Experiments compare TRB with vanilla OPD and stronger or more persistent off-policy baselines across two Qwen3 model-pair settings.The settings distill Qwen3-1.7B-Base from Qwen3-8B and Qwen3-0.6B-Base from Qwen3-4B under a shared protocol.
  • Benchmark results: TRB attains the best average pass@1 score in both model-pair settings and outperforms fixed-ε blending in both.The methods use the same per-prefix solver, while TRB anneals its budget and fixed-ε blending does not.
  • Training dynamics: Several interventions rise faster than vanilla OPD at the start of training on the Qwen3-0.6B-Base ←Qwen3-4B setup.For the plotted SKD run, only about a 0.0093 fraction of generated tokens is replaced by the teacher at the first training step.
  • Training dynamics: Under TRB, teacher token-mean entropy is lower during warmup, then largely aligns with vanilla OPD while the benchmark curve remains higher.The main teacher-side difference therefore appears during warmup rather than after training returns to pure student rollouts.
  • Prefix analysis: TRB prefixes yield higher success than vanilla-OPD prefixes across all tested truncation lengths under both teacher and student continuation models.The controlled step-0 probe changes only the prefix source while holding truncation length and continuation model fixed.

6 Discussion

TRB achieves the strongest average in Table 1 while operating only during warmup. Comparisons with fixed-ε and Figures 2–4 suggest that teacher-guided off-policy behavior is more effective as a warmup intervention than when maintained throughout training, and that faster early improvement alone does not ensure the strongest final result.

  • TRB gives the strongest average in Table 1 while acting only during warmup.
  • The fixed-ε comparison suggests teacher-guided off-policy behavior is more useful during warmup than throughout the full run.The same local solver performs better as a warmup than when it remains active throughout training.
  • Figure 2 suggests faster early improvement or a more direct intervention does not by itself produce the strongest final result.

7 Limitations

The study is limited to two math-reasoning OPD settings with Qwen3-Base student–teacher pairs and correctness-based evaluation, so transfer to other domains or teacher–student gaps is unclaimed. TRB also raises warmup training-time cost through online teacher decoding and student–teacher co-residency.

  • Scope: The evaluation covers only two math-reasoning OPD settings using Qwen3-Base student–teacher pairs and correctness-based evaluation.The study does not claim that the same warmup schedules transfer unchanged to other domains or teacher–student gaps.
  • Overhead: TRB increases training-time cost during warmup because it requires online teacher decoding and student–teacher co-residency.Appendix C analyzes this overhead.

A Experimental Details · A.1 Evaluation Protocol · A.2 Baseline Setup Details

The experiments use a fixed reverse-KL OPD objective while varying rollout behavior during warmup, with evaluation based on pass@1 across mathematical-reasoning benchmarks. Baselines include pure student rollouts, veto, interleaved teacher injection, temperature warmup, fixed-ε blending, and SFT warmup.

  • A Experimental Details: Training uses verl with SGLang rollouts, FSDP2, and 8 NVIDIA H100 GPUs, while sampling 25,600 prompts from OpenThoughts3-1.2M.The system prompt asks models to reason step by step and place final answers within \boxed{}.
  • A Experimental Details: The reverse-KL OPD objective remains fixed across main experiments, with only rollout behavior varied during warmup.This isolates rollout-policy effects from changes to the OPD training objective.
  • A.1 Evaluation Protocol: Evaluation measures mathematical reasoning with pass@1, estimated as c/n for n sampled generations containing c correct answers and averaged over problems.The evaluation budget is large enough to stabilize checkpoint-to-checkpoint comparisons.
  • A.1 Evaluation Protocol: The Qwen3-1.7B-Base ←Qwen3-8B setup is evaluated on MATH500, AIME24, AIME25, AMC, and Olympiad.These benchmarks are listed for the larger student-teacher distillation setting.
  • A.1 Evaluation Protocol: TRB linearly anneals the student-centered KL budget from ε0 to 0, then switches to pure student decoding after warmup, unlike fixed-ε blending.TRB solves per-prefix teacher strength by bisection, whereas fixed-ε keeps its budget active throughout training.
  • A.2 Baseline Setup Details: Vanilla OPD uses pure student rollouts throughout, while veto varies its initial veto coefficient and otherwise keeps OPD hyperparameters fixed.Both baselines inherit the common training stack.
  • A.2 Baseline Setup Details: Interleaved teacher injection replaces student tokens outside the teacher top-K set with fresh teacher samples, while temperature warmup linearly returns rollout temperature to 1.0.Temperature warmup ends at step 15 or 25 before ordinary OPD decoding resumes.
  • A.2 Baseline Setup Details: SFT warmup uses the first 50 OPD prompt steps, with batches of 64 prompts and 4 teacher responses per prompt, then initializes ordinary OPD.Variants use SFT checkpoints after 15, 25, or 50 supervised updates and preserve the corresponding prompt stream and response multiplicity.

B Extended Results · B.1 Additional Warmup Diagnostics · B.2 Illustrative Early Rollouts

Extended results show that TRB consistently outperforms SKD across both model-pair sweeps, while diagnostics link larger trust-region budgets to higher reward and teacher log-probability but lower teacher-support separability. A qualitative early-rollout example further shows TRB preserving prompt structure when pure student behavior quickly goes off-task.

  • B Extended Results: Figures 5 and 6 compare TRB, SKD, and vanilla OPD across sweeps on the two main model-pair settings.Each plotted point is the best-over-training mean score for one hyperparameter setting; persistent fixed-ε blending is omitted because it appears in Table 1 and Figure 2.
  • B Extended Results: The Qwen3-1.7B-Base ←Qwen3-8B sweep groups TRB by warmup horizon and initial budget, SKD by K and teacher temperature τT, and marks vanilla OPD with a dashed red line.Each point represents the best-over-training mean score for one hyperparameter setting.
  • B Extended Results: The Qwen3-0.6B-Base ←Qwen3-4B sweep uses the same best-over-training mean-score comparison and grouping scheme for TRB, SKD, and vanilla OPD.TRB is grouped by warmup horizon and initial budget, while SKD is grouped by K and teacher temperature τT.
  • B Extended Results: Across both setups, the strongest TRB settings exceed the strongest SKD settings, with much of the SKD sweep below the TRB range.On the smaller setup, SKD beats vanilla OPD in only one configuration and does not change Table 1’s overall ranking.
  • B.1 Additional Warmup Diagnostics: As trust-region budget ε increases during the first 25 warmup steps, mean teacher log-probability and mean verifier reward increase, while teacher-support-score AUROC decreases.Figure 7 varies only ε on the Qwen3-1.7B-Base ←Qwen3-8B setup, separating reward level from teacher-support separability.
  • B.2 Illustrative Early Rollouts: In a prompt-matched first-warmup-step example, the pure-student rollout drifts off-task almost immediately, whereas TRB remains attached to the prompt’s arithmetic structure.The example is included as a qualitative sanity check rather than quantitative evidence.

C Efficiency Analysis · D EOS Canonicalization under Tokenizer Mismatch · E Derivation of the Trust-Region Solution

TRB changes teacher computation from post-generation batching to online decoding, increasing peak generation-time memory while preserving the teacher-side FLOP order. EOS canonicalization aligns mismatched stop tokens, and the trust-region optimizer yields a monotone KL-constrained interpolation whose coefficient is found by binary search.

  • C Efficiency Analysis: TRB queries the teacher online during decoding and reuses its statistics for the per-prefix reverse-KL term, rather than computing them afterward in a batched pass.This shifts teacher computation into sequential decoding and increases peak generation-time memory.
  • C Efficiency Analysis: The dominant extra memory comes from resident teacher weights and KV cache; releasing teacher state after warmup restores the student-only generation profile.Both student and teacher must remain resident during online policy construction.
  • C Efficiency Analysis: Trust-region search, interpolation-coefficient bisection, and log-space blending add little computation relative to transformer forwards, while teacher FLOPs remain the same order as standard OPD.The teacher is evaluated once per generated token, sequentially rather than in a batched pass.
  • D EOS Canonicalization under Tokenizer Mismatch: EOS canonicalization maps student and teacher EOS tokens to one shared event before sampling or KL evaluation, preventing the semantic stop event from being split across coordinates.The aligned sampler emits the native student EOS token when the shared event is selected.
  • E Derivation of the Trust-Region Solution: Introducing Lagrange multipliers for the student-centered KL constraint and normalization gives (1 + η) log µ(a) = log πT(a) + η log πS(a) + c.The resulting exponents satisfy β = 1/(1 + η) and 1 − β = η/(1 + η), recovering the interpolation family in Eq. 2.
  • E Derivation of the Trust-Region Solution: The trust-region family is parameterized by β ∈ [0, 1], and the implementation seeks the largest feasible β under the student-centered KL constraint.Binary search is applicable because the relevant KL map is monotone nondecreasing.
  • E Derivation of the Trust-Region Solution: Differentiation shows DKL(µβ ∥πS) is monotone nondecreasing in β, so the feasible set is an interval and β∗ can be found by binary search on [0, 1].This establishes the validity of the solver’s one-dimensional search.

F Small-Budget Efficiency of Trust Regions · G Sequence-Level Control from Token-Level Trust Regions

Trust-region warmup is locally efficient: small movement toward the teacher reduces teacher KL at first order while incurring only second-order student behavior-KL cost, yielding a √ε improvement for budget ε. At sequence level, token-wise KL budgets provide an additive upper bound on rollout-level divergence, including a uniform-budget corollary.

  • F Small-Budget Efficiency of Trust Regions: For small budgets, moving slightly from the student toward the teacher reduces teacher KL at first order while behavior-KL cost grows only second order.This establishes a favorable local trade-off for the trust-region blend path.
  • F Small-Budget Efficiency of Trust Regions: The small-budget analysis applies at each fixed prefix, where the trust-region path moves behavior toward the teacher while remaining centered on the student.The derivation assumes common support and excludes the degenerate case in which the two policies agree up to normalization.
  • F Small-Budget Efficiency of Trust Regions: At a fixed prefix, a small KL budget ε improves teacher closeness by order √ε while costing ε in student-centered behavior KL.The result characterizes trust-region warmup as efficient under a student-centered budget.
  • G Sequence-Level Control from Token-Level Trust Regions: For fixed-length rollouts, the rollout-level KL between TRB behavior and the student decomposes into the sum of token-level conditional KL terms.The same argument extends to stopped sequences by padding with an absorbing EOS state.
  • G Sequence-Level Control from Token-Level Trust Regions: If each prefix’s token-level behavior KL is bounded by ε̄_t, rollout-level KL is bounded by the sum of the ε̄_t budgets.The bound follows by replacing every local KL term in the decomposition with its uniform upper bound.
  • G Sequence-Level Control from Token-Level Trust Regions: With the same budget ε̄_t = ε at every step, the sequence-level divergence is controlled by accumulating that per-step budget across the rollout.This gives a direct sequence-level consequence of token-level trust-region constraints.
Loading 2605.31159v1…