Reinforcement-aware Knowledge Distillation for LLM Reasoning

Resumen Técnico: Destilación de Conocimiento con Conciencia de RL para el Razonamiento de LLM (RLAD)

1. Planteamiento del Problema

Los Grandes Modelos de Lenguaje (LLMs) han alcanzado capacidades de razonamiento significativas mediante el entrenamiento post-RL (Aprendizaje por Refuerzo), particularmente en tareas de cadena de pensamiento (CoT) larga. Sin embargo, los altos costos de inferencia de estos modelos grandes hacen necesario su destilación en modelos estudiantes más pequeños. Los métodos existentes de Destilación de Conocimiento (KD) enfrentan limitaciones críticas cuando se aplican al entrenamiento post-RL:

• Desajuste de Distribución: Los métodos de KD estándar (offline, off-policy o on-policy) suelen depender de trazas de profesor fijas o de una regularización de divergencia KL estática. En RL, la política del estudiante evoluciona continuamente, lo que provoca que la distribución de las trayectorias (rollouts) se desplace. Los objetivos de destilación estáticos se vuelven subóptimos o inestables a medida que el estudiante se aleja de la distribución inicial del profesor.
• Interferencia de Objetivos: Los métodos que integran KD en el bucle de RL (por ejemplo, KDRL) típicamente añaden un término de divergencia KL entre el estudiante y el profesor como una penalización de regularización. Esto crea un problema de optimización multiobjetivo donde el objetivo de maximización de recompensa compite con el objetivo de imitación. Esta competencia requiere un ajuste cuidadoso de hiperparámetros (balanceo de pérdida) y puede conducir a una dinámica de entrenamiento inestable, ya que la supervisión del profesor puede no alinearse con las estimaciones de ventaja actuales del estudiante.

2. Metodología: Destilación con Conciencia de RL (RLAD)

Los autores proponen RLAD, un marco diseñado para realizar una imitación selectiva durante el entrenamiento post-RL. El principio central es que el estudiante solo debe ser guiado por el profesor cuando dicha guía mejore la actualización de la política actual bajo el objetivo de RL.

Componente Clave: Destilación de Ratio de Región de Confianza (TRRD)
En lugar de utilizar un regularizador KL independiente, RLAD introduce TRRD, que reemplaza el ratio de muestreo de importancia estándar en la Optimización de Política Relativa de Grupo (GRPO) por una mezcla profesor-estudiante.

• Mecanismo: TRRD define un nuevo ratio de verosimilitud $r^{TRRD}_{i,t}$ para el $t$-ésimo token de la $i$-ésima respuesta:
  $$r^{TRRD}_{i,t}(\theta_S) = \left( \frac{\pi_{\theta_S}(y_t | x, y_{<t})}{\pi_{\theta_S,old}(y_t | x, y_{<t})} \right)^\alpha \left( \frac{\pi_{\theta_S}(y_t | x, y_{<t})}{\pi_{\theta_T}(y_t | x, y_{<t})} \right)^{1-\alpha}$$
  Aquí, $\pi_{\theta_S}$ es la política actual del estudiante, $\pi_{\theta_S,old}$ es la política previa del estudiante y $\pi_{\theta_T}$ es la política del profesor. $\alpha \in [0, 1]$ es un coeficiente de mezcla.
• Interpretación de Región de Confianza: Este ratio puede verse como una actualización estilo PPO donde la región de confianza está anclada a una mezcla del antiguo estudiante y el profesor ($r_{\pi_{mix}}$), en lugar de solo al antiguo estudiante.
• Selección Consciente de la Ventaja: La actualización se pondera mediante la ventaja relativa de grupo $\hat{A}_{i,t}$.
  • Ventaja Positiva: Si la ventaja es positiva, se incentiva al estudiante a aumentar la probabilidad del token. Si el profesor también asigna una alta probabilidad, la mezcla de anclaje es mayor, permitiendo al estudiante aumentar su probabilidad antes de alcanzar el límite de recorte (permitiendo actualizaciones más grandes). Si el profesor asigna una baja probabilidad, el ancla es menor, restringiendo la magnitud de la actualización.
  • Ventaja Negativa: Si la ventaja es negativa, se incentiva al estudiante a disminuir la probabilidad. Si el profesor asigna una alta probabilidad, el límite de recorte evita que el estudiante suprima agresivamente el token, preservando comportamientos alineados con el profesor.
  • Ventaja Cercana a Cero: La influencia del profesor se vuelve insignificante, evitando la interferencia cuando la señal de recompensa no es informativa.

Este diseño equilibra implícitamente la exploración, la explotación y la imitación al asegurar que la guía del profesor sea modulada por la señal de recompensa (ventaja) y esté constreñida dentro de una región de confianza.

3. Contribuciones Clave

1. Marco RLAD: Un nuevo marco de destilación para el entrenamiento post-RL que integra la guía del profesor mediante la Destilación de Ratio de Región de Confianza (TRRD), reemplazando la regularización KL estática con un objetivo de ratio de verosimilitud recortado y ponderado por la ventaja.
2. Análisis Teórico: Un análisis que demuestra que TRRD permite una imitación selectiva condicionada por la recompensa. Preserva las actualizaciones limitadas por la región de confianza y equilibra naturalmente el compromiso entre seguir al profesor y optimizar las recompensas sin requerir un balanceo delicado de la pérdida.
3. Validación Empírica: Experimentos extensos en razonamiento lógico (K&K Logistics) y razonamiento matemático de contexto largo (AIME, AMC, MATH) que muestran mejoras consistentes sobre las líneas base.

4. Resultados Experimentales

Los autores evaluaron RLAD en tareas de razonamiento lógico y razonamiento matemático complejo utilizando modelos estudiantes que van desde 0.6B hasta 8B de parámetros y profesores de hasta 32B.

• Razonamiento Lógico (K&K Logistics):

  • Usando un profesor Qwen3-8B, RLAD mejoró significativamente la precisión promedio para Qwen3-0.6B (de 0.76 a 0.94 en 8K de contexto) y Qwen3-1.7B (de 0.95 a 0.99 en 8K de contexto) en comparación con GRPO estándar.
  • Las ganancias fueron más pronunciadas en los subconjuntos más difíciles (PPL7 y PPL8), donde RLAD mantuvo un margen claro sobre KDRL.
  • La dinámica de entrenamiento mostró una convergencia más rápida y recompensas finales más altas en comparación con tanto GRPO como KDRL.

• Razonamiento Matemático:

  • Modelos Base (estilo R1-Zero): En Qwen3-8B-Base, RLAD mejoró el rendimiento promedio en +5.5 puntos sobre GRPO en tareas de contexto largo (30K).
  • Modelos Post-entrenados: RLAD superó la destilación de KL explícita (KDRL) y la destilación offline.
  • Especificidad de Métricas: RLAD mostró mayores ganancias en Pass@1 (Mean@32) comparado con KDRL, que tendía a mostrar mayores ganancias relativas en Pass@32. Los autores interpretan esto como que RLAD logra con éxito la mejora de la política alineada con la recompensa (explotación), mientras que las ganancias de KDRL fueron impulsadas más por la pura imitación de soluciones diversas del profesor.
  • Estabilidad: RLAD exhibió dinámicas de entrenamiento más estables con menos oscilaciones en comparación con KDRL, que sufrió problemas de acoplamiento de objetivos.

• Robustez y Eficiencia:

  • Profesor Débil: Cuando el profesor era más débil que el estudiante, RLDO se mantuvo cerca del rendimiento de GRPO, mientras que KDRL se degradó más significativamente.
  • Hiperparámetros: Se encontró que RLAD es mayormente insensible al coeficiente de mezcla $\alpha$ (fijado en 0.5), a diferencia de KDRL que es sensible al peso de KL $\lambda$.
  • Eficiencia: Integrar RLAD añadió solo un ~12% de latencia de entrenamiento comparado con el GRPO estándar, ya que el profesor solo se utiliza para el cálculo de la log-probabilidad, no para la generación de trayectorias.

5. Significado y Afirmaciones

El artículo afirma que RLAD aborda el desalineamiento fundamental entre la supervisión estática del profesor y las actualizaciones dinámicas de la política de RL. Al integrar la guía del profesor dentro del mecanismo de actualización de la región de confianza, RLAD evita el "ajuste frágil" requerido por la regularización KL multiobjetivo.

Los autores aseveran que RLAD permite al estudiante moverse hacia comportamientos de alta ventaja dentro de una distribución de anclaje mientras explota oportunistamente la guía del profesor. Esto resulta en:

• Desempeño Superior: Supera consistentemente a GRPO, KDRL y la destilación offline, particularmente en benchmarks desafiantes.
• Estabilidad: Dinámicas de entrenamiento más estables y mejor convergencia.
• Verdadera Optimización de RL: A diferencia de los métodos que dependen fuertemente de la imitación (que pueden aumentar el Pass@32 pero no el Pass@1), RLAD prioriza el comportamiento alineado con la recompensa, lo que conduce a ganancias más fuertes en las respuestas correctas más probables del modelo (Pass@1).

El artículo concluye que, si bien el enfoque depende de modelos profesor de pesos abiertos y bien calibrados para el acceso a los logits, representa un paso significativo hacia la compatibilidad de la destilación con la naturaleza dinámica del entrenamiento post-RL.