Post-Training is About States, Not Tokens: A State Distribution View of SFT, RL, and On-Policy Distillation

Resumen Técnico: Entrenamiento Posterior como Moldeado de la Distribución de Estados

Enunciado del Problema

Los análisis actuales de los métodos de entrenamiento posterior de Modelos de Lenguaje de Gran Escala (LLM)—específicamente el Ajuste Fino Supervisado (SFT), el Aprendizaje por Refuerzo (RL) y la Distilación—se centran principalmente en los objetivos de optimización (por ejemplo, máxima verosimilitud, gradientes de política, divergencia KL). Sin embargo, esta visión centrada en el objetivo tiene dificultades para explicar varios fenómenos empíricos:

• Por qué un SFT moderado puede mejorar el rendimiento con un olvido mínimo, mientras que un SFT agresivo ("de estrés") provoca un olvido catastrófico y un rendimiento degradado en la tarea objetivo.
• Por qué el RL, a pesar de utilizar recompensas dispersas o ruidosas, a menudo produce mejoras estables.
• Por qué un modelo estudiante en la distilación puede superar a veces a un profesor degradado, incluso cuando el profesor es la única fuente de supervisión.

El artículo argumenta que estas discrepancias surgen porque los análisis estándar pasan por alto un eje crítico de variación: la distribución de estados sobre la cual se aplica la supervisión. En los modelos autoregresivos, un "estado" se define como el prompt más el prefijo generado ($s_t = (x, y_{<t})$). Los diferentes métodos de entrenamiento posterior aplican sus señales a distribuciones de estados fundamentalmente distintas (prefijos de conjuntos de datos fijos frente a trayectorias inducidas por el aprendiz), lo que impacta significativamente el comportamiento del modelo, la localidad y la retención.

Metodología

Los autores proponen un marco unificado que concibe el entrenamiento posterior como moldeado de la distribución de estados. Descomponen cualquier método de entrenamiento posterior en dos ejes distintos:

1. Fuente de Estado ($q_k$): De dónde provienen los contextos de entrenamiento (estados).
2. Fuente de Señal ($S$): Qué supervisión (tokens, recompensas, logits) se proporciona en esos estados.

Formulación Teórica:

• SFT: Ajuste de estados fuera de política. El modelo minimiza la pérdida de tokens en estados de conjuntos de datos fijos ($d_{data}$). Esto crea una discrepancia si las trayectorias del conjunto de datos se desvían de las propias trayectorias del modelo, lo que conduce a un sesgo de exposición y una posible inestabilidad bajo estrés.
• RL: Mejora local en política. El modelo muestrea trayectorias desde su política actual ($d_{\pi_\theta}$) y actualiza basándose en recompensas. Las actualizaciones se localizan naturalmente en los estados que el modelo visita realmente.
• Distilación en Política (OPD): Un híbrido donde el estudiante muestrea estados ($d_{\pi_S}$) pero el profesor proporciona la señal (logits o continuaciones). Esto desacopla la fuente de estados de la fuente de señal, permitiendo que el estudiante aprenda reparaciones locales sin heredar la distribución de trayectorias global del profesor.

Configuración Experimental:

• Modelo: Qwen3-0.6B-Base con adaptadores LoRA.
• Hardware: Una sola GPU RTX 3090.
• Tareas:
  • Objetivo: GSM8K (razonamiento matemático).
  • Retención: TruthfulQA y MMLU (para medir el olvido catastrófico).
• Métricas: Precisión en la tarea objetivo, ratio de retención, magnitud del olvido y deriva de estado (medida mediante la Discrepancia Media Máxima, MMD, en características léxicas de los estados de las trayectorias).

Condiciones Experimentales:

1. SFT Moderado: Entrenamiento estándar en GSM8K.
2. SFT de Estrés: Entrenamiento agresivo (5 épocas, alta tasa de aprendizaje, alto rango LoRA) para inducir degradación.
3. OPD: Entrenar a un estudiante utilizando un profesor degradado de "SFT de Estrés" como fuente de señal, pero muestreando estados desde la propia política del estudiante.
4. RL en Política: RL ligero utilizando recompensas exactas relativas a grupos en GSM8K.

Resultados Clave

1. El SFT es Dependiente del Contexto: El SFT moderado mejoró GSM8K (0.448 $\to$ 0.512) con una pérdida de retención insignificante. Por el contrario, el SFT de estrés degradó tanto la tarea objetivo (0.420) como las tareas de retención (TruthfulQA/MMLU cayeron significativamente), demostrando que el SFT no es inherentemente destructivo, sino que se vuelve tal bajo presión fuera de política que no se alinea con la distribución de estados natural del modelo.
2. OPD Supera a Profesores Degradados: Un estudiante entrenado mediante OPD utilizando el modelo "SFT de Estrés" como profesor superó al profesor en todas las métricas (GSM8K: 0.466 vs. 0.420; MMLU: 0.430 vs. 0.364). Esto confirma que los estudiantes no heredan los fallos globales de trayectoria del profesor si son entrenados sobre sus propios estados, recibiendo solo guía condicional local.
3. El RL Preserva la Localidad: El RL ligero en política mejoró GSM8K (0.472) mientras mantenía niveles de retención cercanos a la base. Esto respalda la visión de que el RL actúa como un procedimiento de mejora local fundamentado en los estados alcanzables por el aprendiz.
4. Las Métricas de Deriva son Insuficientes: Los métodos con una deriva de estado escalar casi idéntica (MMD) exhibieron resultados de retención muy diferentes. Por ejemplo, el SFT de estrés y el OPD desde el profesor de estrés tuvieron valores MMD similares (~0.0109) pero ratios de retención muy diferentes (0.82 vs. 0.95). Esto indica que la fuente y la localidad de los estados de entrenamiento importan más que la magnitud del desplazamiento de distribución por sí sola.

Contribuciones

1. Marco Unificado: Formuló SFT, RL y OPD bajo una visión común de "distribución de estados", separando explícitamente la fuente de estados de la fuente de señal.
2. Pipeline Empírico Controlado: Implementó una configuración experimental de una sola GPU midiendo precisión objetivo, retención, olvido y deriva de estado de trayectorias para aislar los efectos de la distribución de estados.
3. Evidencia para Afirmaciones Centradas en el Estado: Proporcionó evidencia empírica de que:
   • OPD puede recuperar el rendimiento y superar a un profesor degradado al desacoplar el muestreo de estados de la generación de señales.
   • El RL en política mejora el rendimiento objetivo con una pérdida mínima de retención debido a su localidad.
   • Las métricas de deriva escalar (como MMD) no pueden caracterizar completamente la dinámica del entrenamiento posterior sin considerar la fuente de estados.

Significado y Afirmaciones

El artículo afirma que el comportamiento del entrenamiento posterior está gobernado por la interacción entre la señal de supervisión y la distribución de estados de entrenamiento. Los autores argumentan que la "fuente, localidad y dependencia del aprendiz de los estados de entrenamiento" son tan críticas como la forma de la señal de supervisión en sí misma.

• Reencuadre del SFT: El olvido catastrófico no es solo una función de la máxima verosimilitud, sino que surge cuando actualizaciones densas fuera de política empujan al modelo hacia regiones incompatibles con su propia distribución de estados futura.
• Reencuadre de la Distilación: Un estudiante puede superar a un profesor porque los errores del profesor a menudo están acoplados a la propia distribución de trayectorias del profesor. Al consultar al profesor sobre los estados del estudiante, este accede a reparaciones locales útiles sin heredar fallos globales de trayectoria.
• Reencuadre del RL: La estabilidad del RL se atribuye parcialmente a su naturaleza en política, asegurando que las actualizaciones se apliquen solo a estados que el modelo puede alcanzar y de los cuales puede recuperarse realmente.

El artículo concluye que es necesaria una visión centrada en el estado para explicar la estabilidad y las inversiones de rendimiento que el análisis a nivel de objetivo por sí solo no puede justificar. Los autores enfatizan que estos hallazgos se basan en un estudio a pequeña escala y deben considerarse como evidencia para una hipótesis mecanicista en lugar de una afirmación definitiva de referencia.