Resumen Técnico: DISA - Muestreo de Importancia Offline para Ajuste de Distribución en RL de LLM

1. Declaración del Problema

Los agentes de razonamiento modernos y los Modelos de Lenguaje Grandes (LLM) se evalúan cada vez más por su capacidad para generar múltiples trayectorias de solución válidas para un solo prompt, en lugar de una única respuesta correcta. Los objetivos estándar de Aprendizaje por Refuerzo (RL), como PPO, GRPO y GSPO, se centran en la maximización de la recompensa. Aunque efectivos para la precisión de una sola respuesta, estos métodos tienden a sufrir colapso de modos, donde la política concentra la masa de probabilidad en la única trayectoria de alta recompensa más fácilmente reforzable, eliminando la diversidad de rutas de razonamiento válidas requeridas para métricas de evaluación robustas como pass@k.

Para abordar esto, el RL de ajuste de distribución busca adaptar una política a una distribución objetivo donde la probabilidad de una trayectoria es proporcional a su recompensa exponenciada (una distribución de Boltzmann). Sin embargo, este enfoque requiere calcular una función de partición ($Z(q)$) para cada prompt, que normaliza la distribución. Los métodos existentes de ajuste de distribución (por ejemplo, FlowRL) aprenden esta función de partición en línea junto con la política. Este acoplamiento crea dos problemas críticos:

1. Errores de Calibración: Los errores en la función de partición estimada distorsionan directamente las actualizaciones de la política.
2. Falta de Diagnósticos: No existe una señal de calidad independiente para la función de partición; su error está entrelazado con la pérdida de la política, haciendo imposible diagnosticar si un rendimiento deficiente se debe a la política o a la estimación de la partición.

2. Metodología: Anclaje de Muestreo de Importancia Desacoplado (DISA)

La idea central de DISA es que la función de partición $Z(q)$ depende únicamente de la política de referencia ($\pi_{ref}$), la distribución de prompts y la función de recompensa; es independiente de la política. Por lo tanto, puede estimarse offline antes de que comience la optimización de la política, desacoplando el problema de calibración del bucle de RL.

DISA opera en tres etapas distintas:

Etapa 1: Estimación Offline de Muestreo de Importancia (IS)

Para cada prompt $q$ en un conjunto de datos $\mathcal{D}_Z$, DISA extrae $N$ trayectorias propuestas de un modelo proponente fuerte $p_T$ (por ejemplo, un modelo docente más grande). Luego, calcula una estimación insesgada de la función de partición utilizando muestreo de importancia en el espacio logarítmico:
$$\log \hat{Z}_{IS}(q) = \text{logsumexp}_{i} \left( \log \pi_{ref}(o_i|q) - \log p_T(o_i|q) + \beta \tilde{r}(q, o_i) \right) - \log N$$
Este paso produce una etiqueta por prompt ruidosa pero insesgada para la función de partición.

Etapa 2: Amortización mediante Regresión

Para desruidar las etiquetas de IS y generalizar a prompts fuera del conjunto de entrenamiento, se entrena un regresor ligero $g_\psi$ (una pequeña MLP) mediante regresión de mínimos cuadrados para predecir $\log Z(q)$ a partir del prompt $q$:
$$\psi^* = \arg \min_\psi \mathbb{E}_{q \sim \mathcal{D}_Z} \left[ (g_\psi(q) - \log \hat{Z}_{IS}(q))^2 \right]$$
Una vez entrenado, los parámetros $\psi^*$ se congelan. Esto crea un "oráculo congelado" que actúa como una constante por prompt durante la fase posterior de RL.

Etapa 3: Aprendizaje de Política Online Anclada

En la fase online de RL, la política $\pi_\theta$ se optimiza utilizando la pérdida de Equilibrio de Trayectoria (TB), pero con una modificación crucial: la función de partición aprendida $Z_\phi(q)$ se reemplaza por la estimación congelada $g_{\psi^*}(q)$.
$$\mathcal{L}_{DISA}(\theta) = \mathbb{E}_{(q, \{o_i\})} \left[ \frac{1}{G} \sum_{i=1}^G \omega_i \left( g_{\psi^*}(q) + \log \pi_\theta(o_i|q) - \log \pi_{ref}(o_i|q) - \beta \tilde{r}(q, o_i) \right)^2 \right]$$
Críticamente, los gradientes fluyen solo a través de la política $\pi_\theta$. La función de partición actúa como un ancla fija (stop-gradient), separando estrictamente la estimación de la función de partición del aprendizaje de la política en datos, gradientes y diagnósticos.

3. Contribuciones Clave

1. Identificación del Acoplamiento como Cuello de Botella: El artículo identifica el acoplamiento en línea de la política y la función de partición en los métodos existentes de ajuste de distribución como un obstáculo central, señalando que inyecta el error de la función de partición directamente en el gradiente de la política sin una señal de calidad independiente.
2. Pipeline de DISA: Un pipeline de tres etapas que estima la función de partición offline, la amortiza en un regresor condicionado al prompt y congelado, y utiliza esta estimación congelada en el RL de equilibrio de trayectoria online.
3. Garantías Teóricas:
   • Preservación Exacta: Si la función de partición offline es exacta, el minimizador global del objetivo de DISA coincide con la distribución objetivo inclinada por la recompensa.
   • Estacionariedad: Incluso con un sesgo solo de prompt en la función de partición, la distribución objetivo permanece como un punto estacionario en política.
   • Perturbación Acotada: Los errores residuales de estimación offline entran en el objetivo de la política solo a través de un sobre de perturbación acotado (controlado por la desigualdad de Cauchy–Schwarz), asegurando estabilidad.
4. Validación Empírica: Experimentos extensos en benchmarks de matemáticas y código que demuestran que DISA iguala o supera las líneas base acopladas en línea, mientras supera significativamente a las líneas base de maximización de recompensa en métricas de diversidad.

4. Resultados Experimentales

Los autores evaluaron DISA en dos backbones de pesos abiertos (Qwen2.5-7B y Qwen3-4B-Base) a través de seis benchmarks de matemáticas (incluyendo AIME 2024/2025, AMC 2023, MATH-500) y tres benchmarks de código (LiveCodeBench, CodeForces, HumanEval+).

• Razonamiento Matemático: DISA igualó o superó la línea base acoplada en línea FlowRL en rendimiento promedio (Mean@8). En Qwen3-4B-Base, DISA logró 65.5 Mean@8, superando a FlowRL (64.8) y excediendo significativamente a las líneas base de maximización de recompensa GRPO (49.4) y GSPO (62.7).
• Generación de Código: DISA logró las puntuaciones pass@16 más fuertes (una métrica sensible a la diversidad), alcanzando 31.4 en Qwen3-4B-Base, superando a FlowRL (30.9), GSPO (29.9) y GRPO (27.7).
• Preservación de la Diversidad: Una evaluación con LLM como juez en problemas de AIME mostró que DISA retuvo la mayor fracción de diversidad a nivel de estrategia entre los métodos de RL (3.72 en una escala de 1-5), superando significativamente a GRPO y GSPO (~3.0) y superando ligeramente a FlowRL (3.40).
• Distilación vs. RL: Al compararse con el ajuste fino (SFT) en las mismas trayectorias de alta recompensa utilizadas para la estimación offline de DISA, DISA proporcionó 13.8 puntos Mean@8 adicionales en Qwen3-4B-Base, demostrando que ajustar la distribución suave inclinada por la recompensa recupera información del paisaje descartada por los objetivos duros de SFT.
• Análisis de Sensibilidad: Los estudios sobre la fuerza del modelo proponente y la temperatura inversa ($\beta$) siguieron los patrones de sesgo-varianza predichos, confirmando que un proponente más fuerte reduce la varianza en las etiquetas offline, y un $\beta$ óptimo equilibra la fuerza de la señal contra los acantilados de varianza.

5. Significado y Afirmaciones

El artículo afirma que DISA representa un cambio arquitectónico significativo en el RL de ajuste de distribución para LLM. Al desacoplar la estimación de la función de partición del bucle de aprendizaje de la política, DISA:

• Preserva el Óptimo Inclinado por la Recompensa: Mantiene las garantías teóricas del ajuste de distribución sin la inestabilidad del entrenamiento conjunto en línea.
• Habilita Diagnósticos Independientes: La separación permite a los investigadores diagnosticar la calidad de la función de partición (a través de la pérdida de validación del regresor) independientemente del rendimiento de la política.
• Mejora la Diversidad: Recupera con éxito la estructura de solución multimodal de las tareas de razonamiento, previniendo el colapso de modos típico de los métodos de maximización de recompensa mientras mantiene una alta precisión.

Los autores concluyen que DISA ofrece una alternativa robusta y fundamentada teóricamente a los métodos acoplados en línea, recuperando efectivamente los beneficios del ajuste de distribución mientras mitiga los errores de calibración que históricamente han obstaculizado su aplicación en el post-entrenamiento de LLM. Se sugiere trabajo futuro para extender este enfoque a dominios que carecen de un modelo proponente inicial fuerte, potencialmente mediante distilación autoiterativa.