Repensando a Distilação de Modelos de Linguagem de Grande Escala: Uma Perspectiva de Processo de Decisão de Markov Restrito

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

September 26, 2025

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

cs.AI

Resumo

Apresentamos uma nova abordagem para a destilação de modelos de linguagem de grande escala (LLM) ao formulá-la como um problema de aprendizado por reforço com restrições. Embora trabalhos recentes tenham começado a explorar a integração de recompensas específicas à tarefa em processos de destilação, os métodos existentes geralmente dependem de ponderação de recompensas ad-hoc. Propomos um framework de otimização fundamentado que maximiza as recompensas específicas da tarefa enquanto restringe a divergência do modelo professor a permanecer abaixo de um limite especificado. Nossa abordagem adapta o aprendizado por reforço com aumento de estado e restrições ao cenário de destilação, introduzindo uma função de recompensa modificada que mantém garantias teóricas de satisfação das restrições sem exigir aumento de estado ou acesso ao modelo professor durante a implantação, e sem a sobrecarga computacional dos métodos duais de Lagrange. Por meio de experimentos extensos em tarefas de raciocínio matemático, demonstramos que nosso método alcança melhores taxas de satisfação das restrições e melhor raciocínio em comparação com as linhas de base de relaxação suave de Lagrange, mantendo um desempenho competitivo nas tarefas. Nosso framework oferece uma solução teoricamente fundamentada e praticamente eficiente para destilação com consciência de recompensa em cenários com recursos limitados.

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.