Resumen Técnico: Auto-destilación Geométrica para la Generalización del Razonamiento

1. Planteamiento del Problema

La destilación on-policy (OPD) ha emergido como una estrategia de post-entrenamiento práctica para los modelos de lenguaje de gran tamaño (LLM), proporcionando una supervisión densa a nivel de token sobre las trayectorias generadas por el propio modelo estudiante. Una variante específica y efectiva, la auto-destilación con contexto privilegiado (OPSD), condiciona al "maestro" (el mismo modelo) con información adicional, como un rastro completo de la solución o una pista, mientras que el estudiante debe razonar únicamente a partir del contexto original del problema.

Si bien esta configuración proporciona una supervisión abundante, el artículo identifica un modo de fallo crítico: la deriva predictiva (predictive drift). Debido a que el maestro posee un contexto privilegiado, puede comprometerse con confianza en continuaciones que son obvias dada la solución, pero que permanecen inciertas o sin apoyo desde la propia perspectiva del estudiante. Los objetivos de destilación estándar (por ejemplo, la divergencia de Kullback–Leibler) ponderan las preferencias del maestro según su propia masa de probabilidad. En consecuencia, en estados donde el maestro tiene una alta confianza pero el estudiante asigna una probabilidad baja (poca superposición), la señal del gradiente es fuerte. Tras múltiples actualizaciones, estos tirones fuertes fuerzan al estudiante a adoptar las continuaciones confiadas del maestro antes de que la propia información del estudiante pueda justificarlas. Esta acumulación de actualizaciones desajustadas conduce a la derza (drift), donde el modelo suprime la incertidumbre y degrada el rendimiento del razonamiento fuera de la distribución (OOD), resultando a menudo en un "acuerdo confiado en respuestas incorrectas".

Los enfoques existentes intentan mitigar esto filtrando señales del maestro que no sean fiables o reformulando los pesos de la divergencia (por ejemplo, usando KL inversa o mecanismos de filtrado). Sin embargo, los autores argumentan que estos métodos no logran abordar el problema central: qué tan lejos debería mover un modelo el estudiante. La selección por sí sola no puede regular la magnitud de la actualización, y el movimiento sin control se acumula en una deriva global.

2. Metodología: GEOSD

Los autores introducen GEOSD (Auto-destilación Geométrica), un objetivo diseñado para contrarrestar esta deriva tratando el comportamiento predictivo como un movimiento dentro de un espacio geomético específico de distribuciones de siguiente token. El método opera de dos maneras complementarias:

A. Atenuación de la Presión Local (Pérdida de Hellinger)

En lugar de la divergencia KL estándar, GEOSD emplea una pérdida de Hellinger para modular el tirón del maestro en cada estado visitado.

• Mecanismo: La distancia de Hellinger se deriva de la geometría del simplex de probabilidad bajo un mapeo de raíz cuadrada ($p \to \sqrt{p}$), donde las distribuciones residen en un hiperesfera. El gradiente de la pérdida de Hellinger escala cada preferencia del maestro mediante la media geométrica de las probabilidades del estudiante y del maestro ($\sqrt{p_i q_i}$).
• Efecto: Si el maestro asigna una alta probabilidad a un token pero el estudiante asigna una probabilidad cercana a cero, el término de superposición se desvanece y el tirón del gradiente se atenúa. Por el contrario, las preferencias compartidas se transfieren con fuerza. Esto asegura que la supervisión fortalezca gradualmente a medida que el estudiante y el maestro llegan a un acuerdo, en lugar de forzar al estudiante a adoptar de inmediato las continuaciones confiadas pero sin apoyo del maestro.

B. Control de Deriva Global (Término Proximal de Fisher–Rao)

Aunque la ponderación por superposición limita el tirón en estados individuales, no evita la acumulación de pequeñas actualizaciones razonables en una gran deriva global durante el entrenamiento.

• Mercanismo: GEOSD añade un término proximal que penaliza la distancia entre la distribución actual del estudiante y un punto de control reciente ($\theta_{ckpt}$). Esta distancia se mide utilizando la distancia de Fisher–Rao (la longitud de arco en la hiperesfera), que rastrea el desplazamiento acumulado en la geometría de las distribuciones predictivas.
• Efecto: Esto actúa como una "región de confianza móvil", amortiguando desplazamientos abruptos mientras permite un movimiento constante cuando la señal del maestro es consistente. A diferencia de anclarse al modelo inicial (lo que prohibiría la adaptación), anclarse a un punto de control reciente permite que el modelo se adapte mientras evita que se desvíe demasiado de su comportamiento estable reciente.

C. Actualización de Gradiente Natural

Tanto la pérdida de Hellinger como el término proximal de Fisher–Rao son distancias en la misma geometría de las distribuciones de siguiente token. Para optimizar este objetivo, GEOSD utiliza una actualización de gradiente natural.

• Mecanismo: En lugar de actualizar los parámetros en el espacio euclidiano, la actualización está precondicionada por la Matriz de Información de Fisher (aproximada mediante K-FAC para mayor eficiencia).
• Efecto: Esto asegura que cada paso realice el cambio más pequeño en las predicciones del estudiante que mejore el objetivo, en lugar de realizar un paso de tamaño fijo en el espacio de parámetros. El término proximal actúa efectivamente como un "retroceso al punto de control" en esta geometría, devolviendo la política hacia el punto de control basándose en la saliencia conductual de las direcciones de los parámetros.

3. Contribuciones Clave

1. Objetivo GEOSD: Un marco de auto-destilación que pondera las preferencias del maestro por la superposición entre maestro y estudiante (vía pérdida de Hellinger) y penaliza la deriva predictiva de un punto de control reciente (vía distancia de Fisher–Rao), optimizado mediante gradientes naturales en una geometría unificada.
2. Mejora Empírica: GEOSD preserva las ganancias de distribución interna (ID) de la auto-destilación mientras mejora significativamente el razonamiento fuera de la distribución (OOD). A través de tres familias de modelos (Qwen3, Olmo, DeepSeek-R1) y cinco escalas (de 1.7B a 32B), mejora la precisión OOD promedio en 5.7–8.6 puntos sobre el modelo base, mientras que los objetivos KL estándar degradan el rendimiento OOD.
3. Análisis Mecanicista: El artículo caracteriza por qué el emparejamiento estándar falla en OOD: concentra la masa de probabilidad en la opción principal del maestro al "drenar" las alternativas en estados de alta entropía, lo que lleva a un acuerdo confiado en respuestas incorrectas. GEOSD mantiene la dominancia de la opción principal mientras mantiene las continuaciones alternativas "al alcance", evitando el colapso prematuro.

4. Resultados Experimentales

Los autores evaluaron GEOSD en tareas de razonamiento matemático (DAPO-Math para entrenamiento; AIME 2024/2025, AMC 2023, MATH-500 para pruebas OOD).

• Rendimiento: GEOSD logró la mejor precisión OOD en todos los benchmarks y familias de modelos probados. Por ejemplo, en Qwen3-8B, mejoró la precisión OOD en +8.6 puntos respecto al modelo base, comparado con una degradación de -8.1 puntos para la KL directa (Forward KL).
• Robustez: Las ganancias se mantuvieron en diversas cantidades de contexto privilegiado (desde prefijos cortos hasta soluciones completas) y en diversas escalas de modelos (1.7B a 32B).
• Ablación: Eliminar la ponderación de superposición de Hellinger, el término proximal de Fisher–Rao o la actualización de gradiente natural resultó en caídas significativas en el rendimiento OOD, confirmando que los tres componentes son necesarios.
• Falso Consenso: El análisis mostró que los métodos estándar (como la KL directa) exhiben altas tasas de "falso consenso" (muestras que acuerdan fuertemente en respuestas erróneas) en estados de alta entropía. GEOSD redujo esta tasa incluso por debajo del modelo base, manteniendo una autocoherencia informativa.

5. Significado y Reivindicaciones

El artículo sostiene que la degradación del razonamiento OOD en la auto-destilación de contexto privilegiado no es simplemente el resultado de usar señales poco fiables, sino una consecuencia de cómo esas señales mueven al modelo. Los objetivos estándar permiten que la confianza del maestro fuerce al estudiante hacia un modo estrecho prematuramente, drenando la masa de probabilidad de las alternativas viables.

GEOSD demuestra que la destilación de razonamiento efectiva no depende solo de qué señales del maestro confiar, sino de qué tan lejos se permite que esas señales muevan al estudiante. Al regular este movimiento geométricamente —atenuando los tirones donde el estudiante carece de apoyo y anclando la trayectoria a un estado estable reciente— GEOSD logra un equilibrio superior entre el rendimiento de distribución interna y la generalización fuera de la distribución. Los autores concluyen que controlar la deriva predictiva es esencial para escalar la auto-destilación en tareas de razonamiento.

Limitaciones Notadas: Los autores reconocen que sus hallazgos se han establecido actualmente solo para el razonamiento matemático. La aplicabilidad de la transferencia ponderada por superposición y el control de deriva de Fisher–Rao a otros dominios (por ejemplo, código, razonamiento científico) donde los estados de alta entropía pueden tener estructuras diferentes sigue sin probarse. Además, el método asume que la política base es ampliamente competente; para modelos base muy débiles, el mecanismo podría suprimir señales útiles que requieren cambios de comportamiento de gran escala.