The Weakest Link: Distilling LLM Reasoning with Worst-Case Constrained Reinforcement Learning
cs.LG, cs.AI
Submitted: 2026-09-29
Updated: 2026-09-29
License: http://creativecommons.org/licenses/by/4.0/
The gist: Distilling the reasoning capabilities of large language models (LLMs) into smaller students is a central challenge for efficient deployment.
Terminology
Abstract
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.
Sources
- SIKeD: Self-guided Iterative Knowledge Distillation for mathematical reasoning
- Step-wise Adaptive Integration of Supervised Fine-tuning and Reinforcement Learning for Task-Specific LLMs
- Gold-medalist Performance in Solving Olympiad Geometry with AlphaGeometry2
- MiniLLM: On-Policy Distillation of Large Language Models
- Distilling the Knowledge in a Neural Network
- Distilling Step-by-Step! Outperforming Larger Language Models with Less Training Data and Smaller Model Sizes
- Tree-OPO: Off-policy Monte Carlo Tree-Guided Advantage Optimization for Multistep Reasoning
- On Almost Surely Safe Alignment of Large Language Models at Inference-Time
- DistiLLM: Towards Streamlined Distillation for Large Language Models
- DistiLLM-2: A Contrastive Approach Boosts the Distillation of LLMs
- Training language models to follow instructions with human feedback
- Learning to summarize from human feedback
- On a few pitfalls in KL divergence gradient estimation for RL
- Distilling Morphology-Conditioned Hypernetworks for Efficient Universal Morphology Control
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks