Distilling the reasoning capabilities of large language models (LLMs) into smaller students is a central challenge for efficient deployment. Current approaches face a fundamental tension: optimizing purely for verifiable task rewards (e.g., via GRPO) leads to reward hacking, where students arrive at correct final answers through flawed intermediate logic, while regularizing with soft divergence penalties against a teacher (e.g., KL-based distillation) dilutes task performance and, critically, allows the student to compensate for severe logical violations at one step with high teacher agreement at others. We argue that this averaging is fundamentally misaligned with the nature of reasoning: a chain-of-thought is only as valid as its weakest link. Motivated by this observation, we formulate reasoning distillation as a constrained reinforcement learning problem in which the task reward is maximized subject to a worst-case constraint on the teacher log-likelihood along every prefix of the trajectory. To avoid the prohibitive cost of dual Lagrangian solvers and the test-time teacher dependence of state-augmented methods such as Saute, we derive an unaugmented constrained MDP whose reward transformation preserves the hard-constraint semantics, admits a low-variance policy gradient decomposition into single-step and long-term terms, and provably satisfies the worst-case constraint almost surely in the penalty limit. Through extensive experiments on mathematical reasoning and code generation tasks, we demonstrate that our method significantly expands the accuracy-fidelity Pareto front. By matching the high Final Answer Correctness of pure RL and drastically reducing teacher constraint violations, we ultimately achieve the highest rigorous Reasoning Success Rate across all evaluated settings.
Figures & tables
Figure 1: GRPO (right) reaches a wrong answer (4.7619) but applies a learned rounding trick to recover the correct answer (5).
Figure 2: Reasoning Improvement across budget bounds. Constraining the worst-case average prevents local hallucination effectively, yielding strictly higher peak reasoning performance across optimal budget thresholds compared to a cumulative sum budget.
Figure 3: Accuracy-Fidelity Pareto Front on the 3 distillation settings. X-axis: Constraint Satisfaction (higher is better). Y-axis: Final Answer Correctness (higher is better). Soft-regularized baselines (GKD-GRPO line) force a harsh trade-off between fidelity and accuracy. Our worst-case constrained method matches or outperforms pure RL accuracy while retaining high fidelity.
Qwen2.5 (MATH)
Llama-3.2 (MATH)
Qwen2.5 (GSM-Sym.)
Method
RSR
FAC
CS
RSR
FAC
CS
RSR
FAC
CS
SFT on Teacher CoT
24.66
27.02
78.54
13.64
17.48
57.72
55.78
65.9
83.46
GKD
24.00
25.84
85.18
9.34
15.24
65.54
57.58
66.06
86.22
MiniLLM
24.22
26.28
85.68
11.08
15.20
61.88
55.28
63.96
83.88
GKD-GRPO λ=10.0
24.58
26.32
84.38
7.44
16.46
66.88
56.56
65.70
83.70
GKD-GRPO λ=1.0
24.32
26.40
85.14
9.60
17.16
67.76
57.18
65.78
83.98
Table 1: Main distillation results across three settings. All metrics are reported in percentages (%) and higher is better. Our method ( d=−ln(1/4) ) achieves the highest Reasoning Success Rate (RSR) across all settings by breaking the trade-off between Final Answer Correctness (FAC) and Constraint Satisfaction (CS). Best results are in bold , second best are underlined .
Figure 4: Average win rate (%) evaluated by an LLM-as-a-Judge across the three distillation settings.
Figure 5: Distribution of the worst-case prefix average cost at test time. Soft baselines display extended ”tails” of severe localized violations, empirically validating the compensation problem.
Appendix figures & tables5 assets
Supplementary material from the paper’s appendix.
Appendix
Method
RSR
FAC (Pass@1)
FAC − RSR
GRPO
1.8
92.0
90.2
GKD-GRPO λ=0.001
14.4
74.4
60.0
GKD-GRPO λ=0.01
34.0
50.4
16.4
Ours d=−ln(1/4)
34.2
45.6
11.4
Appendix
Table 2: Code generation on the MBPP test set (500 problems, %). The prompt contains the unit tests, making pass@1 easy to game: GRPO inflates pass@1 to 92.0% by copying them into the generated code (RSR 1.8%). Ours achieves the highest RSR and smallest FAC − RSR gap.
Table 3: Pooled Spearman rank correlation between per-trajectory fidelity statistics and RSR (divergence costs are negated so that higher always means higher teacher fidelity).
Figure 6: Example of generated answer with Qwen2.5-1.5B after distillation.
Figure 7: Example of generated answer with Qwen2.5-1.5B after distillation.
Figure 8: Example of generated answer with Qwen2.5-1.5B after distillation.
Language model distillation transfers a larger model’s capabilities to smaller or specialized models using outputs, probabilities, or internal representations.
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.
When distilling reasoning from large language models (LLMs) into smaller ones, teacher rationales for similar problems often vary wildly in structure and strategy. Like a chef who makes the same dish differently each time, this inconsistency burdens the student with noisy supervision that is hard to internalize. We propose Distillation through Reasoning Path Compression (D-RPC), which constrains the teacher to follow a compact, dynamically maintained bank of reusable high-level reasoning paths. For each training question, D-RPC retrieves the most relevant path and conditions the teacher to follow it, producing rationales that are consistent across similar problems yet diverse enough to cover different problem types. A PAC-Bayes analysis formalizes the resulting trade-off between bank size and coverage: smaller banks reduce supervision entropy but risk coverage gaps, and the generalization bound identifies an optimal intermediate size confirmed by our ablations. Across five math and commonsense reasoning benchmarks with two student models, D-RPC consistently outperforms chain-of-thought distillation, freeform rationale generation, direct distillation, and structured-supervision baselines, while using fewer tokens than template-heavy alternatives.
Jialin Yang, Jiankun Wang, Jiajun Wu +3
University of Calgary · Calgary, Canada · University of Michigan +1
Reasoning distillation aims to transfer multi-step reasoning capabilities from large language models to smaller, more efficient ones. While recent methods have shown promising gains, they typically rely on static teacher-student hierarchies and frame distillation as trajectory imitation. This is misaligned with the structure of reasoning, where intermediate steps are often locally under-specified: global correctness constrains the final answer, but does not uniquely determine each intermediate move. We propose validity-calibrated reasoning distillation, a framework that treats reasoning distillation as a problem of local learning-signal allocation rather than path alignment. Instead of enforcing token-level imitation, we compare the student's and teacher's proposed next-step actions under the same prefix and use their relative local validity to modulate the strength of the distillation update. This yields a dynamic, context-dependent supervision mechanism that preserves the teacher's structural guidance while adapting update strength to local reasoning quality. Across mathematical reasoning, code generation, and instruction-following benchmarks, our method consistently outperforms strong distillation baselines. These results indicate that effective LLM reasoning distillation is governed not by rigid trajectory imitation, but by principled, locally calibrated allocation of learning signal.
Khouloud Saadi, Di Wang
Department of Computer Science, KAUST, Thuwal, Saudi Arabia