Resumen Técnico: OPD+: Replanteando el Diseño de la Ventaja para la Destilación On-Policy

1. Planteamiento del Problema

La destilación on-policy (OPD) es una técnica utilizada para transferir capacidades de modelos de lenguaje docentes capaces a modelos estudiantes base mediante la optimización de una señal de supervisión densa a nivel de token derivada de los rollouts generados por el estudiante. Mientras que las implementaciones de OPD existentes (por ejemplo, Tinker) suelen utilizar la divergencia de Kullback-Leibler (KL) como el objetivo principal, emplean una elección de diseño específica: detener el gradiente (stop the gradient) en el término de recompensa (la diferencia de log-razón entre la verosimilitud del estudiante y la del docente).

Este diseño está motivado por el deseo de estabilidad numérica, trazando paralelismos con las redes de objetivos en Deep Q-Networks (DQN). Sin embargo, los autores identifican un fallo teórico crítico en este enfoque:

1. Estimaciones de Gradiente Sesgadas: Para las $f$-divergencias generales, la función de recompensa depende de los parámetros de la política del estudiante $\theta$. Detener el gradiente en este término de recompensa ignora la dependencia directa de la recompensa sobre $\theta$, lo que conduce a estimaciones de gradiente matemáticamente incorrectas.
2. Degradación del Rendimiento: Este sesgo causa un colapso catastrófico en el entrenamiento para funciones de divergencia distintas a la KL Inversa (por ejemplo, KL Directa y Divergencia de Jensen-Shannon), volviéndolas imposibles de entrenar bajo configuraciones estándar de OPD. Incluso para la KL Inversa, aunque el sesgo desaparece teóricamente debido a la identidad del score, los resultados empíricos sugieren que el diseño de detención de gradiente sigue siendo subóptimo para el rendimiento.

2. Metodología: El Marco de Trabajo OPD+

El artículo propone OPD+, un marco de optimización corregido que reevalúa el diseño de la ventaja para la destilación on-policy basada en la $f$-divergencia.

2.1 Fundamento Teórico

Los autores formulan el objetivo de OPD como la minimización de la $f$-divergencia entre la política del estudiante $p_\theta$ y la política del docente $q$. La función objetivo se define como:
$$J_f(\theta) = \mathbb{E}_{x \sim X, y \sim p_\theta(\cdot|x)} \left[ \sum_{n=1}^{L_y} f(u_{\theta,n}) \right]$$
donde $u_{\theta,n} = q(y_n | y_{<n}, x) / p_\theta(y_n | y_{<n}, x)$ es la razón de densidad a nivel de token, y $f$ es una función convexa con $f(1)=0$.

2.2 La Corrección

El artículo deriva el gradiente real de este objetivo. La OPD estándar utiliza un estimador de gradiente de política de "detención de gradiente" donde la recompensa $r = -f(u_{\theta,n})$ se trata como fija. Los autores demuestran que el gradiente real contiene dos términos:

1. El término estándar de la función de score: $\mathbb{E}[r \nabla_\theta \log p_\theta]$.
2. Un término de gradiente de recompensa directo: $\mathbb{E}[\nabla_\theta r]$.

Para una $f$-divergencia general, el término de gradiente de recompensa directo no desaparece. Específicamente, $\nabla_\theta r = u_{\theta,n} f'(u_{\theta,n}) \nabla_\theta \log p_\theta$.
Consecuentemente, la función de ventaja correcta no es simplemente $-f(u)$, sino un término ponderado:
$$w_f(u) = -f(u) + u f'(u)$$

2.3 Caso Especial: KL Inversa

Para la KL Inversa, $f(u) = -\log u$. El término de corrección se convierte en:
$$w(u) = -(-\log u) + u(-1/u) = \log u + 1$$
El término $+1$ adicional es una base constante (baseline). Dado que $\mathbb{E}[\nabla_\theta \log p_\theta] = 0$ (identidad del score), este término desaparece teóricamente en la esperanza. Sin embargo, los autores argumentan que, en la práctica, esta base constante actúa como un mecanismo de reducción de varianza, mejorando la velocidad de convergencia y la estabilidad.

2.4 Implementación

OPD+ requiere solo un cambio de una sola línea de código en las implementaciones de OPD existentes: reemplazar la recompensa de divergencia negativa bruta $-f(u)$ por el peso corregido $w_f(u)$. Este marco soporta varias $f$-divergencias, incluyendo KL Directa, KL Inversa y la Divergencia de Jensen-Shannon (JSD).

3. Contribuciones Clave

1. Marco de Optimización Genérico: El artículo proporciona un marco unificado para la destilación on-policy basado en $f$-divergencias generales, derivando el gradiente y las funciones de pérdida correctos que consideran la dependencia de la recompensa de los parámetros de la política.
2. Prueba Teórica de Sesgo: Demuestra matemáticamente que la operación estándar de detención de gradiente conduce a estimaciones de gradiente sesgadas para las $f$-divergencias generales, invalidando la práctica común de tratar la recompensa como fija.
3. Algoritmo OPD+: Los autores introducen OPD+, el algoritmo corregido que elimina el colapso de entrenamiento observado en la KL Directa y JSD, y mejora el rendimiento incluso para la KL Inversa.
4. Validación Empírica: El trabajo desafía la creencia predominante de que la KL Inversa es la métrica de distancia superior o la única viable para la destilación on-policy, demostando que otras divergencias pueden superar a esta cuando se optimizan correctamente.

4. Resultados Experimentales

Los autores evaluaron OPD+ en la familia de modelos QWEN3 (1.7B y 8B parámetros) en dos dominios distintos: Matemáticas de Competición (DeepMath, AIME 24/25, HMMT 25) y Uso de Herramientas Agénticas (Agentic Tool Use).

4.1 Rescate de Divergencias "Imposibles de Entrenar"

Bajo el diseño estándar de detención de gradiente (OPD), la KL Directa y JSD exhibieron un fallo catastrófico, cayendo al 0% de precisión en tareas de uso de herramientas y fallando en igualar la línea base de inicialización SFT en tareas matemáticas.

• Corrección OPD+: Al aplicar la ventaja corregida $w(u)$, OPD+ eliminó completamente este modo de fallo.
• Rendimiento de JSD: JSD+ no solo se estabilizó, sino que superó la línea base de la KL Inversa. En el desafiante benchmark AIME 25, JSD+ alcanzó un 50.42% de precisión frente al 48.54% de la KL Inversa. En la tarea de Uso de Herramientas, JSD+ alcanzó el 60.29%, superando la línea base de la KL Inversa de 59.97%.

4.2 Mejora en la KL Inversa

Incluso dentro del dominio donde la KL Inversa es tradicionalmente dominante, OPD+ produjo mejoras empíricas estrictas:

• AIME 24: KL Inversa+ mejoró del 64.27% al 64.69%.
• Uso de Herramientas: KL Inversa+ mejoró del 59.97% al 61.44%.
• Convergencia: La dinámica de entrenamiento (Figura 1) muestra que las variantes de OPD+ aceleran la convergencia temprana y previenen los estancamientos o la degradación del rendimiento observados en las líneas base estándar.

4.3 Autodestilación (OPSD)

El método también fue validado en un entorno de Autodestilación On-Policy (OPSD), donde el modelo actúa tanto como estudiante como docente (con el docente condicionado en soluciones de verdad fundamental privilegiadas). OPD+ superó consistentemente a las versiones vanilla de detención de gradiente en todos los tipos de divergencia, siendo la KL Inversa corregida la que mostró el mejor rendimiento entre las divergencias.

5. Significado y Reivindicaciones

El artículo sostiene que el bajo rendimiento histórico de las $f$-divergencias alternativas (KL Directa, JSD) en la destilación on-policy no es una limitación intrínseca de las divergencias mismas, sino un efecto secundario de un coeficiente de gradiente mal especificado causado por el diseño de detención de gradiente.

Al corregir la función de ventaja para que sea matemáticamente fiel al objetivo de la $f$-divergencia, OPD+:

1. Expande el Espacio de Diseño: Permite el uso de una variedad más amplia de distancias estadísticas para la destilación, yendo más allá del estándar de la industria de la KL Inversa.
2. Mejora el Rendimiento: Establece nuevos picos de rendimiento estado del arte en benchmarks de razonamiento y uso de herramientas.
3. Desafía Convenciones: Cuestiona la creencia común de que la KL Inversa es inherentemente superior debido a su estabilidad, sugiriendo que dicha estabilidad era un artefacto del método de optimización y no de la elección de la divergencia.

Los autores concluyen que su trabajo proporciona un algoritmo más matemáticamente correcto y empíricamente efectivo, lo que potencialmente puede desencadenar más investigación en el diseño de funciones de ventaja para métodos de destilación on-policy. Reconocen limitaciones respecto al análisis de escalado de modelos y la comparación entre verosimilitud de un solo token y destilación de logits, señalándolos como áreas para trabajos futuros.