Patch the Distribution Mismatch: RL Rewriting Agent for Stable Off-Policy SFT

1. El Problema: Desajuste de Distribución y Olvido Catastrófico

El artículo aborda un desafío fundamental en la adaptación de Modelos de Lenguaje Grandes (LLMs) a escenarios específicos (dominios downstream): el olvido catastrófico durante el Ajuste Fino Supervisado (SFT).

• Causa Raíz: Cuando los datos de downstream presentan un cambio de distribución significativo respecto a la distribución de entrenamiento previa del modelo, el SFT actúa como un aprendizaje "off-policy". Las demostraciones expertas pueden tener una probabilidad muy baja bajo la política actual del modelo, lo que induce inestabilidad en el entrenamiento y fuerza al modelo a sesgarse hacia la nueva distribución, perdiendo capacidades generales adquiridas anteriormente.
• Limitaciones de los Métodos Actuales: Se han propuesto enfoques centrados en los datos, como la reescritura de datos, para cerrar esta brecha antes del SFT. Sin embargo, los métodos existentes (como SDFT o "Mind the Gap") sufren de dos deficiencias clave:
  1. Desalineación Distribucional: Reescriben los datos condicionados a un prompt específico ($\pi_0(\cdot | x, y^\star, x_{prompt})$), lo que no garantiza que los resultados se alineen con la distribución natural de generación de preguntas y respuestas (QA) del modelo ($\pi_0(\cdot | x)$).
  2. Colapso de Diversidad: El uso de plantillas fijas y prompts rígidos tiende a producir reescrituras templadas y homogéneas, reduciendo la diversidad de las soluciones y exacerbando el colapso de modos.

2. Metodología: Agente de Reescritura Basado en RL

Los autores proponen un marco novedoso que trata la reescritura de datos como un problema de aprendizaje de políticas, utilizando Aprendizaje por Refuerzo (RL) para optimizar la distribución de los datos reescritos.

Arquitectura del Agente

• Base Congelada: Se utiliza un modelo base instruido y congelado ($\pi_0$) como columna vertebral.
• Adaptación Eficiente (LoRA): Se entrena una política de reescritura ligera ($R_\phi$) mediante adaptadores LoRA sobre el modelo congelado. Esto actúa como un "parche" controlado que corrige localmente el comportamiento de generación sin desviarse excesivamente de la distribución original.
• Optimización GRPO: Se emplea Group Relative Policy Optimization (GRPO), un método de gradiente de política que elimina la necesidad de una red de valor separada, estimando la línea base a partir de grupos de salidas muestreadas.

Función de Objetivo Unificada (Recompensas)

El agente se entrena para maximizar una recompensa compuesta que equilibra tres objetivos, utilizando un mecanismo de puerta dura (hard gating) para la consistencia de la tarea:

1. Consistencia de la Tarea (Puerta Dura):

   • Actúa como un filtro binario ($r_{task} \in \{0, 1\}$).
   • Se verifica primero la corrección de la respuesta final (con un verificador basado en reglas) y luego la validez del razonamiento (con un juez LLM).
   • Si la reescritura falla esta verificación, la recompensa total es cero y no se actualiza la política con señales auxiliares.

2. Alineación Distribucional ($r_{dist}$):

   • Busca que las reescrituras sean más probables bajo la distribución natural de QA del modelo base ($\pi_0(\cdot | x)$), sin el condicionamiento del prompt de reescritura.
   • Se calcula como la verosimilitud negativa normalizada (NLL) bajo el modelo base congelado.

3. Diversidad ($r_{div}$):

   • Actúa como regularizador contra el colapso de modos.
   • Se mide la distancia semántica (usando embeddings) entre las reescrituras factibles dentro de un grupo. Se recompensa a las reescrituras que aportan contenido semántico novedoso al conjunto.

Construcción del Dataset

El proceso sigue un pipeline GENERAR-VERIFICAR-RECUPERAR (FALLBACK):

1. El agente genera una reescritura $\tilde{y}$.
2. Se verifica la consistencia de la tarea.
3. Si pasa, se usa $\tilde{y}$ como objetivo de supervisión. Si falla, se recupera la demostración experta original $y^\star$ para no perder supervisión.

3. Contribuciones Clave

1. Insight sobre el Desajuste de Distribución: Identifican que la reescritura basada en prompts induce una distribución condicional artificial que no coincide con la generación natural de QA, y proponen alinear explícitamente la reescritura con esta distribución natural.
2. Objetivo Unificado con RL: Introducen un marco que optimiza simultáneamente la consistencia de la tarea, la alineación distribucional y la diversidad, superando las limitaciones de los métodos heurísticos basados en plantillas.
3. Eficiencia y Estabilidad: Demuestran que un agente de reescritura ligero (LoRA) entrenado con RL puede generar datasets de alta calidad que estabilizan el SFT downstream.

4. Resultados Experimentales

Los experimentos se realizaron en modelos de diferentes escalas (Llama-3.2-1B, 3B y Mistral-7B) enfocándose en razonamiento matemático como tarea downstream.

• Rendimiento Downstream: El método propuesto ("Rewriting-agent") logra mejoras en matemáticas comparables o ligeramente superiores al SFT estándar (Vanilla SFT).
• Reducción del Olvido Catastrófico: La contribución más significativa es la preservación de capacidades generales.
  • En Mistral-7B, mientras el SFT estándar reduce el rendimiento en benchmarks generales en un 19.17%, el método propuesto solo reduce un 7.35%.
  • En promedio, se reduce el olvido en benchmarks no relacionados con matemáticas en un 12.34% en comparación con el SFT estándar.
• Análisis de Pérdida: El entrenamiento SFT sobre los datos reescritos por el agente comienza con una pérdida más baja y converge a un plato más estable en comparación con SFT estándar o reescrituras basadas en plantillas, indicando una mejor alineación con la distribución del modelo.
• Ablaciones:
  • Eliminar la recompensa de alineación distribucional ($r_{dist}$) causa la mayor caída en la retención general, confirmando que la alineación con $\pi_0(\cdot | x)$ es crucial.
  • El uso de una puerta dura (hard gate) es superior al "shaping" suave, ya que evita que señales ruidosas de reescrituras incorrectas degraden la política.

5. Significado e Impacto

Este trabajo ofrece una solución elegante al dilema entre plasticidad (aprender nuevas tareas) y estabilidad (retener conocimientos previos) en el ajuste fino de LLMs.

• Paradigma Data-Centric: En lugar de restringir las actualizaciones del modelo (como hacen métodos de "trust region" o "clipping"), el enfoque actúa en la fuente de los datos, transformando las demostraciones expertas para que sean "más fáciles de aprender" para el modelo sin sacrificar la calidad de la tarea.
• Escalabilidad: Al utilizar LoRA y RL, el método es computacionalmente eficiente y no requiere reentrenar el modelo base completo.
• Generalización: La capacidad de generar datos que respetan la distribución natural del modelo sugiere que este enfoque podría ser aplicable a otros dominios más allá de las matemáticas, mitigando el riesgo de que los modelos se vuelvan rígidos o pierdan su "personalidad" y capacidades generales al especializarse.

En resumen, el artículo demuestra que aprender a reescribir la supervisión mediante un agente de RL es una estrategia superior para estabilizar el SFT off-policy, logrando un equilibrio óptimo entre el rendimiento en la tarea específica y la retención de capacidades generales.