Resumen Técnico: Reinforce Adjoint Matching (RAM)

Enunciado del Problema

Los modelos de difusión y de ajuste de flujo (flow-matching) logran una generación de alta fidelidad porque su preentrenamiento se basa en regresión supervisada: una muestra limpia se añade ruido analíticamente, y el modelo realiza una regresión contra un objetivo de forma cerrada. Sin embargo, alinear estos modelos con funciones de recompensa (por ejemplo, para corrección composicional, renderizado de texto o preferencia humana) mediante el entrenamiento posterior de Aprendizaje por Refuerzo (RL) interrumpe esta estructura eficiente.

Los métodos existentes de entrenamiento posterior por RL enfrentan obstáculos significativos de computación y estructura:

1. Métodos de Gradiente de Política: Tratan la eliminación de ruido como un Proceso de Decisión de Markov, requiriendo simulaciones estocásticas de EDE (Ecuaciones Diferenciales Estocásticas) y optimizando cotas inferiores variacionales. Esto introduce una varianza alta y requiere muchos pasos de muestreo.
2. Métodos de Control Óptico Estocástico (Adjoint): Apuntan exactamente al objetivo de maximización de recompensa regularizado por KL, pero requieren barridos adjuntos hacia atrás a lo largo de trayectorias de EDE almacenadas y gradientes de recompensa retropropagados a través del proceso de muestreo. Esto es numéricamente inestable y computacionalmente costoso, particularmente para modelos de ajuste de flujo donde los programas de ruido divergen cerca del extremo ruidoso.
3. Pérdidas Sustitutas: Los enfoques recientes reemplazan verosimilitudes intratables con sustitutos (por ejemplo, ELBOs o puntos finales contrastivos), sacrificando la estructura de regresión limpia que permite que el preentrenamiento escale.

El desafío central es derivar un objetivo de entrenamiento posterior por RL que mantenga la simplicidad y escalabilidad del marco de regresión de preentrenamiento mientras optimiza eficazmente las recompensas.

Metodología: Reinforce Adjoint Matching (RAM)

Los autores proponen Reinforce Adjoint Matching (RAM), un método que extiende la estructura de regresión de preentrenamiento al entrenamiento posterior por RL sin requerir simulaciones de EDE, barridos adjuntos hacia atrás ni gradientes de recompensa.

Fundamento Teórico

El método se basa en la maximización de recompensa regularizada por KL formulada como un problema de control óptimo estocástico.

1. Estructura del Proceso Óptimo: Bajo regularización KL, el proceso generativo óptimo inclina la distribución de puntos finales limpios ($X_0$) hacia muestras de alta recompensa ($p^*_0 \propto p_{ref}^0 \exp(r)$) pero deja inalterada la ley condicional de los estados ruidosos dado un punto final limpio. Específicamente, la trayectoria desde un $X_0$ fijo hasta un estado ruidoso $X_t$ permanece siendo el mismo núcleo gaussiano analítico utilizado en el preentrenamiento.
2. Condición de Ajuste Adjoint: El control óptimo $u_t$ satisface una condición de punto fijo donde el control es igual al gradiente escalado negativo de la función de valor (el adjunto): $u_t = -\sigma_t \nabla_x V_t(x)$.
3. Descomposición REINFORCE: Para evitar calcular $\nabla r$ (que puede ser no diferenciable), los autores aplican la identidad log-derivada (REINFORCE) a la función de valor. Esto descompone el adjunto en un término de recompensa ponderado por la puntuación del puente hacia atrás y un término de corrección de costo de trayectoria.

El Objetivo RAM

RAM aproxima el adjunto eliminando la corrección de costo de trayectoria (justificada por la baja varianza y la estabilidad a escala) y aprovechando la intuición estructural de que el núcleo de adición de ruido permanece analítico.

1. Muestreo de Puntos Finales: En lugar de simular una EDE controlada, RAM extrae un punto final limpio $X_0$ de la distribución del modelo actual utilizando cualquier muestreador ODE estándar.
2. Adición de Ruido Analítica: El $X_0$ muestreado se añade ruido analíticamente para generar estados de entrenamiento $X_t = (1-t)X_0 + t\epsilon$, exactamente como en el preentrenamiento.
3. Objetivo de Forma Cerrada: El objetivo de regresión para el campo de velocidad $v_\theta$ se deriva sustituyendo el término de recompensa y la puntuación del puente de Bayes (que tiene una expresión de forma cerrada para adición de ruido lineal) en la condición adjunta.

La función de pérdida resultante es:
$$L_{RAM}(\theta) = \mathbb{E}_t \left[ \left\| v_\theta^t(X_t) - \text{sg}\left( v_{ref}^t(X_t) + r(X_0) \left( (\epsilon - X_0) - v_\theta^t(X_t) \right) \right) \right\|^2 \right]$$
donde $X_0 \sim p_\theta^0$, $\epsilon \sim \mathcal{N}(0, I)$, y $\text{sg}(\cdot)$ denota el operador de stop-gradient.

Mecanismos Clave de Eficiencia

• Sin Simulaciones de EDE: Los estados de entrenamiento se generan añadiendo ruido a un solo punto final $K$ veces con vectores de ruido independientes. Esto produce $K$ muestras de entrenamiento condicionalmente independientes por pase forward del modelo, proporcionando más señal de gradiente por paso que los métodos que dependen de trayectorias únicas de EDE.
• Sin Gradientes de Recompensa: El método funciona con recompensas de caja negra.
• Sin Barridos hacia Atrás: La naturaleza analítica del núcleo de adición de ruido elimina la necesidad de integración adjunta hacia atrás.

Contribuciones Clave

1. Extensión Estructural: El artículo demuestra que la estructura de regresión del preentrenamiento se extiende al entrenamiento posterior por RL bajo regularización KL. El proceso óptimo solo cambia la distribución de puntos finales, no la ley condicional de adición de ruido.
2. Algoritmo RAM: Una nueva pérdida de consistencia que corrige el objetivo de preentrenamiento con información de recompensa utilizando una puntuación de puente de Bayes de forma cerrada, evitando los cuellos de botella computacionales de los métodos existentes.
3. Justificación Teórica: Una derivación que muestra que el punto fijo de RAM corresponde al óptimo regularizado por KL, con exactitud garantizada para referencias gaussianas bajo recompensas lineales y acuerdo de primer orden en general.
4. Análisis de Varianza: Un análisis de estimadores de costo de trayectoria que muestra que eliminar la corrección de costo de trayectoria (como hace RAM) reduce significativamente la varianza en comparación con los estimadores de adjunto de horizonte completo o locales, haciéndolo el único estimador viable para el entrenamiento posterior a escala de imágenes.

Resultados Experimentales

Los autores evaluaron RAM en Stable Diffusion 3.5 Medium (SD3.5M) en tres tareas:

1. Generación de Imágenes Composicional (GenEval): Midiendo la corrección de objetos, atributos y relaciones espaciales.
2. Renderizado de Texto Visual (OCR): Midiendo la legibilidad del texto en las imágenes generadas.
3. Alineación con Preferencias Humanas (PickScore): Alineando con preferencias humanas.

Rendimiento:

• Maximización de Recompensa: RAM logró las recompensas de entrenamiento más altas en las tres tareas, superando a líneas base como Flow-GRPO, DiffusionNFT y AWM.
• Eficiencia de Entrenamiento: RAM alcanzó la recompensa máxima de Flow-GRPO en 50 veces menos pasos para GenEval, 48 veces menos para OCR y 34 veces menos para PickScore.
• Calidad de Imagen: A diferencia de algunas líneas base que sufrieron de "hacking de recompensa" (degradación de la calidad estética para métricas de tarea), RAM mantuvo o mejoró las métricas de calidad de imagen (Aesthetic, DeQA, ImageReward, HPSv2) mientras maximizaba las recompensas de tarea.
• Compatibilidad con CFG: RAM permanece compatible con la Guía Libre de Clasificador (CFG) después del entrenamiento posterior, mientras que líneas base como DiffusionNFT y AWM vieron caer sus recompensas de tarea significativamente cuando se aplicó CFG.

Significado y Afirmaciones

El artículo afirma que RAM representa un cambio de paradigma en el entrenamiento posterior por RL para modelos de difusión y de ajuste de flujo. Al aprovechar las propiedades estructurales específicas del problema de control óptimo, RAM restaura el paradigma de "regresión simple" que hizo escalable el preentrenamiento.

Los autores afirman que:

• RAM elimina la necesidad de costosas simulaciones de EDE y barridos adjuntos hacia atrás, que han sido los principales cuellos de botella para escalar RL a la generación de imágenes de alta resolución.
• El método es "simple y escala", logrando resultados de vanguardia con significativamente menos pasos de entrenamiento y costos de computación por paso comparables a los métodos existentes.
• El enfoque alinea exitosamente los modelos con recompensas complejas (composicionalidad, texto, preferencias) sin degradación visible de la calidad ni hacking de recompensa, demostrando que la estructura de preentrenamiento es lo suficientemente robusta para manejar la alineación posterior al entrenamiento cuando se cumplen las condiciones teóricas correctas.