Destilación On-Policy con Región de Confianza

Resumen

La Destilación en Política (OPD, por sus siglas en inglés) es una técnica fundamental para el post-entrenamiento eficiente de modelos de lenguaje de gran escala (LLMs), con amplias aplicaciones en aprendizaje de agentes, mejora multitarea y compresión de modelos. Sin embargo, el entrenamiento OPD se vuelve inestable cuando las distribuciones del profesor y del estudiante difieren sustancialmente, ya que la supervisión del profesor sobre los tokens generados por el estudiante puede producir gradientes de política poco fiables e incluso causar fallos en la optimización. Este trabajo aborda la supervisión fiable a nivel de token en política mediante estrategias de asignación de crédito y propone la Destilación en Política de Región de Confianza, TrOPD. Presenta las siguientes características: 1) Aprendizaje en Política de Región de Confianza: TrOPD realiza OPD solo en regiones donde el profesor proporciona supervisión fiable, mitigando la dificultad de optimización del estimador K1 de divergencia KL inversa bajo desajuste de distribuciones. 2) Estimación de valores atípicos: Para regiones atípicas, exploramos el recorte de gradientes, el enmascaramiento y la estimación de divergencia KL directa para reducir los efectos adversos de la supervisión no fiable. 3) Guía fuera de política: El estudiante continúa la generación a partir de prefijos del profesor y utiliza la divergencia KL directa para imitar la guía fuera de política, fomentando la exploración en política hacia regiones fiables. Los experimentos muestran que TrOPD supera consistentemente a las líneas base OPD de última generación, incluyendo OPD, EOPD y REOPOLD, en razonamiento matemático, generación de código y puntos de referencia de dominio general.

English

On-Policy Distillation (OPD) is a fundamental technique for efficient post-training of large language models (LLMs), with broad applications in agent learning, multi-task enhancement, and model compression. However, OPD training becomes unstable when the teacher and student distributions differ substantially, as teacher supervision on student-generated tokens may yield unreliable policy gradients and even cause optimization failure. This work addresses reliable on-policy token-level supervision through credit assignment strategies, and proposes Trust Region On-Policy Distillation, TrOPD. It features the following characteristics: 1) Trust-Region On-Policy Learning: TrOPD performs OPD only in regions where the teacher provides reliable supervision, mitigating the optimization difficulty of the K1 reverse-KL estimator under distribution mismatch. 2) Outlier Estimation: For outlier regions, we explore gradient clipping, masking, and forward-KL estimation to reduce the adverse effects of unreliable supervision. 3) Off-Policy Guidance: The student continues generation from teacher prefixes and uses forward KL to imitate off-policy guidance, encouraging on-policy exploration toward reliable regions. Experiments show that TrOPD consistently outperforms SoTA OPD baselines, including OPD, EOPD, and REOPOLD, across mathematical reasoning, code generation, and general-domain benchmarks.