Trust-Region Behavior Blending for On-Policy Distillation

Resumen Técnico: Mezcla de Comportamiento en Región de Confianza para la Destilación On-Policy

Declaración del Problema

La destilación on-policy (OPD) aborda el desajuste de prefijos inherente a la destilación offline mediante el entrenamiento de un modelo estudiante sobre prefijos muestreados de su propia política actual mientras se iguala a un profesor más fuerte. Sin embargo, este enfoque introduce una fragilidad crítica en las etapas iniciales del entrenamiento: los modelos estudiantes débiles suelen generar prefijos de baja calidad. Cuando la distribución de trayectorias del estudiante es deficiente, la supervisión del profesor se aplica a prefijos de baja calidad o "inutilizables", lo que limita la efectividad de la señal de aprendizaje.

Los intentos previos para estabilizar este proceso suelen implicar mover la recolección fuera de la política (off-policy) (por ejemplo, inyectando tokens del profesor o cambiando las distribuciones objetivo). Por el contrario, las ejecuciones (rollouts) puras del estudiante preservan la distribución pero sufren de una baja calidad inicial. El desafío central es mejorar la calidad de las ejecuciones tempranas sin desplazar permanentemente el proceso de recolección fuera de la política ni alterar el objetivo fundamental de destilación por prefijo.

Metodología: Mezcla de Comportamiento en Región de Confianza (TRB)

Los autores proponen la Mezcla de Comportamiento en Región de Confianza (TRB), un método de calentamiento (warmup) diseñado para controlar la política de comportamiento durante la fase inicial de recolección de ejecuciones. TRB opera bajo el principio de encontrar una política de comportamiento que sea lo más cercana posible al profesor, pero estrictamente constreñida dentro de una "región de confianza" centrada en la política actual del estudiante.

1. Optimización por Prefijo
En cada prefijo de generación $h$, TRB define una política de comportamiento $\mu^*(\cdot | h)$ resolviendo un problema de optimización con restricciones. El objetivo es minimizar la divergencia de Kullback-Leibler (KL) respecto a la política del profesor $\pi_T$, sujeto a una restricción de divergencia KL respecto a la política del estudiante $\pi_S$:

$$\mu^*(\cdot | h) = \arg \min_{\mu} D_{KL}(\mu \parallel \pi_T) \quad \text{sujeto a} \quad D_{KL}(\mu \parallel \pi_S) \le \epsilon$$

Aquí, $\epsilon$ representa la desviación local permitida respecto al estudiante. Esta formulación asegura que la distribución de muestreo sea atraída hacia los tokens respaldados por el profesor, pero permanezca acotada dentro de una región de confianza centrada en el estudiante.

2. Solución de Forma Cerrada
La optimización produce una solución de forma cerrada donde la política de comportamiento es una interpolación geométrica entre las políticas del estudiante y del profesor:
$$\mu_\beta(a | h) = \frac{\pi_S(a | h)^{1-\beta} \pi_T(a | h)^\beta}{Z_\beta(h)}$$
El coeficiente $\beta^*(h)$ es el valor máximo factible en $[0, 1]$ tal que se satisfaga la restricción de KL. Si el profesor ya se encuentra dentro de la región de confianza, $\mu^*$ se convierte en el profesor; si $\epsilon=0$, revierte al estudiante.

3. Calentamiento con Atenuación (Annealed Warmup)
TRB se aplica únicamente durante una fase inicial de calentamiento. El presupuesto de KL $\epsilon$ se atenúa linealmente a cero a lo largo de un horizonte $K$:
$$\epsilon_k = \epsilon_0 \left(1 - \frac{k}{K}\right)$$
Esto permite una fuerte guía del profesor al inicio del entrenamiento para mejorar la calidad del prefijo, mientras que gradualmente devuelve la recolección de ejecuciones al muestreo puro del estudiante a medida que este mejora. Crucialmente, la pérdida de destilación de KL inversa por prefijo permanece inalterada; solo se modifica la política de recolección de datos.

Contribuciones Clave

• Una Nueva Estrategia de Calentamiento: TRB introduce un método para controlar el comportamiento de las ejecuciones tempranas mediante la optimización de una política guiada por el profesor bajo una restricencia explícita de KL centrada en el estudiante, en lugar de simplemente inyectar tokens o alterar la función de pérdida.
• Preservación de las Propiedades On-Policy: A diferencia de los métodos que mezclan persistentemente las políticas del profesor y del estudiante o cambian la distribución objetivo, TRB asegura que el entrenamiento eventualmente regrese a las ejecuciones puramente on-policy, preservando la distribución objetivo para la fase principal de entrenamiento.
• Formulación Teórica: El artículo proporciona una solución de forma cerrada para la política de comportamiento que equilibra la guía del profesor con la estabilidad del estudiante, justificada por la monotonicidad de la divergencia KL con respecto al coeficiente de interpolación.

Resultados Experimentales

Los autores evaluaron TRB en dos configuraciones de destilación de razonamiento matemático:

1. Qwen3-1.7B-Base destilado desde Qwen3-8B.
2. Qwen3-0.6B-Base destilado desde Qwen3-4B.

Desempeño:

• TRB logró el mayor desempeño promedio entre todos los métodos comparados en ambas configuraciones, superando a la OPD vanilla, la mezcla de $\epsilon$ fijo, Veto, SKD (inyección intercalada del profesor), el calentamiento por temperatura y el calentamiento por SFT.
• En la configuración de Qwen3-1.7B, TRB logró un pass@1 promedio de 33.2%, comparado con el 32.3% de la OPD vanilla.
• En la configuración de Qwen3-0.6B, TRB logró 44.4%, comparado con el 44.0% de la OPD vanilla.

Análisis:

• Dinámica del Entrenamiento Temprano: El análisis de trayectorias (Figura 2) muestra que, aunque algunos métodos (como SKD) aumentan más rápido inicialmente, no necesariamente mantienen el desempeño final más alto. TRB proporciona un "inicio rápido" que conduce a resultados finales superiores.
• Entropía y Calidad del Prefijo: Durante la fase de calentamiento, TRB reduce la entropía media de los tokens del profesor en los prefijos visitados, lo que indica una guía de mayor calidad. Las sondas de paso 0 (Figura 4) demuestran que los prefijos generados bajo TRB producen tasas de éxito más altas cuando son continuados ya sea por el profesor o por el estudiante en comparación con los prefijos de la OPD vanilla, lo que sugiere que TRB logra desplazar con éxito la distribución de estados tempranos hacia trayectorias más aprendibles.

Significancia y Reclamaciones

El artículo afirma que TRB aborda eficazmente la fragilidad de la destilación on-policy temprana al mejorar la calidad de la distribución de datos inicial sin desviarse permanentemente del objetivo on-policy.

• Alcance Modesto: Los autores declaran explícitamente que su estudio se limita a dos entornos de razonamiento matemático con pares Qwen3-Base y un protocolo de evaluación basado en la corrección. No afirman que los esquemas de calentamiento específicos se transfieran sin cambios a otros dominios o brechas entre profesor y estudiante.
• Compromiso de Costo: El artículo reconoce que TRB aumenta los costos de tiempo de entrenamiento durante la fase de calentamiento debido al requerimiento de decodificación online del profesor y la co-residencia estudiante-profesor. Sin embargo, señalan que estos costos son temporales, ya que el entrenamiento regresa al perfil de ejecución estándar y más eficiente de la OPD tras el horizonte de calentamiento.
• Perspectiva Central: Los resultados sugieren que el comportamiento off-policy guiado por el profesor es más beneficioso cuando los prefijos visitados por el estudiante están significativamente desalineados con el profesor. Una vez resuelto este desajuste, la guía off-policy continua puede ofrecer rendimientos decrecientes, lo que respalda la estrategia de atenuar la política de comportamiento hacia el estudiante.