Trust-Region Behavior Blending (TRB)

• Trust-Region Behavior Blending (TRB) is a method that blends teacher and student policies within a KL trust region to generate high-quality training prefixes.
• The method uses an annealed KL budget to gradually transition from teacher-guided rollouts to pure student rollouts, mitigating early instability.
• Empirical evaluations in mathematical reasoning tasks show that TRB outperforms other OPD baselines by significantly improving pass@1 accuracy.

Trust-Region Behavior Blending (TRB) is a warmup methodology for on-policy distillation (OPD) designed to address the limitations of early student rollouts during knowledge distillation from a stronger teacher model. Unlike offline distillation which suffers from prefix mismatch and exposure bias, OPD uses the student policy to generate training prefixes. However, in early stages, student-generated prefixes may be low-quality—impairing the effectiveness of teacher supervision. TRB proposes the use of a blended behavior policy, defined within a Kullback–Leibler (KL) trust-region centered on the student, to sample higher-quality prefixes that are as close as possible to the teacher’s distribution but constrained not to move beyond a preset KL divergence from the student. The KL budget is annealed to zero, so training transitions back to pure student rollouts. Empirical results demonstrate that TRB yields superior pass@1 accuracy in mathematical reasoning tasks relative to various strong OPD baselines (Plyusov et al., 29 May 2026).

1. On-Policy Distillation and Motivation

On-policy distillation (OPD) minimizes the reverse-KL divergence between a teacher policy π_T and a student policy π_S(·; θ), evaluated over prefixes h sampled from the student’s own rollouts:

$$L_{\text{OPD}}(\theta) = \mathbb{E}_{h\sim P_{\pi_S}} \left[ D_{\mathrm{KL}} \left( \pi_T(\cdot|h) \parallel \pi_S(\cdot|h; \theta) \right)\right].$$

This methodology avoids the prefix mismatch of offline distillation but introduces instability early in training due to weak student-generated prefixes. Consequently, applying strong teacher supervision on these low-quality prefixes can slow or destabilize optimization.

Trust-Region Behavior Blending (TRB) seeks to mitigate this by replacing the pure student rollout policy with a blended behavior policy during an initial warmup phase. The behavior policy is constructed to closely match the teacher while being restricted within a KL trust region centered at the student.

2. Mathematical Framework of TRB

The per-prefix reverse-KL loss is unchanged:

$$L(\theta) = \mathbb{E}_{h \sim P_{\pi_\text{blend}}}\left[D_{\mathrm{KL}}\left(\pi_T(\cdot|h) \parallel \pi_S(\cdot|h; \theta)\right)\right].$$

For each prefix h, the behavior policy π_blend(·|h) is the solution to:

$$\mu^* = \arg\min_{\mu} D_{\mathrm{KL}}(\mu \parallel \pi_T(\cdot|h)) \quad \text{subject to} \quad D_{\mathrm{KL}}(\mu \parallel \pi_S(\cdot|h)) \leq \epsilon_t,$$

with normalization constraints $\sum_a \mu(a) = 1$, $\mu(a) \geq 0$. The KL budget $\epsilon_t$ is time-varying, enabling gradual transition from teacher-guided to student-only rollouts.

For a warmup horizon K, the KL budget follows a linear schedule:

$$\epsilon_t = \epsilon_0 \cdot \left(1 - \frac{t}{K}\right), \quad t = 0,1,\ldots,K,$$

where $\epsilon_0$ is the initial KL budget and $\epsilon_t \to 0$ as $t \to K$.

3. Behavior Policy Derivation and Construction

The closed-form solution for $$L(\theta) = \mathbb{E}_{h \sim P_{\pi_\text{blend}}}\left[D_{\mathrm{KL}}\left(\pi_T(\cdot|h) \parallel \pi_S(\cdot|h; \theta)\right)\right].$$0 lies within a one-parameter geometric family:

$$L(\theta) = \mathbb{E}_{h \sim P_{\pi_\text{blend}}}\left[D_{\mathrm{KL}}\left(\pi_T(\cdot|h) \parallel \pi_S(\cdot|h; \theta)\right)\right].$$1

Alternatively, in exponential coordinates:

$$L(\theta) = \mathbb{E}_{h \sim P_{\pi_\text{blend}}}\left[D_{\mathrm{KL}}\left(\pi_T(\cdot|h) \parallel \pi_S(\cdot|h; \theta)\right)\right].$$2

where $$L(\theta) = \mathbb{E}_{h \sim P_{\pi_\text{blend}}}\left[D_{\mathrm{KL}}\left(\pi_T(\cdot|h) \parallel \pi_S(\cdot|h; \theta)\right)\right].$$3. The blending ratio $$L(\theta) = \mathbb{E}_{h \sim P_{\pi_\text{blend}}}\left[D_{\mathrm{KL}}\left(\pi_T(\cdot|h) \parallel \pi_S(\cdot|h; \theta)\right)\right].$$4 is selected so that $$L(\theta) = \mathbb{E}_{h \sim P_{\pi_\text{blend}}}\left[D_{\mathrm{KL}}\left(\pi_T(\cdot|h) \parallel \pi_S(\cdot|h; \theta)\right)\right].$$5. In practice, a binary search over $$L(\theta) = \mathbb{E}_{h \sim P_{\pi_\text{blend}}}\left[D_{\mathrm{KL}}\left(\pi_T(\cdot|h) \parallel \pi_S(\cdot|h; \theta)\right)\right].$$6 is used to meet the KL constraint per prefix.

4. Training Procedure and Pseudocode

The TRB warmup and annealing procedure can be concisely described as follows.

$$L(\theta) = \mathbb{E}_{h \sim P_{\pi_\text{blend}}}\left[D_{\mathrm{KL}}\left(\pi_T(\cdot|h) \parallel \pi_S(\cdot|h; \theta)\right)\right].$$9

Hyperparameters include the initial KL budget $$L(\theta) = \mathbb{E}_{h \sim P_{\pi_\text{blend}}}\left[D_{\mathrm{KL}}\left(\pi_T(\cdot|h) \parallel \pi_S(\cdot|h; \theta)\right)\right].$$7 and warmup horizon $$L(\theta) = \mathbb{E}_{h \sim P_{\pi_\text{blend}}}\left[D_{\mathrm{KL}}\left(\pi_T(\cdot|h) \parallel \pi_S(\cdot|h; \theta)\right)\right].$$8. Other training settings—such as optimizer (AdamW), learning rate (1e−5), batch size, and top-k (k=16) for reverse-KL—are as in standard OPD.

5. Empirical Evaluation and Ablations

TRB was empirically validated in two math-reasoning knowledge distillation setups:

1. Qwen3-1.7B student distilled from Qwen3-8B teacher: Evaluated on MATH500, Olympiad, AMC, AIME24, AIME25
2. Qwen3-0.6B student distilled from Qwen3-4B teacher: Evaluated on GSM8K, MATH500, AMC, Olympiad

The principal metric is pass@1 mean accuracy under a large decoding budget. The following baselines were compared:

Baseline	Description
Vanilla OPD	Pure student rollouts
Fixed-ε blending	TRB solver with constant KL budget
Veto	Adaptive target reformulation
SKD	Interleaved teacher token injection
Temperature warmup	Gradual rollout temperature schedule
SFT warmup	Offline supervised pre-training

TRB achieves the highest average pass@1 in both settings (e.g., 33.2% vs. 32.3% for vanilla OPD on 1.7B→8B), consistently outperforming all baselines. Fixed-ε blending underperforms TRB, indicating that annealing the KL budget is crucial. Ablation studies confirm that only annealed TRB (rather than persistent fixed-ε) yields the strongest final accuracies, as teacher-guided rollouts are beneficial early but detrimental if continued throughout training. Optimal results were found around ε₀ ≈ 0.01 and K ≈ 50 steps.

6. Significance and Context

Trust-Region Behavior Blending provides a principled mechanism to blend teacher supervision in early OPD training by sampling from the most teacher-like policy within a student-centered KL ball, annealing the radius to zero, and then reverting to pure student rollouts. The approach does not alter the core OPD objective but systematically guides the sampling distribution to increase early prefix quality. This yields robust downstream improvements in mathematical reasoning accuracy, especially for model compression and student model initialization in knowledge distillation regimes (Plyusov et al., 29 May 2026). The clear demonstration that annealing—and not persistent constraint—is fundamentally important provides substantive guidance for future OPD approaches.