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

Abstract: We introduce a novel approach to LLM distillation by formulating it as a constrained reinforcement learning problem. While recent work has begun exploring the integration of task-specific rewards into distillation processes, existing methods typically rely on ad-hoc reward weighting. We propose a principled optimization framework that maximizes task-specific rewards while constraining the divergence from the teacher model to remain below a specified threshold. Our approach adapts constrained state augmented reinforcement learning to the distillation setting, introducing a modified reward function that maintains theoretical guarantees of constraint satisfaction without requiring state augmentation or teacher model access during deployment and without the computational overhead of the dual Lagrangian methods. Through extensive experiments on mathematical reasoning tasks, we demonstrate that our method achieves better constraint satisfaction rates and better reasoning compared to the soft Lagrangian relaxation baselines while maintaining competitive task performance. Our framework provides a theoretically grounded and practically efficient solution for reward-aware distillation in resource-constrained settings.

• The paper introduces a constrained RL formulation that maximizes task reward while strictly bounding divergence from the teacher.
• It removes the need for explicit state augmentation and dual Lagrangian methods, thereby improving computational efficiency and stability.
• Experimental results show enhanced reasoning quality and superior constraint satisfaction on mathematical reasoning tasks.

Rethinking LLM Distillation: A Constrained Markov Decision Process Perspective

Introduction

This paper presents a principled framework for LLM distillation, recasting the problem as a constrained reinforcement learning (RL) task. Traditional distillation approaches focus on minimizing divergence (typically KL) between student and teacher models, often neglecting task-specific reward signals that are crucial for complex reasoning tasks. Recent RL-based distillation methods introduce reward terms but rely on ad-hoc hyperparameter balancing, which is unstable and computationally expensive. The authors propose a constrained optimization approach that maximizes task reward while strictly bounding the divergence from the teacher, leveraging a modified state-augmented RL formulation that eliminates the need for teacher access at deployment and avoids the computational overhead of dual Lagrangian methods.

Constrained RL Formulation for Distillation

The core contribution is the formulation of LLM distillation as a constrained MDP, where the objective is to maximize expected task reward subject to a hard constraint on the cumulative divergence from the teacher policy. The authors adapt the Saute state-augmentation method, but crucially remove the need for explicit state augmentation by exploiting the history-conditioned nature of LLM policies. The constraint is enforced via a modified reward function that penalizes trajectories violating the divergence budget, with the penalty scaled by the degree of violation (using an $f$-divergence such as KL).

This approach yields several advantages:

• Direct control over fidelity: The constraint is specified in terms of KL scale, eliminating the need for delicate hyperparameter tuning.
• Efficient optimization: The method avoids the computational cost of dual ascent and teacher model access at test time.
• Theoretical guarantees: The authors prove optimality equivalence, Bellman optimality, and almost sure constraint satisfaction for the proposed formulation.

Policy Gradient Optimization

The paper provides a detailed derivation of the policy gradient for the unaugmented constrained objective, decomposing the gradient into a likelihood-ratio term and an explicit dependence term arising from the penalty function. Under mild regularity assumptions, the gradient estimator is shown to be well-behaved and variance-reducing, with the explicit dependence term vanishing on feasible trajectories and providing informative feedback on constraint-violating ones.

Experimental Evaluation

Extensive experiments are conducted on mathematical reasoning tasks, distilling Qwen2.5-1.5B-Math and Llama3.2-3B students from larger teacher models. The evaluation covers final answer correctness (FAC), reasoning win/loss rates (RWR/RLR), constraint satisfaction, and KL divergence. Baselines include pure reward optimization (GRPO), pure KL minimization (GKD, Mini-LLM), and hybrid Lagrangian relaxation (GKD-GRPO).

The results demonstrate:

• Superior constraint satisfaction: The proposed method consistently achieves higher rates of constraint satisfaction compared to baselines, occupying a unique region of the Pareto front balancing FAC and fidelity.

  Figure 1: Pareto frontier analysis showing the trade-off between final answer correctness and constraint satisfaction across different methods and hyperparameter settings for Qwen2.5-1.5B-Math.

• Improved reasoning quality: While pure reward optimization yields high FAC, it suffers from poor reasoning quality and frequent constraint violations. The constrained RL approach achieves a balanced profile, with strong reasoning win rates and competitive FAC.

  Figure 2: Pairwise comparison heatmap showing the relative performance of the proposed method against baselines, averaged across all evaluation datasets on Llama3.2-3B.

• Necessity of teacher signal: Reward maximization alone cannot guarantee correct reasoning steps; the teacher signal is essential for transferring reasoning capabilities.

  Figure 3: Example illustrating that checking the final answer alone is insufficient for evaluating reasoning. GRPO (right) makes mistakes and reaches a wrong answer (4.76) but takes an extra rounding step to the correct one (5), likely as a learned strategy through training.

• Sample efficiency and stability: The penalty refinement using the policy-dependent discrepancy term $\phi$ is critical for training stability and constraint satisfaction.

Implementation Details

The method is implemented using standard policy gradient optimization with the following settings:

• Batch size: 64 responses (8 questions × 8 responses per question)
• Learning rate: $1 \times 10^{-5}$, AdamW optimizer
• Discount factor: $\gamma=1$
• Constraint threshold: $d=0.35$
• Penalty parameter: $n=20$
• Training epochs: 20

KL divergence is computed per token, and the reward is assigned based on final answer correctness. The method is computationally efficient, requiring less than two days of training per method on a single accelerator.

Theoretical and Practical Implications

The constrained RL perspective provides a robust and interpretable framework for LLM distillation, enabling direct control over the fidelity-performance trade-off. The elimination of state augmentation and dual optimization makes the approach scalable and practical for large models. Theoretical guarantees ensure that the student model operates reliably within a defined trust region of the teacher, which is critical for deployment in safety-critical or resource-constrained environments.

Future Directions

Potential avenues for future research include:

• Extending the framework to multi-teacher or multi-task distillation settings.
• Incorporating richer, process-aware or logit-aware supervision signals.
• Exploring adaptive constraint thresholds and dynamic penalty scaling.
• Applying the method to other domains requiring strict fidelity guarantees, such as medical or legal reasoning.

Conclusion

This work advances the state of LLM distillation by introducing a constrained RL framework that is both theoretically sound and practically efficient. The approach enforces strict fidelity constraints while leveraging task-specific rewards, resulting in student models that are compact, reliable, and capable of high-quality reasoning. The framework is broadly applicable and sets a foundation for future developments in controllable and deployable LLMs.