Resumen Técnico: Lightning OPD 2.0

Declaración del Problema

La destilación on-policy (OPD) ha emergido como una alternativa estable y rentable al aprendizaje por refuerzo para el post-entrenamiento de grandes modelos de razonamiento. Esta aprovecha un profesor de alta capacidad para proporcionar supervisión densa a nivel de token sobre las trayectorias generadas por el estudiante. Sin embargo, un prerrequisito crítico para una OPD efectiva, identificado en trabajos previos (Lightning OPD), es la consistencia del profesor: el modelo que proporciona la supervisión de la OPD debe ser el mismo modelo que generó las demostraciones utilizadas para la política de referencia del Ajuste Fino Supervisado (SFT) inicial.

En la práctica, esta condición se viola frecuentemente. Los datos de SFT suelen tener una procedencia mixta o no documentada, o los profesionales pueden desear utilizar un conjunto de datos de SFT específico generado por un modelo mientras destilan con un modelo profesor diferente y potencialmente más fuerte. Este escenario de profesor cruzado introduce un desajuste donde el desacuerdo bruto entre el profesor y la referencia contiene no solo señales de razonamiento útiles y específicas del contexto, sino también un componente recurrente asociado a diferencias estilísticas (redacción, formato, cadencia de razonamiento). Cuando no se corrige, este "sesgo de estilo-token" se acumula a través de los rollouts, interfiriendo sistemáticamente con la señal de razonamiento y causando que el rendimiento de la destilación se degrade, incluso cuando el nuevo profesor es más fuerte que el generador original de SFT.

Metodología: Lightning OPD 2.0

Para abordar esto, los autores introducen Lightning OPD 2.0, un método que mitiga el sesgo de estilo mediante la residualización de estilo mediante ajuste cruzado (cross-fitted style residualization). La idea central es que el desacuerdo relacionado con el estilo es predecible a través de los rollouts almacenados en caché basándose en coordenadas léxicas simples y de contexto grueso, mientras que el desacuerdo relacionado con el razonamiento depende más del estado específico del problema.

El método procede de la siguiente manera:

1. Descomposición del Desacuerdo: La diferencia bruta de log-probabilidad entre el profesor y la referencia ($d_{it} = \ell^T_{it} - \ell^R_{it}$) se modela como la suma de un componente de sesgo de estilo recurrente $b(z_{it})$ y una variación residual $v_{it}$ (que representa la evidencia específica del contexto).
   $$d_{it} = b(z_{it}) + v_{it}$$
2. Coordenadas Observables: Para estimar el componente recurrente $b(z_{it})$, el método utiliza dos tipos de coordenadas:
   • Identidad del Token: Captura preferencias léxicas.
   • Contexto Grueso: Definido por la posición de la respuesta normalizada y la sorpresa de la política de referencia (una medida de qué tan atípico es un token bajo el SFT de referencia).
3. Estimación mediante Ajuste Cruzado (Cross-Fitted): Los rollouts almacenados en caché se particionan en $K$ pliegues (folds). Para un pliegue mantenido fuera (held-out fold) $k$, el método construye dos tablas de búsqueda utilizando únicamente los otros $K-1$ pliegues:
   • Una tabla de búsqueda de tokens que estima el desacuerdo medio por identidad de token.
   • Una tabla de búsqueda de contexto que estima el desacuerdo medio por par de coordenadas de contexto.
   • Estas tablas se suavizan para manejar grupos poco frecuentes.
4. Residualización: Para cada token en el pliegue mantenido fuera, el sesgo de estilo estimado $\hat{b}_{it}$ se calcula como el promedio de las predicciones de las tablas de búsqueda de token y contexto. Este estimado se resta del desacuerdo bruto para producir una señal residualizada:
   $$d^*_{it} = d_{it} - \hat{b}_{it}$$
5. Optimización: El desacuerdo residualizado $d^*_{it}$ reemplaza al desacuerdo bruto en el objetivo de Lightning OPD. La política del estudiante se actualiza utilizando la ventaja derivada de esta señal residualizada, eliminando efectivamente el sesgo de estilo recurrente mientras se preserva la evidencia de razonamiento específica del contexto.

Contribuciones Clave

1. Formulación de la OPD de Profesor Cruzado: El artículo define formalmente el escenario donde el generador de datos de SFT y el profesor de OPD se seleccionan de forma independiente, demostrando que la aplicación ingenua de la OPD en este escenario conduce a una degradación sustancial del rendimiento.
2. Identificación del Sesgo de Estilo-Token: Los autores identifican un componente recurrente del desacuerdo entre profesor y referencia que es predecible a través de los rollouts y está asociado con diferencias estilísticas, explicando por qué la supervisión de profesor cruzado bruta puede ser engañosa.
3. Algoritmo Lightning OPD 2.0: La introducción de una técnica de residualización de estilo mediante ajuste cruzado que estima y resta este sesgo utilizando tablas de búsqueda de tokens y contexto. Esto permite que la referencia de SFT y el profesor de destilación se elijan independientemente sin reconstruir el pipeline de SFT.

Resultados Experimentales

El método fue evaluado en tareas de razonamiento matemático y generación de código utilizando dos escenarios distintos de profesor cruzado:

• Escenario 1: Referencia de SFT de Qwen3-4B (generado por Qwen3-8B) con un profesor de OPD de Qwen3-30B-A3B-Thinking-2507.
• Escenario 2: Referencia de SFT de Klear-Reasoner-8B (entrenado con datos destilados de DeepSeek-R1) con el mismo profesor de OPD Qwen3-30B.

Hallazgos Clave:

• Ganancias de Rendimiento: Lightning OPD 2.0 supera consistentemente al baseline Lightning OPD y a otras variantes de OPD (IW-OPD, TA-OPD) en escenarios de profesor cruzado.
  • En el escenario de Qwen3-4B, mejoró el razonamiento matemático promedio en 3.1 puntos y la generación de código en 1.4 puntos respecto al baseline Lightning OPD.
  • En el escenario de Klear-Reasoner-8B, mejoró el razonamiento matemático promedio en 1.0 punto y la generación de código en 1.4 puntos.
• Estado del Arte: Partiendo de Klear-Reasoner-8B-SFT, Lightning OPD 2.0 alcanzó un 82.4% en AIME 2024 y un 63.0% en LiveCodeBench v5.
• Análisis de Mecanismo: Un análisis post-hoc mostró que la residualización redujo significamente la desviación de la señal de entrenamiento respecto a una señal de referencia diagnóstica (que representa la señal ideal consistente con el profesor). Por ejemplo, en el escenario de Qwen3-4B, la tasa de tokens con grandes desviaciones (>1 nat) cayó del 8.14% al 3.85% tras la corrección.
• Ablación: Eliminar la búsqueda de tokens, la búsqueda de contexto o el ajuste cruzado resultó en un menor rendimiento, confirmando la necesidad de combinar las estimaciones léxicas y contextuales con la división de muestras para prevenir el auto-ajuste (self-fitting).

Significado y Reivindicaciones

El artículo afirma que Lightning OPD 2.0 establece un enfoque práctico para la destilación on-policy de profesor cruzado. Al relajar el requisito estricto de consistencia del profesor, permite a los profesionales:

• Seleccionar el modelo más adecuado para generar las demostraciones de SFT (para maximizar la calidad de la política de referencia) de forma independiente del modelo seleccionado para la destilación posterior (para maximizar el resultado final de post-entrenamiento).
• Utilizar conjuntos de datos de SFT existentes sin el costo sustancial de regenerar las demostraciones para cada candidato a profesor.

Los autores posicionan este trabajo como una solución a un cuello de botella práctico en el escalado de modelos de razonamiento, permitiendo pipelines de post-entrenamiento más flexibles y eficientes donde el generador de SFT y el profesor de OPD están desacoplados. El método se presenta como un marco robusto que logra un rendimiento de estado del arte en benchmarks de razonamiento mientras mantiene la estabilidad y la rentabilidad de la OPD offline.