PaperlayerView exact revision on arXiv

Source-linked AI summary

KDRL: Post-Training Reasoning LLMs via Unified Knowledge Distillation and Reinforcement Learning

Hongling Xu, Qi Zhu, Heyuan Deng, Jinpeng Li, Lu Hou, Yasheng Wang, Lifeng Shang, Ruifeng Xu, Fei Mi

arXiv:2506.02208v1cs.LGcs.AIcs.CL

TL;DR

RL and KD offer complementary post-training approaches, but RL is compute-intensive and limited by base-model capacity, while KD is teacher-constrained and can generalize poorly out of domain. KDRL jointly optimizes on-policy reverse-KL distillation from a teacher and RL reward feedback through policy gradients. In experiments, KDRL largely outperforms standard SFT approaches and GRPO by 5.7% and 2.6%, respectively, and KD-RKL by 1.1%.

  • Problem

    RL and KD offer complementary post-training approaches, but RL is compute-intensive and limited by base-model capacity, while KD is teacher-constrained and can generalize poorly out of domain.

    Source: §1 · p. 1

  • Method

    KDRL jointly optimizes on-policy reverse-KL distillation from a teacher and RL reward feedback through policy gradients.

    Source: §1 · p. 2

  • Results

    In experiments, KDRL largely outperforms standard SFT approaches and GRPO by 5.7% and 2.6%, respectively, and KD-RKL by 1.1%.

    Source: §1 · p. 2

  • Takeaways & Limitations

    KDRL achieves greater training and reasoning-token efficiency while outperforming GRPO and various KD baselines in performance.

    Source: §6 · p. 10

Abstract

from arXiv · show

Recent advances in large language model (LLM) post-training have leveraged two distinct paradigms to enhance reasoning capabilities: reinforcement learning (RL) and knowledge distillation (KD). While RL enables the emergence of complex reasoning behaviors, it often suffers from low sample efficiency when the initial policy struggles to explore high-reward trajectories. Conversely, KD improves learning efficiency via mimicking the teacher model but tends to generalize poorly to out-of-domain scenarios. In this work, we present \textbf{KDRL}, a \textit{unified post-training framework} that jointly optimizes a reasoning model through teacher supervision (KD) and self-exploration (RL). Specifically, KDRL leverages policy gradient optimization to simultaneously minimize the reverse Kullback-Leibler divergence (RKL) between the student and teacher distributions while maximizing the expected rule-based rewards. We first formulate a unified objective that integrates GRPO and KD, and systematically explore how different KL approximations, KL coefficients, and reward-guided KD strategies affect the overall post-training dynamics and performance. Empirical results on multiple reasoning benchmarks demonstrate that KDRL outperforms GRPO and various KD baselines while achieving a favorable balance between performance and reasoning token efficiency. These findings indicate that integrating KD and RL serves as an effective and efficient strategy to train reasoning LLMs.

1 Introduction

RL can improve reasoning through self-exploration, while KD offers efficient teacher guidance; each also has constraints, and using them in separate stages decouples their benefits. KDRL jointly optimizes both with on-policy reverse-KL distillation and reward feedback, and the paper reports improved benchmark performance and efficiency.

1 Introduction

  • KD learns efficiently from a teacher but is constrained by teacher capability and can generalize poorly out of domain, while RL is compute-intensive and limited by base-model capacity.The paper describes RL as able to amplify useful reasoning patterns through self-exploration and reward guidance.Source: §1 · p. 1
  • Figure 1 depicts policy-gradient integration: a rule-based verifier rewards on-policy rollouts while a teacher supplies token-level KL supervision.The study examines KL approximations, KL scheduling, and reward-guided distillation.Source: Figure 1 · p. 2, §1 · p. 2
  • KDRL jointly optimizes teacher supervision and self-exploration, addressing the decoupling of KD and RL in separate post-training stages.It minimizes reverse KL between the policy and teacher using on-policy rollouts, enabling integration with RL algorithms such as GRPO.Source: §1 · pp. 1–2
  • 5.7% and 2.6%: KDRL outperforms standard SFT and GRPO, respectively, and exceeds KD-RKL by 1.1%.The paper also reports a favorable trade-off between performance and training cost and reasoning-token usage.Source: §1 · p. 2

2 Preliminary

The preliminary section introduces GRPO as the RL approach used in the study, explains teacher imitation and on-policy reverse-KL distillation, and describes KL-estimation options. It also presents the training metrics used to compare methods.

2.1 Reinforcement Learning

  • GRPO samples response groups per question and uses outcome rewards, advantage-based variance reduction, and constrained policy updates; this study removes the KL constraint and uses token-level loss.GRPO is selected for computational efficiency.Source: §2.1 · pp. 2–3

2.2 Knowledge Distillation

  • Teacher-generated SFT minimizes forward KL, whereas training on student-generated prefixes minimizes reverse KL and can alleviate exposure bias, at higher computational cost.In the reverse-KL formulation, the teacher’s relative preference for a trajectory acts as a reward signal.Source: §2.2 · p. 3

2.3 KL Approximation

  • The k2 approximation provides an unbiased estimate of the reverse-KL gradient, while the direct estimator is described as high variance.The section contrasts gradient estimation with estimation of the scalar divergence.Source: §2.3 · p. 3
  • Top-K approximation retains and renormalizes the teacher’s top-K logits for full-vocabulary KL computation, avoiding the memory overhead of storing full distributions.Figure 2 tracks training rewards, response length, response clip ratio, and AIME24 accuracy across methods.Source: §2.3 · p. 3, Figure 2 · p. 4

3 KDRL: Unifying KD and RL for Post-training Reasoning LLMs

KDRL combines GRPO with reverse-KL distillation, then examines how KL estimation, the KD/RL balance, and reward-guided masking affect training and reasoning efficiency. Across these studies, joint loss and response-level masking are supported strategies, while results vary with the KL estimator and coefficient.

3 KDRL: Unifying KD and RL for Post-training Reasoning LLMs

  • KDRL jointly trains with teacher supervision and environment feedback, using reverse KL as its distillation objective.The study investigates how to integrate the objectives, approximate reverse KL, set its coefficient, and mask samples for its computation.Source: §3 · p. 4

3.1 Combination of GRPO and KD-RKL

  • Joint loss makes training stable and outperforms both GRPO and KD-RKL in reward and validation accuracy, whereas reward shaping collapses early.Joint-loss KDRL also produces more stable response-length growth than GRPO or KD-RKL, whose responses respectively grow gradually and inflate rapidly with truncation approaching 15%.Source: §3.1 · p. 4

3.2 Approximations of KL Divergence

  • The k2 estimator improves AIME24 accuracy by 1.4% over k3, while Top-K destabilizes training and produces no accuracy gain.The Top-K runs also show a sharp rise in repetition within overlong responses; k2 and k3 yield similar training rewards.Source: §3.2 · pp. 4–5, Figure 3 · p. 5

3.3 Balance of RL and KD

  • The linearly annealed KL coefficient achieves the best overall performance, matching the response length of β = 2e-3 while maintaining a lower clip ratio.A larger constant β = 5e-3 is associated with rapid length inflation and lower reward acquisition, while a smaller β = 1e-3 yields more stable length growth but underuses teacher supervision.Source: §3.3 · pp. 5–6, Figure 4 · p. 5

3.4 Reward-Guided KD

  • Response-level reward-guided masking maintains comparable performance with shorter responses, whereas group-level masking yields limited improvement.Response-level masking omits KD loss for positively rewarded responses; group-level masking suppresses it for an entire sampled group if any response receives a reward, applying KD only when all responses fail.Source: §3.4 · p. 6, Figure 5 · p. 6

4 Experiments

Across mathematical reasoning benchmarks, KDRL combines teacher supervision and reward-driven learning, outperforming standalone KD and RL methods. Its efficiency analyses, ablations, and R1-Zero-like experiments examine the balance between accuracy, training cost, and reasoning length.

4.1 Experimental Setup

  • Stage 1 trains R1-Distill-Qwen-1.5B on DeepScaleR to create the DeepScaleR-1.5B-8K backbone; stage 2 applies KDRL with Skywork-OR1-Math-7B as teacher.Stage 2 uses 16 rollouts per prompt, a 20K context length, and difficulty-controlled training data.Source: §4.1 · p. 7

4.2 Main Results

  • KDRL-Annealing improves overall accuracy over SFT, GRPO, and KD-RKL by 4.7%, 2.6%, and 1.1%, respectively.SFT falls 3.1% below GRPO, while KD-RKL exceeds GRPO by 1.5%.Source: §4.2 · pp. 7–8

4.3 Efficiency Analysis

  • KDRL and KDRL-Annealing offer the best AIME24 accuracy–reasoning-token balance; KD-RKL uses over 6K additional tokens compared with KDRL.At roughly 80 hours, KDRL consistently outperforms GRPO and KD-RKL under equal training budgets.Source: §4.3 · p. 8, Figure 6 · p. 8

4.4 Ablation Study

  • KDRL-Masking reduces output length by over 1.4K tokens, or more than 10%, while maintaining comparable accuracy.The k3 estimator produces shorter outputs but slightly lower accuracy than k2; using the weaker R1-Distill-Qwen-7B teacher yields only a 1.1% gain over GRPO.Source: §4.4 · pp. 8–9

4.5 KDRL for R1-Zero-Like Training

  • In R1-Zero-like training, KDRL-Annealing achieves the best performance on all datasets, surpassing GRPO and KD-RKL by 1.2% and 1.7%, respectively.GRPO slightly exceeds KD-RKL, with 31.3% versus 30.8%.Source: §4.5 · p. 9, Table 2 · p. 9

5 Related Work

Related work spans RL for reasoning, off- and on-policy distillation, and approaches combining RL with KD. KDRL is positioned as an on-policy, rule-based RL framework that incorporates teacher supervision during reinforcement fine-tuning.

Reinforcement Learning for Reasoning LLMs

  • Prior reasoning RL spans direct training of pretrained models, enhancement of distilled models, and backbone-agnostic strategy improvements; KDRL falls into the third category.The paper reports experiments in both Zero-RL and RL-from-distilled settings.Source: §5 · p. 10

Knowledge Distillation

  • Standard KD uses teacher-generated outputs, whereas on-policy KD learns from student samples and teacher logits; KDRL combines on-policy KD-RKL with RL.The paper notes prior reasoning distillation largely focused on off-policy SFT-style approaches, while Qwen3 explores on-policy distillation.Source: §5 · p. 10

Combination of RL and KD

  • Unlike DPO-based guidance, GKD, or LUFFY, KDRL studies on-policy distillation with rule-based RL for reasoning and adds KL annealing and reward-guided masking.The authors describe this as an early step toward unified post-training for long-CoT models.Source: §5 · p. 10

6 Conclusion

The paper presents KDRL as an on-policy framework that combines teacher guidance and reinforcement learning through policy gradients. Across reasoning benchmarks, it reports gains over GRPO and KD baselines alongside improved training and reasoning-token efficiency.

6 Conclusion

  • KDRL integrates knowledge distillation and reinforcement learning through policy gradients under on-policy optimization.The paper analyzes KL approximations, KL coefficients, and reward-guided distillation.Source: §6 · p. 10
  • Across several reasoning benchmarks, KDRL outperforms GRPO and various KD baselines while achieving greater training and reasoning-token efficiency.Source: §6 · p. 10

A Mathematical Derivations

The appendix derives the RL and distillation objectives, then describes reverse-KL estimators and the combined KDRL gradient. It also introduces a Top-K approximation for using teacher signals over a reduced vocabulary.

A.1 Reinforcement Learning

  • The appendix derives gradients for REINFORCE and GRPO as its RL starting point.Source: §A.1 · p. 15

A.2 Knowledge Distillation

  • Forward KL minimization is equivalent to maximizing student log-likelihood on teacher-generated data, with sequence-length normalization used in practice.Source: §A.2 · p. 15

A.3 KL Approximation

  • Reverse-KL distillation samples from the current policy, and the appendix presents direct, k2, k3, and Top-K approximations.The k2 estimator has an unbiased RKL gradient; the k3 estimator has an unbiased scalar divergence estimate but a biased gradient. Top-K retains and renormalizes the teacher's top-K logits.Source: §A.2 · p. 16, §A.3 · p. 17

A.4 KDRL

  • The KDRL objective’s gradient combines the GRPO and reverse-KL gradients.Source: §A.4 · p. 17

B.1.1 Traning Dataset

The two-stage pipeline uses a reused open-source math dataset for Stage 1 and a difficulty-filtered Skywork-OR1 dataset for KDRL training.

B.1.1 Traning Dataset

  • Stage 1 reuses DeepScaleR’s open-source dataset of 40K math question–answer pairs from sources including AIME, AMC, and Omni-Math.Source: §B.1.1 · p. 18
  • After difficulty and prompt-length filtering, Stage 2 uses 61K Skywork-OR1 math problems, retaining 10% unsolved examples for teacher-supervision signals.Filtering removes problems with pass@16 of at least 15 and problems unanswered by all three R1-Distill-Qwen scales.Source: §B.1.1 · p. 18

B.1.2 Implementation Details

KDRL uses a binary reward requiring both structured output and a verified correct answer, while the baselines vary in their use of teacher supervision.

B.1.2 Implementation Details

  • KDRL’s binary reward is 1 only when the response follows the <think>...</think> format and its extracted answer passes math-verify checking.If direct matching fails, the gold answer is wrapped in \boxed{} and checked again.Source: §B.1.2 · p. 18
  • Training uses VeRL and a shared step-by-step prompt template that requests the final answer within \boxed{} and applies the R1-Distill-Qwen chat template.Source: §B.1.2 · p. 18

B.2 Baseline Descriptions

  • SFT uses eight sampled teacher responses and reject sampling for correct answers, while KD-RKL minimizes student–teacher RKL with a fixed 2e-3 coefficient.SFT uses a 20K-token context window, global batch size 128, and learning rate 2e-7.Source: §B.2 · p. 18
  • GRPO is the pure-RL baseline without teacher supervision; its hyperparameters match KDRL except for KL computation.The passage reports faster training but no imitation signals.Source: §B.2 · p. 18

B.3 Details of R1-Zero-like Training

R1-Zero-like training uses separate difficulty filters for teacher and student data, and specifies accuracy-reward and KL settings for its training variants.

B.3 Details of R1-Zero-like Training

  • Difficulty filtering yields 58K teacher-training and 46K student-training examples from Skywork-OR1, using pass-rate cutoffs of 25% and 50%, respectively.Source: §B.3 · p. 19
  • Reinforcement learning uses only the rule-based accuracy reward; KD uses a 5e-3 KL coefficient, while KDRL-Annealing sets βinit=2e-1, δ=8e-4, and βmin=2e-3.Qwen2.5-3B models train for 480 steps with an 8K context length and sampling temperature 1.0.Source: §B.3 · p. 19
  • The teacher model is produced by applying GRPO to Qwen2.5-7B for 300 training steps.Source: Table 4 · p. 19
  • Table 5 summarizes training hyperparameters and notes that GRPO and KD-RKL share KDRL settings except for KL-related differences.Source: Table 5 · p. 19
Loading 2506.02208v1…