Resumen Técnico: Desacoplamiento de KL y Trayectorias en la Distilación de LLM

Enunciado del Problema

La distilación de conocimiento se ha convertido en una piedra angular del post-entrenamiento de Modelos de Lenguaje Grandes (LLM), particularmente para tareas de razonamiento, operando a menudo junto con el Aprendizaje por Refuerzo (RL). Sin embargo, el espacio de diseño de la distilación sigue siendo poco comprendido. Las pipelines actuales suelen seguir uno de dos paradigmas: distilación off-policy (entrenamiento sobre trazas generadas por el profesor, a menudo mediante Ajuste Fino Supervisado) o distilación on-policy (OPD) (entrenamiento sobre despliegues generados por el estudiante con retroalimentación del profesor).

El artículo argumenta que estos paradigmas predominantes acoplan implícitamente dos elecciones de diseño ortogonales:

1. Fuente del Prefijo: Si la supervisión se aplica a prefijos generados por el profesor (off-policy) o por el estudiante (on-policy).
2. Dirección KL a Nivel de Token: Si el objetivo minimiza la KL hacia adelante ($KL(p_{profesor} \parallel q_{estudiante})$) o la KL inversa ($KL(q_{estudiante} \parallel p_{profesor})$).

La práctica convencional acopla prefijos del profesor con KL hacia adelante y prefijos del estudiante con KL inversa. Los autores sostienen que este acoplamiento no es intrínseco a nivel de token y que desacoplar estos ejes revela un espacio de diseño más amplio con cuatro objetivos válidos, cada uno correspondiente a regímenes de aprendizaje distintos.

Metodología

Marco Teórico: Desacoplamiento de los Ejes

Los autores descomponen la divergencia KL a nivel de secuencia utilizando la regla de la cadena autoregresiva. Demuestran que:

• KL hacia adelante ($KL(p_T \parallel q_S)$) se asocia naturalmente con prefijos del profesor (off-policy) porque la esperanza se toma sobre la distribución de trayectorias del profesor.
• KL inversa ($KL(q_S \parallel p_T)$) se asocia naturalmente con prefijos del estudiante (on-policy) porque la esperanza se toma sobre la distribución de trayectorias del estudiante.

Al tratar estos como ejes independientes, los autores definen cuatro objetivos:

1. Prefijo del Profesor + KL hacia adelante: Corresponde a SFT off-policy (distilación estándar).
2. Prefijo del Estudiante + KL hacia adelante: Corresponde a SFT on-policy estilo DAgger.
3. Prefijo del Profesor + KL inversa: Corresponde a un objetivo de distilación estilo RL offline.
4. Prefijo del Estudiante + KL inversa: Corresponde a Distilación On-Policy (OPD), vista como RL on-policy con recompensa densa.

El artículo establece identidades a nivel de gradiente que muestran que:

• KL hacia adelante produce un gradiente de entropía cruzada estilo SFT con objetivos suaves del profesor.
• KL inversa produce un gradiente de política estilo RL con una recompensa densa definida por el log-ratio profesor-estudiante ($r = \log p_T - \log q_S$).

Configuración Experimental

Los autores realizan un estudio controlado sobre razonamiento matemático utilizando modelos Qwen3 (estudiante de 0.6B, profesores de 4B/8B). Evalúan los cuatro objetivos en dos regímenes:

1. Distilación Autónoma: Comparando el rendimiento final de los modelos entrenados exclusivamente con cada objetivo.
2. Distilación-Posterior-RL: Utilizando los puntos de control distilados como inicializaciones para la Optimización de Política Relativa por Grupos (GRPO).

Rastrean la precisión, la entropía predictiva y la longitud de la respuesta a través de diversas longitudes de secuencia (128 vs. 4096 tokens) y benchmarks (MATH500, GSM8K, AIME24, AMC23).

Contribuciones Clave

1. Identificación de Tres Compensaciones

El análisis empírico revela tres compensaciones fundamentales que gobiernan el diseño de la distilación:

• Compensación Precisión–Entropía (Dirección KL): La KL inversa logra consistentemente una mayor precisión autónoma (Avg@k), pero impulsa la entropía predictiva hacia el colapso, reduciendo la diversidad (menor Pass@k) y haciendo que el modelo sea menos estable para el RL posterior. La KL hacia adelante preserva la entropía, apoyando una mejora de RL más estable, aunque con una menor precisión autónoma inicial.
• Compensación Calidad–Cómputo (Fuente del Prefijo): Los prefijos del estudiante (on-policy) producen mayor precisión bajo pasos de entrenamiento emparejados al supervisar estados que el estudiante visita realmente. Sin embargo, los prefijos del profesor (off-policy) son más eficientes en cómputo bajo FLOPs emparejados porque permiten reutilizar logits del profesor en caché y trayectorias offline, evitando el costo de la generación online del estudiante.
• Compensación Precisión–Estabilidad (Longitud de Entrenamiento): La distilación de secuencias largas (4096 tokens) mejora la precisión de razonamiento, pero, particularmente bajo KL inversa, causa un colapso severo de entropía e "inflación de longitud" (respuestas que alcanzan el límite de generación). El entrenamiento de secuencias cortas es más estable pero menos preciso.

2. Métodos Propuestos para Equilibrar Compensaciones

Motivados por estos hallazgos, los autores proponen dos métodos prácticos:

• Mezcla de KL: Una combinación ponderada de KL hacia adelante y KL inversa en cada prefijo. Los autores encuentran que en la distilación de secuencias largas, una mezcla con predominio hacia adelante es crucial. Mantiene los beneficios de precisión de la KL inversa mientras previene el colapso de entropía y la inflación de longitud, estabilizando la dinámica de entrenamiento.
• Plan de Estudios de Longitud Puerta de Entropía: Un cronograma de entrenamiento que comienza con horizontes cortos y aumenta la longitud de la secuencia solo mientras la entropía predictiva permanezca por encima de un umbral de estabilidad. Este enfoque evita el colapso de entropía asociado con el entrenamiento fijo de horizonte largo.

Resultados

Distilación Autónoma

• KL Inversa generalmente supera a la KL hacia adelante en precisión autónoma (Avg@k), con ganancias de +2.45 puntos en promedio.
• Prefijos del Estudiante superan a los prefijos del profesor en precisión, especialmente en configuraciones de secuencias largas (+3.55 Avg@k de ganancia).
• Colapso de Entropía: La KL inversa pura en configuraciones de secuencias largas impulsa la entropía cerca de cero y causa que las longitudes de respuesta se inflen hasta el límite máximo (por ejemplo, 8192 tokens), degradando el Pass@k.

Distilación-Posterior-RL (GRPO)

• Estabilidad de Inicialización: Los modelos inicializados con KL hacia adelante (incluso si tienen menor precisión autónoma) muestran una mejora más fiable durante GRPO porque retienen mayor entropía.
• Peligro de KL Inversa: Los modelos inicializados con KL inversa a menudo comienzan con mayor precisión pero se degradan o se estancan durante el RL debido a la baja entropía y la capacidad reducida de exploración.
• Éxito del Plan de Estudios: El plan de estudios de longitud puerta de entropía mejora significativamente el rendimiento. Comparado con el entrenamiento fijo de 4096 tokens, mejoró Avg@k en 3.6 puntos y Pass@k en hasta 5.8 puntos, mientras que reducía la longitud promedio de respuesta en aproximadamente 3× (por ejemplo, de ~8185 a ~2688 tokens).

Resultados de la Mezcla de KL

En configuraciones de secuencias largas, una mezcla con un peso alto de KL hacia adelante (por ejemplo, 0.8 Hacia adelante / 0.2 Inversa) preservó la mayor parte de la precisión de la KL inversa mientras estabilizaba la entropía y prevenía la inflación de longitud.

Significado y Afirmaciones

El artículo afirma proporcionar un marco teórico unificado que conecta recetas de entrenamiento dispares (SFT, DAgger, RL Offline, OPD) como instancias específicas de dos elecciones de diseño ortogonales. Al desacoplar la fuente del prefijo y la dirección KL, los autores argumentan que los practicantes pueden ir más allá de paradigmas rígidos para diseñar objetivos de distilación de razonamiento que equilibren explícitamente:

• Precisión (mediante KL inversa y secuencias largas),
• Diversidad/Exploración (mediante KL hacia adelante y preservación de entropía),
• Eficiencia de Cómputo (mediante prefijos/caché del profesor), y
• Entrenabilidad de RL Aguas Abajo (mediante entropía estable).

Los métodos propuestos (mezcla de KL y plan de estudios puerta de entropía) se presentan como herramientas prácticas para navegar las compensaciones identificadas, permitiendo el entrenamiento de modelos de razonamiento más pequeños que sean tanto precisos como estables para una optimización de política posterior. Los autores enfatizan que el objetivo autónomo "mejor" no es necesariamente la mejor inicialización para una pipeline de RL, destacando la necesidad de una visión holística del proceso de post-entrenamiento.