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

arXiv: https://arxiv.org/abs/2509.22921 · HTML: https://mlsanity.com/paper/2509.22921/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Placed in Reinforcement learning for reasoning, one of 17 directions in LLM reasoning.
- Builds on 18 of its 52 references that are in the map.
- First posted on arXiv 26 Sep 2025.

## Placed in

- LLM reasoning / [Reinforcement learning for reasoning](https://mlsanity.com/lr/reinforcement-learning-for-reasoning/)

## Builds on (most cited first, 18 of 18)

- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning](https://mlsanity.com/paper/2501.12948/) (2025)
- [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter](https://mlsanity.com/paper/1910.01108/) (2019)
- [Learning to summarize from human feedback](https://mlsanity.com/paper/2009.01325/) (2020)
- [Distilling Step-by-Step! Outperforming Larger Language Models with Less Training Data and Smaller Model Sizes](https://mlsanity.com/paper/2305.02301/) (2023)
- [DistiLLM: Towards Streamlined Distillation for Large Language Models](https://mlsanity.com/paper/2402.03898/) (2024)
- [MiniLLM: On-Policy Distillation of Large Language Models](https://mlsanity.com/paper/2306.08543/) (2023)
- [Gold-medalist Performance in Solving Olympiad Geometry with AlphaGeometry2](https://mlsanity.com/paper/2502.03544/) (2025)
- [DistiLLM-2: A Contrastive Approach Boosts the Distillation of LLMs](https://mlsanity.com/paper/2503.07067/) (2025)
- [Asymptotics of Language Model Alignment](https://mlsanity.com/paper/2404.01730/) (2024)
- [Pangu-Agent: A Fine-Tunable Generalist Agent with Structured Reasoning](https://mlsanity.com/paper/2312.14878/) (2023)
- [On a few pitfalls in KL divergence gradient estimation for RL](https://mlsanity.com/paper/2506.09477/) (2025)
- [Step-wise Adaptive Integration of Supervised Fine-tuning and Reinforcement Learning for Task-Specific LLMs](https://mlsanity.com/paper/2505.13026/) (2025)
- [On Almost Surely Safe Alignment of Large Language Models at Inference-Time](https://mlsanity.com/paper/2502.01208/) (2025)
- [SIKeD: Self-guided Iterative Knowledge Distillation for mathematical reasoning](https://mlsanity.com/paper/2410.18574/) (2024)

## Built on by (most cited first, 3 of 3)

- [A Survey of On-Policy Distillation for Large Language Models](https://mlsanity.com/paper/2604.00626/) (2026)
- [Decoding as Optimisation on the Probability Simplex: From Top-K to Top-P (Nucleus) to Best-of-K Samplers](https://mlsanity.com/paper/2602.18292/) (2026)
- [A Safe Action Is Not Enough: Feasible-Future Decoding for Vision-Language-Action Policies](https://mlsanity.com/paper/2610.05166/) (2026)

## Abstract (arXiv)

We introduce a novel approach to large language model (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.