Replanteamiento de la Distilación de Modelos de Lenguaje a Gran Escala: Una Perspectiva de Procesos de Decisión de Markov con Restricciones

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

Resumen

Presentamos un enfoque novedoso para la destilación de modelos de lenguaje de gran escala (LLM, por sus siglas en inglés) al formularlo como un problema de aprendizaje por refuerzo con restricciones. Si bien trabajos recientes han comenzado a explorar la integración de recompensas específicas de tareas en los procesos de destilación, los métodos existentes suelen depender de ponderaciones de recompensas ad-hoc. Proponemos un marco de optimización fundamentado que maximiza las recompensas específicas de la tarea mientras restringe la divergencia del modelo maestro para que permanezca por debajo de un umbral especificado. Nuestro enfoque adapta el aprendizaje por refuerzo con aumento de estado restringido al contexto de destilación, introduciendo una función de recompensa modificada que mantiene garantías teóricas de satisfacción de restricciones sin requerir aumento de estado ni acceso al modelo maestro durante la implementación, y sin la sobrecarga computacional de los métodos duales de Lagrange. A través de experimentos exhaustivos en tareas de razonamiento matemático, demostramos que nuestro método logra mejores tasas de satisfacción de restricciones y un razonamiento superior en comparación con los baselines de relajación suave de Lagrange, manteniendo un rendimiento competitivo en las tareas. Nuestro marco proporciona una solución teóricamente fundamentada y prácticamente eficiente para la destilación consciente de recompensas en entornos con 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.