Ripensare la distillazione dei grandi modelli linguistici: una prospettiva basata sui processi decisionali di Markov vincolati

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

September 26, 2025

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

cs.AI

Abstract

Introduciamo un nuovo approccio alla distillazione di modelli linguistici di grandi dimensioni (LLM) formulandolo come un problema di apprendimento per rinforzo vincolato. Mentre lavori recenti hanno iniziato a esplorare l'integrazione di ricompense specifiche per il compito nei processi di distillazione, i metodi esistenti si basano tipicamente su ponderazioni ad hoc delle ricompense. Proponiamo un framework di ottimizzazione basato su principi che massimizza le ricompense specifiche per il compito, vincolando la divergenza dal modello insegnante a rimanere al di sotto di una soglia specificata. Il nostro approccio adatta l'apprendimento per rinforzo vincolato con aumento dello stato al contesto della distillazione, introducendo una funzione di ricompensa modificata che mantiene garanzie teoriche di soddisfazione dei vincoli senza richiedere l'aumento dello stato o l'accesso al modello insegnante durante il deployment, e senza il sovraccarico computazionale dei metodi Lagrangiani duali. Attraverso esperimenti estesi su compiti di ragionamento matematico, dimostriamo che il nostro metodo raggiunge tassi di soddisfazione dei vincoli e capacità di ragionamento migliori rispetto ai baseline con rilassamento Lagrangiano soft, mantenendo al contempo prestazioni competitive nel compito. Il nostro framework fornisce una soluzione teoricamente fondata e praticamente efficiente per la distillazione consapevole delle ricompense in contesti con risorse limitate.

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.