Self-Distilled Policy Gradient

Resumen Técnico: Gradiente de Política Autodestilado (SDPG)

Planteamiento del Problema
El Aprendizaje por Refuerzo con Recompensas Verificables (RLVR) se ha convertido en un estándar para el post-entrenamiento de Grandes Modelos de Lenguaje (LLMs) en tareas de razonamiento como la generación de matemáticas y código. Algoritmos como la Optimización de Política Relativa de Grupo (GRPO) optimizan contra recompensas de resultados basadas en reglas (por ejemplo, respuestas correctas/incorrectas) sin requerir anotaciones de preferencia humana. Sin embargo, el RLVR enfrenta dos limitaciones primarias:

1. Recompensas Dispersas: La señal de recompensa se aplica a nivel de secuencia, lo que dificulta la asignación de crédito a pasos de razonamiento específicos (tokens), lo que conduce a una exploración ineficiente.
2. Inestabilidad en el Entrenamiento: Las etapas tempranas del entrenamiento suelen sufrir de ventajas negativas, lo que causa inestabilidad y una potencial degradación de las capacidades lingüísticas fundamentales.

Aunque se ha propuesto la Destilación On-Policy (OPD) para proporcionar supervisión densa a nivel de token, los enfoques tradicionales dependen de modelos maestros externos más grandes, lo que incurre en altos costos de memoria e introduce desajustes de distribución. Los métodos recientes de "autodestilación", donde el modelo actúa tanto como estudiante como maestro utilizando contexto privilegiado (por ejemplo, soluciones de verdad fundamental), abordan el problema de la memoria pero corren el riesgo de reforzar tokens localmente plausibles en trayectorias globalmente incorrectas si no se controlan adecuadamente. Además, la autodestilación pura puede conducir al colapso de modo y a una reducción de la exploración.

Metodología: Marco de Trabajo SDPG
El artículo propone el Gradiente de Política Autodestilado (SDPG), un marco que integra la autodestilación on-policy de vocabulario completo exacta con la optimización de política regularizada por KL. El objetivo central combina tres componentes:

1. Gradiente de Política Basado en Resultados ($L_{out}$):
   SDPG conserva la señal de recompensa binaria dispersa de un verificador. Utiliza una ventaja relativa de grupo ($A_{out}$) normalizada dentro de un grupo de ejecuciones (rollouts) para estabilizar el aprendizaje, de manera similar a GRPO pero formulado como una pérdida de sustituto directa tipo REINFORCE sin el recorte de la razón de importancia estilo PPO.

2. Destilación On-Policy de Vocabulario Completo Exacta ($L_{OPD}$):
   En lugar de muestrear tokens para aproximar la señal del maestro, SDPG calcula la divergencia de Kullback-Leibler (KL) inversa exacta entre la distribución del estudiante ($p_t$) y una distribución maestra privilegiada ($q_t$) condicionada a un contexto adicional $c$ (por ejemplo, soluciones o rutas de razonamiento).

   • Identidad del Gradiente: Los autores derivan un resultado teórico clave (Proposición 3.1): el gradiente de la pérdida de KL inversa de vocabulario completo con respecto a la política del estudiante es localmente idéntico a una actualización de gradiente de política de muestreo desvinculado (detached-sampling). La "ventaja" en esta actualización es un log-ratio centrado de las probabilidades del maestro y del estudiante.
   • Mecanismo de Puerta (Gating): Para evitar que el maestro privilegiado refuerce trayectorias incorrectas, la pérdida de destilación está controlada por una puerta basada en la ventaja de resultado. La destilación solo se aplica cuando la ejecución es verificada como correcta ($A_{out} > 0$).

3. Regularización de la Política de Referencia ($L_K$):
   Para evitar que la política se aleje demasiado de sus capacidades iniciales o colapse, SDPG ancla la política actualizada a una política de referencia fija ($\pi_{ref}$). El artículo emplea la divergencia KL No Normalizada (UKL) (hacia adelante o inversa) para esta regularización, la cual ofrece gradientes más elegantes y simétricos en comparación con la KL estandarizada, particularmente cuando se trata de medidas no normalizadas.

Estrategias de Estabilización
Para gestionar el ruido de la señal privilegiada y el equilibrio entre exploración y explotación, SDPG introduce dos estabilizadores:

• Puerta de Ventaja Positiva: Asegura que la destilación solo ocurra en trayectorias respaldadas por el verificador.
• Programación de Calentamiento-Decaimiento ($\beta(k)$): El coeficiente de destilación $\beta$ se incrementa gradualmente (calentamiento) para permitir que el modelo encuentre trayectorias correctas antes de depender de la destilación, y luego se reduce (decaimiento) hacia el final del entrenamiento. Esto evita que la señal privilegiada sobre-constriña al modelo una vez que este ha internalizado el conocimiento necesario, fomentando la exploración en la etapa final.

Contribuciones Clave

1. Derivación Teórica: El artículo deriva una forma exacta de gradiente de política local para la destilación on-policy de vocabulario completo de KL inversa, mostrando que es equivalente a una actualización de muestreo desvinculado con una ventaja de log-ratio centrado.
2. Objetivo Novedoso: Propone un objetivo unificado que combina recompensas de resultados binarios, destilación privilegiada de vocabulario completo con puerta de control y regularización KL no normalizada contra una referencia fija.
3. Técnicas de Estabilización: La introducción de la puerta de ventaja positiva y un programa de calentamiento-decaimiento para el coeficiente de destilación para equilibrar la supervisión densa con la exploración.

Resultados Experimentales
Los autores evaluaron SDPG en benchmarks de razonamiento matemático (AIME2024, AIME2025, AMC23) utilizando Qwen3-4B.

• Desempeño: Ambas variantes de SDPG (usando KL Hacia Adelante No Normalizada y KL Inversa No Normalizada) superaron a la línea base GRPO y a la línea base de autodestilación RLSD en todos los benchmarks.
• Estabilidad: Las variantes de SDPG alcanzaron mesetas de alta recompensa significativamente más rápido que GRPO. Notablemente, SDPG mantuvo una entropía de actor más alta durante todo el entrenamiento en comparación con RLSD, el cual exhibió un colapso de entropía (un signo de colapso de modo) alrededor del paso 250.
• Eficiencia: Los métodos SDPG produjeron longitudes de respuesta que fueron suficientes para el razonamiento de múltiples pasos pero más concisas que las salidas verbosas de GRPO.

Significado
El artículo afirma que SDPG logra cerrar la brecha entre el RLVR de recompensa dispersa y la autodestilación de supervisión densa. Al tratar la autodestilación como una actualización de gradiente de política y controlar rigurosamente la señal privilegiada mediante puertas y programación, SDPG logra una estabilidad y un desempeño superiores. Demuestra que la supervisión densa a nivel de token derivada del propio contexto privilegiado del modelo puede guiar eficazmente a los modelos de razonamiento sin la sobrecarga de memoria de maestros externos o la inestabilidad de la autodestilación pura.