Rethinking Large Language Model Distillation: A Constrained Markov
  Decision Process Perspective
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.
68dba231d2bf1f4b15ec7937
Matthieu Zimmer
claimed_verified
2025-09-30T11:42:52.564Z
66c5beb29fc647b08321cc20
/avatars/d84e497cbcd44e7e06449092033f5bbd.svg
Matthieu Zimmer
user
MatthieuZ
68dba231d2bf1f4b15ec7938
Xiaotong Ji
68dba231d2bf1f4b15ec7939
Tu Nguyen
claimed_verified
2025-09-30T11:42:50.275Z
64b50bea372d4340778befca
/avatars/b1fa3ca96ce9bd3419c9c87dd57b3b9d.svg
Tu Nguyen
user
tumeteor
68dba231d2bf1f4b15ec793a
Haitham Bou Ammar
Rethinking Large Language Model Distillation: A Constrained Markov
  Decision Process Perspective
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.
A novel constrained reinforcement learning framework for LLM distillation maximizes task-specific rewards while maintaining constraint satisfaction without state augmentation or dual Lagrangian methods.
large language model (LLM) distillation constrained reinforcement learning task-specific rewards reward weighting constrained state augmented reinforcement learning reward function constraint satisfaction soft Lagrangian relaxation mathematical reasoning tasks resource-constrained settings