Heroverdenken van Distillatie van Grote Taalmodellen: Een Beperkt Markov Beslissingsproces Perspectief

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

September 26, 2025

Auteurs: Matthieu Zimmer, Xiaotong Ji, Tu Nguyen, Haitham Bou Ammar

cs.AI

Samenvatting

We introduceren een nieuwe aanpak voor distillatie van grote taalmodellen (LLM) door het te formuleren als een geoptimaliseerd reinforcement learning-probleem met beperkingen. Hoewel recent werk begonnen is met het integreren van taakspecifieke beloningen in distillatieprocessen, vertrouwen bestaande methoden doorgaans op ad-hoc gewichtstoekenning van beloningen. Wij stellen een principieel optimalisatiekader voor dat taakspecifieke beloningen maximaliseert, terwijl de afwijking van het leraarmodel onder een gespecificeerde drempel wordt gehouden. Onze aanpast past constrained state augmented reinforcement learning aan voor de distillatie-instelling, waarbij een aangepaste beloningsfunctie wordt geïntroduceerd die theoretische garanties biedt voor het voldoen aan de beperkingen, zonder dat toegang tot het leraarmodel of state augmentation nodig is tijdens de implementatie, en zonder de rekenkundige overhead van duale Lagrangiaanse methoden. Door uitgebreide experimenten op wiskundige redeneertaken tonen we aan dat onze methode betere beperkingsvoldoening en beter redeneren bereikt in vergelijking met de soft Lagrangiaanse relaxatie-baselines, terwijl het competitieve taakprestaties behoudt. Ons kader biedt een theoretisch onderbouwde en praktisch efficiënte oplossing voor beloningsbewuste distillatie in omgevingen met beperkte middelen.

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.