Trust-Region Behavior Blending for On-Policy Distillation

Abstract: 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.

• The paper introduces TRB, which uses a warmup strategy where a blended teacher-student behavior is constrained by a KL trust region to overcome early prefix mismatch.
• It details a closed-form policy blending method and evaluates TRB on Qwen3 student-teacher setups, consistently achieving superior pass@1 scores on math reasoning benchmarks.
• Empirical results and theoretical analysis confirm that early, strictly controlled teacher guidance significantly improves the overall efficacy of on-policy distillation.

Trust-Region Behavior Blending for On-Policy Distillation: Methodology and Empirical Study

Introduction and Motivation

On-policy distillation (OPD) has become the dominant technique for transferring knowledge from large teacher LLMs to smaller student models by collecting supervision over prefixes generated by the student itself. This avoids the distributional shift between training and inference present in standard offline knowledge distillation, where teacher-forced or teacher-generated prefixes dominate. However, in the early stages of OPD, the student typically samples low-quality, misaligned prefixes, restricting the effectiveness of teacher supervision and potentially impeding convergence.

The "Trust-Region Behavior Blending" (TRB) framework addresses this issue by proposing a principled warmup strategy: during initial training steps, rollouts are collected from a behavior policy that interpolates between student and teacher actions but is strictly constrained to remain within a prescribed KL trust region centered on the student. By annealing this trust region to zero over the course of warmup, the method ensures that training reverts to standard (pure student) OPD, but after having benefited from improved prefix quality and more informative supervision in the early regime.

Figure 1: Schematic of TRB—at each prefix, the behavior policy is the most teacher-like distribution within a KL constraint from the student policy.

Methodology: Trust-Region Behavior Policy Construction

Given a student policy $\pi_S$ and teacher policy $\pi_T$, for each prefix $h$, TRB solves:

$$\mu^*(\cdot|h) = \arg\min_{\mu} D_{\mathrm{KL}}(\mu \| \pi_T) \text{ s.t. } D_{\mathrm{KL}}(\mu \| \pi_S) \leq \varepsilon, \sum_a \mu(a) = 1, \mu(a) \geq 0$$

This yields a closed-form blended policy family:

$$\mu_\beta(a|h) = \frac{\pi_S(a|h)^{1-\beta} \pi_T(a|h)^\beta}{Z_\beta(h)}$$

where $\beta \in [0,1]$ is the largest value for which the KL constraint is met. Annealing $\varepsilon$ from an initial value $\varepsilon_0$ to 0 over a fixed warmup horizon $K$ ensures a smooth transition: training begins with more teacher-guided rollouts and ends with pure student sampling.

This method maintains the standard reverse-KL OPD loss but modulates the prefix distribution, distinguishing it from approaches that modify the target distribution, inject teacher tokens, or alter the loss.

Experimental Design

The evaluation benchmarks two primary distillation scenarios:

• Qwen3-1.7B-Base student distilled from Qwen3-8B teacher
• Qwen3-0.6B-Base student distilled from Qwen3-4B teacher

Baselines include:

• Vanilla OPD: Standard student rollouts with fixed reverse-KL loss
• Veto: Adaptive target reformulation at each visited prefix
• SKD: Speculative Knowledge Distillation, which injects teacher tokens when student actions fail a teacher-side rule
• Temperature Warmup and SFT Warmup: Heuristic approaches lowering student temperature or inserting supervised warmup
• Fixed-$\varepsilon$ Blending: Like TRB but with non-annealed trust region

The main metric is pass@1 on math reasoning benchmarks (MATH500, Olympiad, AMC, AIME, GSM8K), and ablations cover a wide range of hyperparameters and warmup lengths.

Empirical Results

Extensive experiments demonstrate that TRB consistently achieves the strongest average pass@1 scores across both model sizes and all benchmark suites, outperforming both traditional OPD and alternative teacher-guidance strategies.

Figure 2: Training curves for several interventions. TRB and SKD demonstrate faster early progress than vanilla OPD; only TRB consistently secures the best final performance.

TRB changes the behavior policy only during warmup, yet produces a persistent improvement. Early in training, SKD and temperature warmup can yield rapid ascent but fail to match the sustained gains of TRB in later trajectories.

Figure 3: Teacher token entropy and evaluation pass@1 over the course of training. TRB reduces teacher entropy (higher teacher confidence) during warmup.

Teacher token entropy on visited prefixes confirms that TRB improves prefix quality in the early regime, with lower entropy reflecting higher teacher certainty on student-reachable trajectories. After warmup, teacher entropy and student performance plateau alongside vanilla OPD, indicating that the prefix correction is transient.

Probing early rollouts reveals that prefixes generated with TRB yield superior continuation success rates for both teacher and student models, confirming that the prefixes are not only more aligned but also exhibit improved transferability of teacher signal.

Figure 4: TRB-initiated prefixes exhibit higher continuation success, for both teacher and student, at all tested truncation lengths.

Hyperparameter sweeps further show that:

• TRB’s warmup achieves better results than persistent (non-annealed, fixed-$\pi_T$0) blending.
• Even when alternative methods match early growth, they often fail to achieve comparable final scores.

  Figure 5: Performance sweeps on the 1.7B-to-8B setting for TRB, SKD, and vanilla OPD. The best TRB settings surpass all other methods.

  

  Figure 6: Analogous performance sweeps on 0.6B-to-4B; TRB is robustly superior across parameterizations.

Theoretical and Practical Implications

Theoretical analysis shows the TRB blend path is a form of exponential tilting, providing efficient teacher supervision without sacrificing coverage of student-reachable states. For small KL budgets, moving toward the teacher is first-order efficient—significantly improving teacher alignment for minimal off-policy cost. Sequence-level KL control is established via additive bounds from token-level constraints across the rollout.

Practically, this work suggests that early but strictly controlled off-policy guidance is necessary and sufficient to overcome student-teacher trajectory mismatch in OPD, negating the need for persistent teacher injection or ad hoc heuristics. The KL trust region ensures that rollouts remain student-reachable, thereby preserving learnability and avoiding over-regularization.

Limitations and Future Directions

The experimental scope is math-centric Qwen3 student-teacher pairs; it remains to be established how TRB scales to broader domains, larger teacher gaps, vision-LLMs, or data distributions with sparser supervision. Furthermore, TRB incurs elevated compute and memory costs during warmup due to online teacher-student co-decoding, with transient but material overhead. Addressing efficiency for larger teacher models or devising asynchronous approximations are relevant next steps.

The approach also leaves open optimal annealing schedules and adaptive strategies for determining when the trust region can be safely collapsed, as well as the interaction of trust-region blending with alternative divergence choices or RL-style objectives.

Conclusion

TRB provides a formal and empirically validated solution to the challenge of early prefix mismatch in on-policy distillation. By collecting rollouts from a teacher-guided behavior policy constrained by a student-centered KL trust region, TRB dramatically improves final pass@1 distillation performance across a suite of math reasoning benchmarks without requiring persistent off-policy intervention. The method is straightforward to implement, theoretically sound, and effective, with immediate relevance for practitioners scaling LLM distillation in the presence of teacher-student gaps.

Reference: "Trust-Region Behavior Blending for On-Policy Distillation" (2605.31159)