Reinforcement-aware Knowledge Distillation for LLM Reasoning

Technical Summary: Reinforcement-aware Knowledge Distillation for LLM Reasoning (RLAD)

1. Problem Statement

Large Language Models (LLMs) have achieved significant reasoning capabilities through Reinforcement Learning (RL) post-training, particularly in long chain-of-thought (CoT) tasks. However, the high inference costs of these large models necessitate distillation into smaller student models. Existing Knowledge Distillation (KD) methods face critical limitations when applied to RL post-training:

• Distribution Mismatch: Standard KD methods (offline, off-policy, or on-policy) often rely on fixed teacher traces or static KL divergence regularization. In RL, the student's policy evolves continuously, causing the rollout distribution to shift. Static distillation targets become suboptimal or unstable as the student moves away from the teacher's initial distribution.
• Objective Interference: Methods that integrate KD into the RL loop (e.g., KDRL) typically add a KL divergence term between the student and teacher as a regularization penalty. This creates a multi-objective optimization problem where the reward maximization objective competes with the imitation objective. This competition requires careful hyperparameter tuning (loss balancing) and can lead to unstable training dynamics, as the teacher supervision may not align with the student's current advantage estimates.

2. Methodology: RL-Aware Distillation (RLAD)

The authors propose RLAD, a framework designed to perform selective imitation during RL post-training. The core principle is that the student should only be guided by the teacher when such guidance improves the current policy update under the RL objective.

Key Component: Trust Region Ratio Distillation (TRRD)
Instead of using a standalone KL regularizer, RLAD introduces TRRD, which replaces the standard importance sampling ratio in Group Relative Policy Optimization (GRPO) with a teacher-student mixture.

• Mechanism: TRRD defines a new likelihood ratio $r^{TRRD}_{i,t}$ for the $t$-th token of the $i$-th response:
  $$r^{TRRD}_{i,t}(\theta_S) = \left( \frac{\pi_{\theta_S}(y_t | x, y_{<t})}{\pi_{\theta_S,old}(y_t | x, y_{<t})} \right)^\alpha \left( \frac{\pi_{\theta_S}(y_t | x, y_{<t})}{\pi_{\theta_T}(y_t | x, y_{<t})} \right)^{1-\alpha}$$
  Here, $\pi_{\theta_S}$ is the current student policy, $\pi_{\theta_S,old}$ is the previous student policy, and $\pi_{\theta_T}$ is the teacher policy. $\alpha \in [0, 1]$ is a mixing coefficient.
• Trust Region Interpretation: This ratio can be viewed as a PPO-style update where the trust region is anchored to a mixture of the old student and the teacher ($r_{\pi_{mix}}$), rather than just the old student.
• Advantage-Aware Selection: The update is weighted by the group-relative advantage $\hat{A}_{i,t}$.
  • Positive Advantage: If the advantage is positive, the student is encouraged to increase the probability of the token. If the teacher also assigns high probability, the mixture anchor is higher, allowing the student to increase its probability more before hitting the clipping bound (allowing larger updates). If the teacher assigns low probability, the anchor is lower, restricting the update magnitude.
  • Negative Advantage: If the advantage is negative, the student is encouraged to decrease the probability. If the teacher assigns high probability, the clipping bound prevents the student from aggressively suppressing the token, preserving teacher-aligned behaviors.
  • Near-Zero Advantage: Teacher influence becomes negligible, preventing interference when the reward signal is uninformative.

This design implicitly balances exploration, exploitation, and imitation by ensuring teacher guidance is modulated by the reward signal (advantage) and constrained within a trust region.

3. Key Contributions

1. RLAD Framework: A new distillation framework for RL post-training that integrates teacher guidance via Trust Region Ratio Distillation (TRRD), replacing static KL regularization with a clipped, advantage-weighted likelihood-ratio objective.
2. Theoretical Analysis: An analysis demonstrating that TRRD enables reward-conditioned, selective imitation. It preserves trust-region-bounded updates and naturally balances the trade-off between following the teacher and optimizing for rewards without requiring delicate loss balancing.
3. Empirical Validation: Extensive experiments on logical reasoning (K&K Logistics) and long-context mathematical reasoning (AIME, AMC, MATH) showing consistent improvements over baselines.

4. Experimental Results

The authors evaluated RLAD on logical reasoning and complex math reasoning tasks using student models ranging from 0.6B to 8B parameters and teachers up to 32B.

• Logical Reasoning (K&K Logistics):

  • Using a Qwen3-8B teacher, RLAD significantly improved average accuracy for Qwen3-0.6B (from 0.76 to 0.94 at 8K context) and Qwen3-1.7B (from 0.95 to 0.99 at 8K context) compared to standard GRPO.
  • Gains were most pronounced on the hardest subsets (PPL7 and PPL8), where RLAD maintained a clear margin over KDRL.
  • Training dynamics showed faster convergence and higher final rewards compared to both GRPO and KDRL.

• Mathematical Reasoning:

  • Base Models (R1-Zero style): On Qwen3-8B-Base, RLAD improved average performance by +5.5 points over GRPO on long-context (30K) tasks.
  • Post-trained Models: RLAD outperformed explicit-KL distillation (KDRL) and offline distillation.
  • Metric Specifics: RLAD showed larger gains on Pass@1 (Mean@32) compared to KDRL, which tended to show larger relative gains on Pass@32. The authors interpret this as RLAD successfully driving reward-aligned policy improvement (exploitation), whereas KDRL's gains were driven more by pure imitation of diverse teacher solutions.
  • Stability: RLAD exhibited more stable training dynamics with fewer oscillations compared to KDRL, which suffered from objective coupling issues.

• Robustness and Efficiency:

  • Weak Teacher: When the teacher was weaker than the student, RLAD remained close to GRPO performance, whereas KDRL degraded more significantly.
  • Hyperparameters: RLAD was found to be largely insensitive to the mixing coefficient $\alpha$ (fixed at 0.5), unlike KDRL which is sensitive to the KL weight $\lambda$.
  • Efficiency: Integrating RLAD added only ~12% training latency compared to standard GRPO, as the teacher is only used for log-probability computation, not trajectory generation.

5. Significance and Claims

The paper claims that RLAD addresses the fundamental misalignment between static teacher supervision and dynamic RL policy updates. By folding teacher guidance into the trust-region update mechanism, RLAD avoids the "brittle tuning" required by multi-objective KL regularization.

The authors assert that RLAD allows the student to move toward high-advantage behaviors within an anchor distribution while opportunistically exploiting teacher guidance. This results in:

• Superior Performance: Consistent outperformance of GRPO, KDRL, and offline distillation, particularly on challenging benchmarks.
• Stability: More stable training dynamics and better convergence.
• True RL Optimization: Unlike methods that rely heavily on imitation (which may boost Pass@32 but not Pass@1), RLAD prioritizes reward-aligned behavior, leading to stronger gains in the model's most probable correct responses (Pass@1).

The paper concludes that while the approach relies on open-weight, well-calibrated teacher models for logit access, it represents a significant step forward in making distillation compatible with the dynamic nature of RL post-training.