Repenser la distillation des grands modèles de langage : une perspective de processus de décision markovien contraint

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

September 26, 2025

papers.authors: Matthieu Zimmer, Xiaotong Ji, Tu Nguyen, Haitham Bou Ammar

cs.AI

papers.abstract

Nous introduisons une nouvelle approche pour la distillation de modèles de langage de grande taille (LLM) en la formulant comme un problème d'apprentissage par renforcement contraint. Bien que des travaux récents aient commencé à explorer l'intégration de récompenses spécifiques aux tâches dans les processus de distillation, les méthodes existantes reposent généralement sur une pondération ad hoc des récompenses. Nous proposons un cadre d'optimisation fondé sur des principes qui maximise les récompenses spécifiques aux tâches tout en contraignant la divergence par rapport au modèle enseignant à rester en dessous d'un seuil spécifié. Notre approche adapte l'apprentissage par renforcement contraint avec augmentation d'état au contexte de la distillation, en introduisant une fonction de récompense modifiée qui maintient des garanties théoriques de satisfaction des contraintes sans nécessiter d'augmentation d'état ni d'accès au modèle enseignant pendant le déploiement, et sans la surcharge computationnelle des méthodes duales de Lagrange. À travers des expériences approfondies sur des tâches de raisonnement mathématique, nous démontrons que notre méthode atteint de meilleurs taux de satisfaction des contraintes et un meilleur raisonnement par rapport aux bases de référence de relaxation douce de Lagrange, tout en maintenant une performance compétitive sur les tâches. Notre cadre offre une solution théoriquement fondée et pratiquement efficace pour la distillation consciente des récompenses dans des contextes à ressources limitées.

English

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.