Resumen Técnico: Hacia una Internalización Robusta en la Destilación On-Policy

1. Planteamiento del Problema

Los avances recientes en la destilación on-policy (OPD) han permitido que los modelos estudiante internalicen el contexto privilegiado (por ejemplo, prompts de sistema, pistas de tareas o descripciones de estados de juego) en sus parámetros, permitiéndoles desempeñarse bien sin dicho contexto en el momento de la inferencia. Sin embargo, este artículo identifica un fenómeno previamente no estudiado denominado degradación inducida por el contexto.

Contrario a la intuición de que el conocimiento internalizado debería hacer que el contexto sea redundante (una operación idempotente), los autores observan que reintroducir el contexto privilegiado a un estudiante destilado a menudo degrada su rendimiento e infla la longitud de las respuestas, incluso en instancias que el modelo resuelve correctamente sin contexto. Esto ocurre porque la OPD estándar optimiza exclusivamente la fidelidad privilegiada (coincidir con la salida del profesor condicionada al contexto), pero no impone restricciones sobre el comportamiento del estudiante cuando el contexto está presente. En consecuencia, el estudiante no logra alcanzar la removibilidad de contexto: la propiedad donde la salida del modelo permanece estable independientemente de si se proporciona o no el contexto privilegiado.

Los autores categorizan las interacciones de contexto post-destilación en tres regímenes:

• Régimen A: El estudiante funciona bien sin contexto, pero la reintroducción del contexto causa caídas significativas en la precisión y altas "tasas de daño" (cambiar predicciones correctas por incorrectas). Este es el régimen más prevalente en sus experimentos.
• Régimen B: El contexto sigue actuando como un andamio útil, lo que indica una internalización incompleta.
• Régimen C: El estudiante ha internalizado con éxito el contexto, y el rendimiento es estable con o sin él.

2. Metodología: Anclaje Sin Contexto (NCA)

Para abordar la degradación inducida por el contexto, los autores proponen el Anclaje Sin Contexto (NCA), un regularizador de consistencia ligero. La idea central es forzar que la salida del estudiante condicionada al contexto ($q_c$) sea consistente con su salida sin contexto ($q_x$), tratando a esta última como un ancla.

Función Objetivo

La OPD estándar minimiza la divergencia KL entre la salida sin contexto del estudiante y la distribución del profesor condicionada al contexto:
$$\mathcal{L}_{OPD} = \mathbb{E}_{x \sim D} [\text{KL}(q_x \parallel f(y \mid x, c))]$$

NCA aumenta esto con un término de consistencia:
$$\mathcal{L} = \mathcal{L}_{OPD} + \beta \cdot \mathcal{L}_{NCA}$$
donde $\mathcal{L}_{NCA}$ se define como:
$$\mathcal{L}_{NCA} \triangleq \text{KL}(\text{sg}[g_\theta(\cdot \mid x)] \parallel g_\theta(\cdot \mid x, c))$$

Decisiones de Diseño Clave

1. Anclaje de Parada de Gradiente (Stop-Gradient): La vista sin contexto $q_x$ pasa a través de un operador de parada de gradiente (sg). Esto asegura que $q_x$ reciba gradientes solo de la pérdida primaria de OPD, mientras que $q_c$ recibe gradientes solo de la pérdida de NCA. Esto elimina la competencia de gradientes entre ambos objetivos.
2. Divergencia KL hacia Adelante (Forward KL): Los autores eligen $\text{KL}(q_x \parallel q_c)$ (KL hacia adelante) en lugar de la KL inversa. Esta elección es de tipo "cobertura de modos" (mode-covering), lo que evita que $q_c$ colapse los modos de $q_x$. Además, la esperanza cae bajo $q_x$, lo que permite la reutilización de las muestras on-policy generadas para la OPD sin requerir un nuevo muestreo de $q_c$.
3. Eficiencia Computacional: NCA requiere solo una pasada adicional hacia adelante por cada paso de entrenamiento (para calcular $q_c$). Los logits de $q_x$ se desprenden del grafo de computación y las secuencias on-policy se reutilizan, lo que hace que la sobrecarga sea insignificante.

3. Configuración Experimental

Los autores evaluaron NCA a través de 12 configuraciones que abarcan:

• Dominios: QA de prompts de sistema (MedMCQA, conjuntos de datos de seguridad) y juegos de texto (Sokoban, FrozenLake).
• Modelos: Varios modelos Llama-3 y Qwen (que oscilan entre 1.7B y 8B de parámetros), incluyendo escenarios de autodestilación y destilación entre modelos distintos.
• Métricas: Precisión/Tasa de victoria con y sin contexto, Tasa de daño (probabilidad de que el contexto cambie una respuesta correcta a incorrecta), Tasas de "Ambos-Correctos/Ambos-Ganados", y cambios relativos en la longitud de la respuesta.

4. Resultados Clave

A través de las 12 configuraciones, NCA demostró los siguientes efectos:

• Mitigación de la Degradación Inducida por el Contexto: NCA redujo el daño inducido por el contexto en 11 de las 12 configuraciones. En los entornos de Régimen A (donde la degradación era severa), NCA redujo significamente la caída de precisión tras la reintroducción del contexto.
• Mejora de la Estabilidad: NCA aumentó consistentemente las tasas de "AmbosCorrectos" o "AmbosGanados", indicando que el modelo mantiene un comportamiento correcto bajo ambas vistas.
• Eliminación de la Inflación de la Longitud de Respuesta: La OPD estándar a menudo causaba que el modelo generara respuestas significativamente más largas cuando se reintroducía el contexto (hasta un aumento de longitud del +39.5%). NCA eliminó eficazmente esta inflación en la mayoría de los casos, reduciendo el cambio de longitud a casi cero.
• Mejora del Rendimiento Sin Contexto: En 7 de las 12 configuraciones, NCA mejoró el objetivo principal de entrenamiento (precisión sin contexto) entre 0.6 y 3.6 puntos porcentuales. Los autores atribuyen esto a que la restricción de consistencia actúa como un regularizador que evita que las actualizaciones de parámetros acomoden comportamientos conflictivos entre las vistas.
• Casos Límite: En un escenario de Régimen B (FrozenLake, destilación entre modelos), NCA aumentó ligeramente el daño, sugiriendo que cuando el contexto aún porta información única y no internalizada, forzar una removibilidad estricta puede sacrificar parte del rendimiento condicionado al contexto en favor de la estabilidad.

5. Análisis Mecanicista

El artículo proporciona un estudio de caso mecanicista sobre la tarea de Sokoban para explicar cómo el NCA logra la removibilidad de contexto:

• Nivel de Representación: El análisis de los estados ocultos reveló que, si bien la OPD estándar redujo la divergencia entre las vistas de contexto y sin contexto, persistía una brecha notable en las capas medias y tardías. Sin embargo, el NCA mantuvo una similitud de coseno casi perfecta ($\ge 0.997$) en todas las capas, confirmando que la removibilidad de contexto se logra a nivel de representación.
• Alineación de Gradientes: En la OPD estándar, los gradientes para el objetivo de destilación y la vista condicionada al contexto eran casi ortogonales. El NCA introduce una nueva dirección de gradiente que equilibra activamente las fuerzas de aprovechar el contexto para la destilación y permanecer invariante ante él.
• Actualizaciones de Parámetros: Tanto la OPD como el NCA concentran las actualizaciones de parámetros en las capas superficiales (particularmente en $k\_proj$ y $q\_proj$). Sin embargo, las actualizaciones específicas impulsadas por el NCA son distintas de las de la OPD, lo que sugiere que la removibilidad de contexto se logra principalmente mediante modificaciones en las representaciones de conocimiento de las capas superficiales que la OPD por sí sola no produce.

6. Significancia y Reivindicaciones

Los autores afirman que su trabajo cambia el paradigma de evaluación para la destilación de contexto privilegiado. Argumentan que una internalización robusta requiere no solo coincidir con el comportamiento del profesor (fidelidad privilegiada), sino también asegurar la removibilidad de contexto.

• Novedad: Este es el primer estudio que identifica y aborda el fenómeno donde la reintroducción del contexto internalizado degrada el rendimiento del modelo.
• Practicidad: El método propuesto (NCA) es simple, requiere solo una pasada hacia adelante y un regularizador de consistencia estándar, pero resuelve eficazmente el problema de la degradación sin sacrificar los objetivos primarios de la destilación.
• Generalizabilidad: El método mejora el rendimiento a través de diversas familias de modelos, tareas y configuraciones de destilación (auto y entre modelos).

El artículo concluye instando a la comunidad a evaluar los modelos destilados no solo por su rendimiento sin contexto, sino también por su robustez ante la reintroducción del contexto, asegurando que el conocimiento internalizado no se convierta en una carga cuando el contexto se proporciona accidental o intencionadamente.