Resumen Técnico: OISD – Auto-distilación Interna en Política de Modelos de Lenguaje

1. Declaración del Problema

Los enfoques recientes de aprendizaje por refuerzo (RL) para el post-entrenamiento de Modelos de Lenguaje de Gran Escala (LLM), particularmente aquellos que utilizan la Optimización de Política Relativa por Grupos (GRPO) con recompensas verificables, han avanzado significativamente el razonamiento matemático. Sin embargo, estos métodos son predominantemente centrados en la salida final. Optimizan la política de actuación final utilizando recompensas escasas a nivel de resultado (por ejemplo, la corrección de la respuesta final), mientras que pasan por alto en gran medida las señales predictivas, los procesos de razonamiento y las representaciones internas codificadas en las capas intermedias.

Aunque la Distilación de Conocimiento (KD) ofrece un mecanismo para proporcionar supervisión densa a nivel de token, los paradigmas existentes sufren limitaciones específicas:

• Distilación Fuera de Política: Crea una discrepancia de distribución porque el estudiante se entrena con trayectorias generadas por un profesor fijo, no con sus propias generaciones.
• Distilación En Política: Reduce la discrepancia, pero típicamente depende de modelos profesores externos.
• Auto-distilación En Política: Los métodos recientes eliminan los profesores externos, pero dependen de información privilegiada inyectada externamente (por ejemplo, prompts privilegiados, soluciones de verdad fundamental o retroalimentación auxiliar) para supervisar el modelo.

Existe una brecha en aprovechar los propios cálculos internos del modelo como fuente de supervisión sin requerir datos privilegiados externos ni múltiples políticas de actuación. Además, investigaciones previas indican que las representaciones de razonamiento evolucionan sistemáticamente a través de las capas (desde predicciones candidatas amplias en las primeras capas hasta decisiones convergidas en las capas profundas), lo que sugiere que las capas intermedias podrían ser guiadas por las capas más profundas para mejorar la coherencia del razonamiento.

2. Metodología: Marco OISD

Los autores proponen OISD (Auto-distilación Interna En Política), un marco que trata la capa final del modelo como un "profesor interno desconectado" y una capa intermedia seleccionada como un "estudiante interno auxiliar". Esto ocurre dentro de un solo pase forward durante el proceso de despliegue (rollout) y optimización GRPO.

Arquitectura Central

• Política de Actuación Única: La capa final ($L$) permanece como la única política para generar despliegues y optimizar el objetivo GRPO.
• Profesor-Alumno Interno: Durante el mismo pase forward, el estado oculto de una capa intermedia seleccionada ($l < L$) se proyecta al espacio de tokens utilizando la matriz de desincrustación (siguiendo el marco de la lente de logits). Esto crea una distribución provisional de "alumno" derivada del mismo contexto que el "profesor" (capa final).
• Sin Privilegio Externo: Las señales de supervisión se derivan enteramente de los propios estados internos del modelo en el mismo despliegue en política, sin requerir soluciones de verdad fundamental externas ni prompts privilegiados.

Objetivos de Alineación

OISD emplea dos mecanismos de alineación interna complementarios para transferir señales predictivas desde la capa final hacia la capa intermedia:

1. Alineación de Logits ("Cómo Pensar"):

   • Objetivo: Transferir comportamientos de razonamiento de alto nivel y creencias predictivas.
   • Mecanismo: Se calcula la divergencia de Jensen–Shannon (JS) entre las distribuciones de tokens escaladas por temperatura de la capa intermedia ($p^l_t$) y la capa final ($q^L_t$).
   • Ponderación: La pérdida se pondera por la ventaja firmada de GRPO ($A_t$). Los tokens de alta ventaja atraen a la capa intermedia hacia la distribución de la capa final, mientras que los tokens de baja ventaja la empujan hacia fuera para evitar reforzar continuaciones pobres.
   • Fórmula: $\mathcal{L}^{(l)}_{\text{think}} = \mathbb{E}_t [ A_t \cdot \text{JS}(p^{l,\tau}_t, \text{sg}(q^{L,\tau}_t)) ]$

2. Alineación de Atención ("Dónde Mirar"):

   • Objetivo: Imponer patrones de atención consistentes y comportamientos de enrutamiento de evidencia.
   • Mecanismo: Se calcula la divergencia de JS entre los pesos de atención de la capa intermedia y la capa final sobre un conjunto muestreado de claves causales (combinando ventanas locales y posiciones escaladas globales para gestionar la memoria).
   • Ponderación: Similar a la alineación de logits, la pérdida se pondera por la ventaja firmada. Las ventajas positivas incentivan a la capa intermedia a atender a la misma evidencia que la capa final; las ventajas negativas reducen el acuerdo con patrones de baja recompensa.
   • Fórmula: $\mathcal{L}^{(l)}_{\text{attn}} = \mathbb{E}_t [ A_t \cdot \frac{1}{H} \sum_{h=1}^H \text{JS}(\tilde{a}^l_{t,h}, \text{sg}(\tilde{a}^L_{t,h})) ]$

Objetivo de Entrenamiento

La pérdida total combina la pérdida GRPO estándar con las dos pérdidas de alineación interna:
$$\mathcal{L}_{\text{OISD}} = \mathcal{L}_{\text{GRPO}} + \lambda_{\text{think}} \mathcal{L}_{\text{think}} + \lambda_{\text{attn}} \mathcal{L}_{\text{attn}}$$
donde $\lambda_{\text{think}}$ y $\lambda_{\text{attn}}$ controlan la fuerza de la distilación.

3. Contribuciones Clave

1. Nuevo Paradigma: Introducción de la Auto-distilación Interna En Política, un paradigma de RL de razonamiento que aprovecha los propios cálculos internos del modelo para la supervisión mientras mantiene una única política de actuación.
2. Marco OISD: Una implementación específica que utiliza la capa final como profesor desconectado y una capa intermedia como alumno, transfiriendo señales predictivas en política sin información privilegiada externa.
3. Mecanismos de Alineación Dual: Desarrollo de dos objetivos complementarios:
   • Alineación de Logits: Supervisa el proceso de razonamiento ("cómo pensar").
   • Alineación de Atención: Supervisa el enrutamiento de evidencia ("dónde mirar").
4. Validación Empírica: Demostración de mejoras sustanciales y consistentes sobre bases de referencia fuertes de RL de razonamiento (incluyendo GRPO y BuPO) en cuatro benchmarks de razonamiento matemático.

4. Resultados Experimentales

Los autores evaluaron OISD en los modelos Qwen3-4B, Qwen3-8B y OctoThinker en cuatro benchmarks: AMC23, MATH500, AIME24 y AIME25.

• Mejoras de Rendimiento:
  • En Qwen3-4B, OISD logró una puntuación promedio de 64.75, superando a GRPO (55.08) por 9.67 puntos y a BuPO (58.51) por 6.24 puntos.
  • Las mejoras fueron particularmente significativas en los benchmarks AIME (por ejemplo, +14.31 puntos sobre GRPO en AIME24).
  • En Qwen3-8B, OISD logró el mejor promedio general de 66.72, superando a GRPO (64.23) y a BuPO (66.36).
• Análisis Pass@K: OISD logró consistentemente tasas Pass@K más altas para todos los valores de $K$ en comparación con GRPO, lo que indica una mejora en la calidad general del conjunto de soluciones muestreadas, no solo en la precisión promedio.
• Estudios de Ablación:
  • Tanto la alineación de Logits como la de Atención contribuyeron positivamente. La alineación de Logits mostró la mayor ganancia en AIME24, mientras que la alineación de Atención tuvo un rendimiento ligeramente mejor en promedio.
  • La combinación de ambas produjo los mejores resultados, confirmando su naturaleza complementaria.
• Comparación con OPSD: OISD (64.75) superó significativamente a OPSD (56.63), un método de auto-distilación en política de última generación que depende de soluciones de verdad fundamental privilegiadas. Esto resalta la efectividad de la supervisión puramente interna de OISD.
• Análisis Cualitativo:
  • Dinámicas de Entrenamiento: OISD mostró una mejora constante en la recompensa y un aumento inicial en la entropía (sugiriendo una exploración más amplia) antes de estabilizarse.
  • Visualización de Atención: Las capas intermedias en OISD se alinearon gradualmente con la atención de la capa final hacia los tokens de evidencia crítica (por ejemplo, puntos finales, supremo) durante el entrenamiento.
  • Visualización Lente de Logits: OISD exhibió una alineación más temprana entre las predicciones intermedias y finales en tokens críticos para el razonamiento en comparación con GRPO, que a menudo producía predicciones intermedias fragmentadas.

5. Significado y Afirmaciones

El artículo afirma que OISD inicia la investigación de la auto-distilación interna en política. Su significado principal radica en:

• Eliminación de la Dependencia Externa: Proporciona supervisión interna densa sin requerir profesores externos, prompts privilegiados o soluciones de verdad fundamental durante el entrenamiento.
• Consistencia de la Política: Preserva una única política de actuación tanto para la generación de despliegues como para la optimización, evitando los problemas de discrepancia de distribución de los métodos fuera de política.
• Razonamiento Mejorado: Al alinear las representaciones intermedias con los patrones de razonamiento y atención de la capa final, OISD incentiva al modelo a aprender representaciones intermedias más fuertes y globalmente consistentes.

Los autores señalan limitaciones, incluido el enfoque en el razonamiento matemático con respuestas verificables y la falta de un estudio en profundidad sobre la selección óptima de capas. Sugieren que el trabajo futuro debería explorar la generalización a tareas de razonamiento más complejas y la supervisión interna de múltiples capas.