Rethinking Large Language Model Distillation: A Constrained Markov Decision Process Perspective

1 INTRODUCTION

• KL-focused distillation can force smaller students to imitate complex teacher reasoning, while reward-only optimization cannot guarantee correct reasoning steps.The introduction motivates balancing teacher fidelity with task-specific reward rather than relying on either alone.Source: §1 · p. 1
• The method formulates distillation as constrained RL, maximizing task reward while keeping student–teacher divergence below a predefined threshold.The threshold is specified directly in KL scale; once the constraint is satisfied, the objective reduces to reward maximization.Source: §1 · pp. 2–3
• Figure 1 illustrates that reward maximization alone can yield flawed reasoning even when a response reaches the correct answer.The GRPO example makes an erroneous intermediate result before an extra rounding step reaches 5.Source: Figure 1 · p. 2, §1 · p. 3
• Experiments show better reasoning quality and constraint satisfaction with competitive final-answer correctness.The comparison covers Qwen2.5-1.5B-Math and Llama-3.2-3B across three evaluation datasets.Source: Figure 2 · p. 2, §1 · p. 3
• The method adapts Saute by reformulating the reward without state augmentation, retaining theoretical guarantees while avoiding teacher access at test time.The paper presents this as a more efficient and practical solution for distillation.Source: §1 · pp. 2–3

2.1 DISTILLATION

• Standard distillation minimizes student–teacher divergence, but on complex reasoning tasks it can miss better-performing solution paths suited to a smaller student's capacity.Reward-based distillation instead combines expected task reward with a divergence regularizer, though choosing its weighting parameter can require extensive task-specific retraining.Source: §2.1 · p. 3

2.2 CONSTRAINED REINFORCEMENT LEARNING

• Constrained RL frames distillation as maximizing task-specific reward while keeping expected student–teacher divergence below a predefined budget.The formulation models prompts with partial responses as states and generated next tokens as actions.Source: §2.2 · pp. 3–4
• Saute tracks remaining budget in an augmented state and penalizes trajectories that exhaust it, yielding feasible policies as the penalty grows under standard assumptions.This avoids the computational overhead of Lagrange multipliers and dual ascent.Source: §2.2 · p. 4
• Direct state augmentation would require computing divergence with the teacher at every step, undermining the goal of a student that operates independently.The paper identifies this deployment requirement as impractical for distillation.Source: §2.2 · p. 4

3.1 CONSTRAINED RL FOR LLM DISTILLATION

• The unaugmented formulation retains hard-constraint semantics by reconstructing the remaining budget from the full interaction history and penalizing infeasible trajectories more when they diverge further from the teacher.The discrepancy term may use an f-divergence such as KL or Jensen–Shannon divergence; feasible trajectories retain task-specific reward.Source: §3.1 · pp. 4–5

3.2 POLICY GRADIENT OPTIMIZATION

• Policy-gradient optimization decomposes the update into a likelihood-ratio term and an explicit reward-dependence term, under differentiability and integrability assumptions.In practice, the discrepancy uses KL with a small probability floor to ensure finiteness and differentiability.Source: §3.2 · p. 5

3.3 THEORETICAL GUARANTEE FOR CONSTRAINT SATISFACTION

• Under standard assumptions, the unaugmented and augmented formulations have equivalent optimal values, and finite-value optimal policies satisfy the constraint almost surely.The student can operate at test time without teacher access, computing cumulative reward from its own output distribution and environment feedback.Source: §3.3 · pp. 5–6

4.1 EXPERIMENTAL SETUP

• Two settings distill Qwen2.5-1.5B-Math from Qwen2.5-7B-Math-Instruct on GSM8K training data and Llama-3.2-3B from Llama-3.2-11B-Instruct on MATH training data.Checkpoints are evaluated after 20 epochs on GSM-Symbolic (main), GSM8K test, and the whole MATH test set.Source: §4.1 · p. 6
• Baselines are reimplemented under GRPO with the same batch composition, learning rate, and AdamW optimizer for fair comparison.The baseline set covers task-reward optimization, KL-only distillation, and hybrid reward–KL objectives; the hybrid multiplier is grid-searched over five values.Source: §4.1 · pp. 6–7
• Evaluation measures final-answer correctness, judged reasoning quality, constraint satisfaction, and average student–teacher KL divergence.The reasoning judge receives the correct final answer to focus its pairwise assessment on the reasoning path.Source: §4.1 · p. 7

4.2 EXPERIMENT RESULTS

• The results section is organized around the best general method, constraint satisfaction, external reward for distillation, and whether distillation improves reasoning.Source: §4.2 · p. 7

• The constrained RL approach excels in reasoning quality and constraint satisfaction while maintaining competitive final-answer correctness.Source: A. What is the best method in general? · p. 7
• Pure reward optimization (GRPO λ=0.0) has the highest final-answer correctness but poor reasoning quality and severe constraint violations.Source: A. What is the best method in general? · p. 7
• KL-only methods GKD and Mini-LLM maintain good constraint satisfaction but have lower final-answer correctness.Source: A. What is the best method in general? · p. 7

• Without ϕ, the method would struggle to satisfy a strict constraint because all trajectories receive the same penalty and training would diverge.Source: B. Is our method able to achieve… · p. 7
• The method achieves superior constraint satisfaction while maintaining competitive final-answer correctness across the compared configurations.Source: B. Is our method able to achieve… · p. 7, Figure 3 · p. 8
• The Pareto analysis compares final-answer correctness against constraint satisfaction across methods and hyperparameter settings for Qwen2.5-1.5B-Math.Each plotted point is a method configuration, and red points mark the Pareto front.Source: Figure 3 · p. 8

C. Can external reward help achieve better distillation?

• Pure KL minimization methods achieve lower final answer correctness than reward-based methods on every dataset for each model.The comparison includes GRPO, GKD-GRPO variants, and the proposed method versus GKD and Mini-LLM.Source: C. Can external reward help achieve… · p. 8
• Reward-based methods also achieve higher reasoning win rates than purely KL-based methods.The passage attributes this improvement to the use of the reward function.Source: C. Can external reward help achieve… · p. 8
• Task-specific rewards help students learn more effective reasoning strategies rather than merely mimic the teacher’s surface-level outputs.Source: C. Can external reward help achieve… · p. 8

D. Does the distillation signal help to better reason?

• In the test example, both methods give the correct final answer, but only the proposed method produces logically valid reasoning steps.More examples are provided in Appendix F.Source: D. Does the distillation signal help… · p. 8
• 12–19% reasoning win rates and 39–55% loss rates accompany GRPO’s 75–80% raw final answer correctness, while the constrained method improves reasoning and maintains competitive success rates.Figure 4 averages pairwise comparisons across Apple/GSM-Symbolic, GSM8K, and MATH with Qwen; darker colors indicate superior row-to-column performance.Source: Figure 4 · p. 8, D. Does the distillation signal help… · p. 8
• Across the reported evaluation dimensions, the constrained approach maintains high constraint satisfaction, leverages external rewards for task performance, and preserves teacher reasoning capabilities in the student.The paper says existing approaches typically excel in only one aspect of the distillation objective.Source: D. Does the distillation signal help… · pp. 8–9

5 RELATED WORKS

• Conventional distillation aligns teacher and student output distributions, often using reverse KL, but does not explicitly guarantee preservation of reasoning abilities on complex tasks (Hinton et al., 2015; Gu et al., 2024; Zimmer et al., 2025).RL-based distillation treats KL divergence as a dense token-level reward; general-purpose methods use REINFORCE-style updates, while Mini-LLM separates long-term reward from a single-step objective (Agarwal et al., 2024; Ko et al., 2024; 2025; Gu et al., 2024).Source: §5 · p. 9
• Task-aware approaches add process supervision, task-relevant logit weighting, or retrieval, but can require fine-grained annotations, complex weighting heuristics, or external knowledge bases (Hsieh et al., 2023; Adarsh et al., 2024; Chen et al., 2025; Li et al., 2025; 2024; Saadi & Wang, 2025; Kang et al., 2023; Tian et al., 2025).Source: §5 · p. 9
• Task-specific constrained RL distillation remains relatively under-explored, and fixed KL penalties can be brittle when students exploit rewards through shallow or degenerate reasoning.In distillation, KL is intended to constrain divergence to preserve the teacher’s reasoning capabilities, unlike its regularization role in standard RLHF settings.Source: §5 · p. 9
• The paper treats KL divergence as an explicit trust-region constraint and reformulates the dual Lagrange problem within a state-augmented MDP to provide an efficient solution practical at LLM scale (Schulman et al., 2015; Achiam et al., 2017; Calvo-Fullana et al., 2024; Sootla et al., 2022a;b).The passage identifies compute, memory, and variance costs from teacher forward passes, cached logits, and inner-loop or dual updates as practical challenges (Achiam et al., 2017).Source: §5 · p. 9

6 CONCLUSION

• The framework maintains constraint-satisfaction guarantees without classic state augmentation and navigates the trade-off between task-specific performance and teacher fidelity.The conclusion also describes avoiding brittle ad-hoc reward weighting and the costs of traditional dual max-min optimization.Source: §6 · pp. 9–10
• Experiments on mathematical reasoning show that strict KL divergence constraints can be enforced with high fidelity while maintaining competitive task rewards.Source: §6 · p. 10

A DERIVATION OF POLICY GRADIENT

• Assumption A1 requires the discrepancy function and its gradient to be well-behaved; the implementation uses KL divergence and overlapping support, for example via a probability floor.These conditions keep the discrepancy finite and its gradient bounded across states.Source: A DERIVATION OF POLICY GRADIENT · p. 15
• The explicit-dependence gradient term vanishes on strictly feasible trajectories but reduces to the negative discrepancy gradient after budget violation.Source: A DERIVATION OF POLICY GRADIENT · p. 15
• At the constraint boundary, the reward is nondifferentiable, so the derivation uses a generalized Mordukhovich subgradient and takes a limiting subgradient from the infeasible side with tolerance ε ↓0.The passage notes that hitting the boundary exactly is unlikely in the continuous constraint setting.Source: A DERIVATION OF POLICY GRADIENT · p. 15
• Near the optimum, the explicit-dependence term disappears on feasible trajectories while continuing to signal trajectories that violate or approach the constraint boundary.The passage also reports that this term contributes to variance reduction during training.Source: A DERIVATION OF POLICY GRADIENT · p. 15

B PROOFS OF CONSTRAINT SATISFACTION GUARANTEE

• Under B1–B3, the unaugmented MDP satisfies the Bellman equation, and an optimal policy with finite value is optimal for the constrained MDP and satisfies its constraint almost surely.The optimal value functions also converge monotonically.Source: B PROOFS OF CONSTRAINT SATISFACTION… · pp. 16–17

C A PERSPECTIVE OF LLM DISTILLATION AS CONTEXTUAL MDPS

• With context fixed within an episode, the contextual and standard MDP formulations are equivalent in return and have coinciding optimal policies on the reachable set.The paper motivates the standard MDP by freezing the student within each episode and fixing the teacher during distillation.Source: C A PERSPECTIVE OF LLM DISTILLATION AS… · p. 17

D.1 REWARD FUNCTION DESIGN

• For mathematical reasoning, the reward is binary and assigned only after the complete solution: 1.0 for a correct final answer and 0.0 otherwise.Intermediate steps are not evaluated in this sparse-reward setup.Source: §D.1 · p. 18

D.2 KL DIVERGENCE COMPUTATION

• The student–teacher KL divergence is computed at each time step over the LLM vocabulary.The supplied passage identifies the vocabulary as V.Source: §D.2 · p. 18

D.3 HYPERPARAMETER SETTINGS

• Training used a constraint threshold d = 0.35 and 20 epochs; Llama3.2-3B received KL-only distillation for its first three epochs to bootstrap unstable GRPO training.The threshold was selected using preliminary KL-minimization experiments.Source: §D.3 · p. 18

D.4 TRAINING TIME

• Each method trained in less than two days on a single accelerator, with similar overall training time across methods.GRPO was only a little faster because backward passes and generation dominated total time.Source: §D.4 · p. 18

E MORE EXPERIMENTS RESULTS

• Supplementary results cover Qwen2.5-1B on GSM8K and Llama3.2-3B on MATH; Figure 5 summarizes pairwise reasoning comparisons across three datasets.The table captions identify FAC, RWR, constraint satisfaction, KL, and RLR for GSM8K, and success rate, constraint satisfaction, and KL for MATH; they do not provide cell values.Source: Table 1 · p. 19, Table 2 · p. 19, Figure 5 · p. 19

F GENERATED ANSWERS

• Figures 6–8 provide generated-answer examples from Qwen2.5-1.5B-Math after distillation.The supplied captions do not describe the answers’ contents.Source: Figure 6 · p. 20, Figure 7 · p. 21, Figure 8 · p. 22