PaperlayerView exact revision on arXiv

Source-linked AI summary

Reinforcement-aware Knowledge Distillation for LLM Reasoning

Zhaoyang Zhang, Shuli Jiang, Yantao Shen, Yuting Zhang, Dhananjay Ram, Shuo Yang, Zhuowen Tu, Wei Xia, Stefano Soatto

arXiv:2602.22495v3cs.LGcs.AI

TL;DR

Distilling RL-trained reasoning LLMs is difficult because fixed-trace and KL-based methods can mismatch evolving student rollouts and interfere with reward optimization. RLAD uses advantage-aware, trust-region-bounded teacher guidance during RL and consistently outperforms GRPO, KDRL, and offline distillation across logic and math benchmarks.

  • Problem

    Fixed KL-based teacher guidance can imitate behaviors misaligned with high-reward trajectories, motivating selective teacher supervision during reinforcement learning.

  • Method

    RLAD uses TRRD, a clipped importance-ratio objective anchored on a teacher–old student mixture, to integrate selective imitation into advantage-weighted trust-region policy updates.

  • Results

    RLAD consistently outperforms GRPO, KDRL, and offline distillation across logical and long-context mathematical reasoning, with larger gains on harder benchmarks and improved training stability.

  • Takeaways & Limitations

    Advantage-aware distillation provides a practical way to balance teacher imitation with reward optimization during RL post-training for smaller reasoning students.

  • Takeaways & Limitations

    The framework assumes a smaller student policy is trained with access to a more powerful teacher policy, alongside teacher behavior information.

Abstract

from arXiv · show

Reinforcement learning (RL) post-training has recently driven major gains in long chain-of-thought reasoning large language models (LLMs), but the high inference cost of such models motivates distillation into smaller students. Most existing knowledge distillation (KD) methods are designed for supervised fine-tuning (SFT), relying on fixed teacher traces or teacher-student Kullback-Leibler (KL) divergence-based regularization. When combined with RL, these approaches often suffer from distribution mismatch and objective interference: teacher supervision may not align with the student's evolving rollout distribution, and the KL regularizer can compete with reward maximization and require careful loss balancing. To address these issues, we propose RL-aware distillation (RLAD), which performs selective imitation during RL -- guiding the student toward the teacher only when it improves the current policy update. Our core component, Trust Region Ratio Distillation (TRRD), replaces the teacher-student KL regularizer with a PPO/GRPO-style likelihood-ratio objective anchored to a teacher--old-policy mixture, yielding advantage-aware, trust-region-bounded distillation on student rollouts and naturally balancing exploration, exploitation, and imitation. Across diverse logic reasoning and math benchmarks, RLAD consistently outperforms offline distillation, standard GRPO, and KL-based on-policy teacher-student knowledge distillation.

1. Introduction

RL-aware distillation addresses the limitations of offline teacher-trace distillation by selectively incorporating teacher guidance during student reinforcement learning. Its TRRD objective uses a clipped likelihood ratio anchored to a teacher–old student mixture, and RLAD consistently improves reasoning performance across math and logic benchmarks.

  • Motivation: Scaling improves LLM generation and reasoning, but inference costs motivate compressing larger teachers into smaller students with knowledge distillation.KD targets improved quality under fixed inference budgets.
  • Limitations: Offline distillation first trains a strong teacher with RL, then SFT-trains the student on teacher-generated traces without student RL optimization or adaptive teacher guidance.This limits transfer of adaptive reasoning because the teacher does not respond to the student’s evolving policy.
  • RL-aware distillation: RLAD performs selective imitation during RL, guiding the student toward the teacher only when imitation benefits the current policy under the RL objective.The framework is tailored to RL post-training of LLMs.
  • Trust Region Ratio Distillation: TRRD replaces standalone KL regularization with a PPO/GRPO-style clipped likelihood-ratio objective anchored on a teacher–old student mixture.The objective is designed for trust-region-bounded, advantage-aware distillation on student rollouts.
  • Results: RLAD consistently improves over prior RL-based distillation baselines across math and logic benchmarks, using students from 0.5B–7B and teachers up to 32B.On K&K Logistics with a Qwen3-8B teacher, Qwen3-0.6B accuracy rises from 0.76 to 0.94 at 8K context and from 0.70 to 0.90 at 2K.
  • Results: Compared with KDRL, RLAD is distinguished by reward-driven policy improvement and more stable validation dynamics, whereas KDRL’s gains favor Pass@32 over Pass@1.The passage attributes KDRL’s pattern primarily to teacher imitation rather than reward-driven policy improvement.

2. Related Work

Prior work spans offline, off-policy, and on-policy knowledge distillation, alongside methods combining teacher supervision with RL reward optimization. RLAD differs by integrating teacher guidance into an RL trust region, applying imitation only when it benefits the current update rather than adding a competing explicit-KL objective.

  • Knowledge Distillation for LLMs: Offline distillation trains students with supervised fine-tuning on fixed teacher-generated traces and requires no teacher queries during student training.This paradigm is widely used for its simplicity and stability.
  • Knowledge Distillation for LLMs: Off-policy KD computes teacher supervision on data not generated by the current student policy, including static prompts and teacher-sampled sequences.Methods operate at sequence or token-logit level, with some using top-k logits to mitigate long-tail noise.
  • Knowledge Distillation for LLMs: On-policy KD learns from student-generated outputs augmented with teacher feedback, reducing exposure bias and better matching the student’s test-time distribution.Representative methods use policy gradients or divergence objectives such as RKLD and JSD.
  • Position of our RLAD: RLAD folds teacher guidance into an RL trust region, applying imitation only when it benefits the current update instead of adding an explicit-KL-style auxiliary objective.This design aims to reduce brittle tuning between imitation and exploration, addressing the mismatch between offline distillation and RL post-training dynamics.

3. Background and Notations

This section formalizes autoregressive language-model policies, rewards, and KL divergence, then defines the student–teacher reinforcement-learning objective. It also introduces GRPO and KDRL as the baseline optimization and teacher-regularization formulations.

  • Problem Statement: Autoregressive responses are modeled as stochastic policies over vocabulary tokens, with sequence probability factorized across token positions.The response is y = (y1, y2, . . . , yT) ∈ V^T, conditioned on prompt x and preceding tokens y<t.
  • KDRL: The KDRL KL term coincides with the reverse-KL distillation objective used in classical and SFT-based knowledge distillation.The KL divergence is defined as an expectation under πθ1 of the log-probability difference between πθ1 and πθ2.
  • Problem Statement: The student policy πθS maximizes expected reward over prompts while transferring knowledge from a more powerful teacher policy πθT.The prompt distribution is denoted D, and the reward objective is J(θS) = E[R(x, y)].
  • Group Relative Policy Optimization (GRPO): GRPO samples G responses from the current student policy, assigns scalar rewards, and normalizes them using prompt-level reward statistics.The method also uses an earlier student policy for importance sampling and a fixed reference policy with clipping and regularization controls.
  • KDRL: KDRL replaces GRPO’s fixed reference policy in the KL regularization term with the teacher policy πθT.This explicitly regularizes the student toward the teacher and yields an objective denoted JKDRL(θS).

4. RL-Aware Distillation (RLAD)

RLAD introduces TRRD, replacing unconditional teacher–student KL regularization with advantage-aware, clipped likelihood-ratio distillation anchored to a teacher–old-policy mixture. This selectively applies teacher guidance on student rollouts while preserving reward-directed updates and balancing exploration, exploitation, and imitation.

  • Fixed KL distillation uniformly guides the student, including where teacher behavior conflicts with high-reward trajectories, potentially degrading task performance.
  • TRRD incorporates the teacher policy into the GRPO importance ratio instead of imposing it through an unconditional KL penalty.The ratio interpolates token-level on-policy improvement and teacher guidance using α ∈[0, 1].
  • TRRD centers a PPO-style clipped update on a teacher–old-policy mixture rather than the previous student policy alone.Clipping bounds the ratio and constrains changes relative to the mixture anchor; α = 0 recovers standard GRPO, while α = 1 yields a fully teacher-anchored objective related to DPO.
  • TRRD implicitly combines KL regularization toward the previous student policy and teacher policy, activating imitation when supported by advantage while preserving the reward-determined update direction.This adaptive trust-region mechanism balances exploration, exploitation, and imitation rather than enforcing teacher alignment as a fixed constraint.
  • Teacher influence is advantage-modulated: positive or negative advantages adjust student probabilities under clipping, while near-zero advantage makes teacher influence negligible.The mixture anchor permits larger or smaller updates depending on the teacher’s assigned probability, while reward-uninformative tokens contribute little.

5. Experiments

Experiments evaluate RLAD on logical and mathematical reasoning, comparing it with GRPO, KL-based KDRL, and offline SFT distillation. They also examine optimization stability, contrasting KDRL’s potentially excessive early KL term with RLAD’s clipped trust-region updates.

  • RLAD is evaluated in two reasoning settings: logical reasoning and mathematical reasoning.
  • The experiments compare RLAD against GRPO without teacher guidance, KDRL with KL distillation, and offline SFT on reject-sampled teacher outputs.
  • Stability Considerations and Hyperparameter Ablations: KDRL’s student-teacher KL term can become excessively large early in training, potentially destabilizing optimization when policies are substantially mismatched.
  • Stability Considerations and Hyperparameter Ablations: RLAD mitigates such outliers by anchoring updates to a teacher–previous-student mixture and clipping the TRRD ratio within a trust region.

5.1. Logical Reasoning Task

On K&K logical reasoning, RLAD trains Qwen3-0.6B and Qwen3-1.7B with GRPO-based selective distillation from Qwen3-8B. It converges faster and achieves higher rewards and accuracy than vanilla GRPO and KDRL, with especially large gains on difficult subsets.

  • Training dynamics: RLAD achieves faster convergence and higher final training and validation rewards than vanilla GRPO and KDRL for both student sizes.The comparison concerns RL training dynamics on the logical reasoning task.
  • Results: +18% average accuracy for Qwen3-0.6B at 8K (0.94 vs. 0.76 under GRPO) and +4% for Qwen3-1.7B (0.99 vs. 0.95), with RLAD outperforming KDRL.At 2K, the corresponding gains are +20% for Qwen3-0.6B (0.90 vs. 0.70) and +7% for Qwen3-1.7B (0.93 vs. 0.86).
  • Results: On PPL8, RLAD improves Qwen3-0.6B from 0.63 to 0.83 at 8K and from 0.50 to 0.69 at 2K, compared with GRPO.These are the most challenging-subset results for Qwen3-0.6B.
  • Results: On PPL8, RLAD improves Qwen3-1.7B from 0.88 to 0.98 at 8K and from 0.52 to 0.75 at 2K, while maintaining a 6%–15% advantage over KDRL.All reported improvements are compared with GRPO.

5.2. Complex Math Reasoning Task

On complex math reasoning, RLAD improves average and hard-benchmark performance over GRPO across base and post-trained students, with faster convergence and greater gains on difficult tasks. Additional analyses show training stability, weak-teacher robustness, active trust-region clipping, and comparable or lower decoding variation.

  • Training dynamics: RLAD converges faster and reaches higher Mean@32 on AIME24, whereas explicit-KL distillation such as KDRL can show noticeable training instability.The instability is attributed to coupling reward maximization with the auxiliary KL regularizer.
  • Results: RLAD improves Avg by +1.9 for Qwen3-1.7B-Base and +5.5 for Qwen3-8B-Base over GRPO.The reported averages are 38.4 vs. 36.5 for Qwen3-1.7B-Base and 66.5 vs. 61.0 for Qwen3-8B-Base.
  • Results: AIME24-Pass@32 rises from 77.8 to 85.4 and AIME25-Pass@32 from 48.5 to 66.4 for Qwen3-8B-Base, while MATH500 changes from 68.2 to 68.6.For post-trained models, Avg also rises from 51.2 to 57.8 for Qwen3-1.7B and from 54.2 to 56.7 for Qwen2.5-1.5B-DS.
  • Weak-teacher robustness: With a weaker Qwen3-0.6B teacher for a Qwen3-1.7B student, RLAD stays close to GRPO while KDRL degrades more noticeably.TRRD shifts the trust-region center and clipping thresholds without adding an independent gradient term that can dominate the advantage direction.
  • Clipping-ratio behavior: RLAD clips sampled tokens more often than GRPO early in training, and the clipping frequency decreases over time.This pattern indicates that the teacher-centered trust-region constraint is active initially and fewer sampled tokens later violate the threshold.
  • Robustness and variance: RLAD has STD@32 comparable to or lower than GRPO in most settings and lower than KDRL on harder long-context math cases.STD@32 measures variation across 32 decoding runs rather than multi-seed training variance.

6. Conclusion

The paper proposes RLAD, a reinforcement-aware knowledge-distillation framework that integrates teacher guidance into advantage-weighted, trust-region policy updates. Across logical reasoning and long-context math tasks, RLAD outperforms GRPO and KDRL, with its largest gains on challenging AIME24/25 evaluations in Best@32.

  • Framework: RLAD is proposed as a framework for reinforcement-learning post-training of reasoning models.It targets reinforcement-aware knowledge distillation for reasoning.
  • Method: TRRD replaces static KL-style teacher regularization with teacher guidance integrated directly into advantage-weighted, trust-region policy updates.This is the core mechanism used by RLAD.
  • Results: RLAD consistently outperforms GRPO and KDRL across logical reasoning and long-context math reasoning tasks.The reported advantage is broad across both task categories.
  • Results: The largest gains occur on challenging evaluations such as AIME24/25 in Best@32.The passage identifies these evaluations as examples of the strongest improvements.

A. Ablation on Distillation Coefficients

The ablation study examines how RLAD’s mixing coefficient α and KDRL’s KL-loss weight λ affect distillation during GRPO training. It uses a Qwen3-1.7B student and Qwen3-8B teacher on GSM8K and MATH under a 1K context length.

  • A. Ablation on Distillation Coefficients: The study varies RLAD’s mixing coefficient α and KDRL’s KL-loss weight λ to examine their effects on math reasoning distillation.Both methods are evaluated during GRPO training.
  • A. Ablation on Distillation Coefficients: The experiments train a Qwen3-1.7B student with GRPO using a Qwen3-8B teacher for both KDRL and RLAD distillation.The same teacher is used while varying λ for KDRL and α for RLAD.
  • A. Ablation on Distillation Coefficients: The ablation is conducted on the GSM8K and MATH math reasoning benchmarks using their training splits and a 1K context length.GSM8K is attributed to Cobbe et al. (2021), and MATH to Hendrycks et al. (2021a).

B. Stabilizing TRRD Training with Clipping

TRRD training is unstable early because teacher and student policies may be substantially mismatched, making the token-level ratio approximation unreliable and producing extreme ratios. The method addresses this with PPO-style clipping to stabilize optimization.

  • Early in training, teacher and student policies can be substantially mismatched, unlike the typically close old and current policies in PPO/GRPO.
  • The token-level ratio is only a first-order Taylor approximation when teacher and student policies are close.
  • This mismatch can produce extreme distillation ratios, leading to unstable gradients and optimization.
  • TRRD stabilizes training by applying PPO-style clipping to the distillation ratios.

C. TRRD as an Interpolation Between GRPO and DPO

TRRD interpolates between purely on-policy GRPO and teacher-referenced imitation through the mixing coefficient α. As α approaches zero, its token-level updates recover the DPO logit structure while retaining on-policy, advantage-weighted, clipped optimization dynamics.

  • C. TRRD as an Interpolation Between GRPO and DPO: When α = 1, TRRD reduces to standard GRPO with the student’s old policy as anchor and purely on-policy reinforcement learning.The update remains driven by reward/advantage signals.
  • C. TRRD as an Interpolation Between GRPO and DPO: As α →0, TRRD approaches a teacher-referenced ratio πθS/πθT and becomes a KL-regularized imitation objective.The anchor increasingly trusts the teacher rather than the student’s old policy.
  • C. TRRD as an Interpolation Between GRPO and DPO: In the α →0 limit, TRRD’s token-level log-ratio yields updates depending on ∆log πθS −∆log πθT, matching the DPO logit.DPO’s logit is defined by the difference between student and teacher log-probability gaps on the same preference pair.
  • C. TRRD as an Interpolation Between GRPO and DPO: Despite the shared logit structure, DPO uses explicit pairwise sigmoid weighting, whereas TRRD uses on-policy ratio and advantage weighting with GRPO-style clipping.TRRD therefore preserves reinforcement-learning update dynamics rather than becoming an explicit pairwise preference loss.
Loading 2602.22495v3…