Resumen Técnico: Aprendizaje por Refuerzo Guiado por Discriminador para Flow Matching

Planteamiento del Problema

Los modelos basados en flujo y de coincidencia de puntuación (score-matching) se han convertido en el paradigma dominante para el modelado generativo debido a sus objetivos de entrenamiento basados en regresión y libres de simulación. Sin embargo, estos modelos suelen entrenarse en un proceso de múltiples etapas: primero, se ajusta un modelo base a una distribución de datos $q$ mediante flujo o score matching (FSM); segundo, se aplica aprendizaje por refuerzo (RL) para inclinar el modelo hacia regiones deseables definidas por una recompensa $r(x)$.

El artículo identifica un desajuste estructural en este paradigma. Mientras que las pérdidas de FSM minimizan el error de regresión $\ell_2$ sobre los campos de velocidad o puntuación bajo los marginales de entrenamiento, a menudo no logran recuperar las propiedades intrínsecas de la distribución de datos —como el realismo visual, la estructura coherente de los objetos o el movimiento físicamente correcto— que están teóricamente presentes en los datos. En consecuencia, los profesionales dependen del RL basado en preferencias para recuperar estas propiedades. Este enfoque es problemático por dos razones:

1. Fallo Estructural: Resulta desconcertante que los modelos no puedan aprender propiedades del mundo real directamente de los datos mediante pérdidas de coincidencia, pero requieran RL para recuperarlas.
2. Costo y Confluencia: El entrenamiento de recompensas suele requerir costosos datos de preferencia humana, lo que conlave la realización objetiva del realismo de los datos con las inclinaciones subjetivas de los anotadores.

Los autores argumentan que el fallo de los objetivos de coincidencia al capturar estas propiedades se debe a un "desajuste entre entrenamiento y prueba" (train-test mismatch, donde los errores se acumulan a lo largo de la distribución de despliegue del modelo en lugar de los marginales de entrenamiento) y un "desajuste geométrico" (donde pequeños errores en el espacio de velocidad pueden conducir a grandes errores en la calidad de la muestra/espacio de recompensa).

Metodología: Aprendizaje por Refuerzo Guiado por Discriminador (DRL)

El artículo propone el Aprendizaje por Refuerzo Guiado por Discriminador (DRL), un método para corregir los modelos de flow matching sin depender de datos de preferencia humana. La idea central es que la recompensa óptima para alinear un modelo base $p_{base}$ con una distribución objetivo $q$ es la razón de densidad logarítmica $\log(q(x)/p_{base}(x))$. Dado que esta razón es intratable, el DRL la estima utilizando un discriminador.

El flujo de trabajo de DRL consta de dos etapas:

Etapa 1: Estimación de la Recompensa en el Espacio de Representación

En lugar de entrenar un discriminador directamente en el espacio de píxeles (lo cual es estadísticamente difícil y semánticamente poco fiable), el DRL entrena un discriminador en un espacio de representación de aprendizaje autosupervisado (SSL) preentrenado.

• Un codificador congelado $\phi$ (por ejemplo, DINOv2) mapea las imágenes a un espacio de características $Z$.
• Un discriminador $D$ se entrena para distinguir entre muestras de datos reales $x \sim q$ y muestras del modelo base $x_{fake} \sim p_{base}$ utilizando sus representaciones codificadas $\phi(x)$.
• La función de recompensa se define como el logit de este discriminador: $\hat{r}(x) = \log(D(\phi(x)) / (1 - D(\phi(x))))$.
• Justificación Teórica: Al restringir la recompensa al espacio SSL, el DRL resuelve un problema de optimización restringida: encuentra la distribución $p$ más cercana a $p_{base}$ (en divergencia KL) tal que el pushforward de $p$ a través de $\phi$ coincida con el pushforward de $q$ a través de $\phi$. Esto asegura que las correcciones se realicen a lo largo de ejes semánticamente significativos visibles para el codificador.

Etapa 2: RL Regularizado por KL mediante Coincidencia Adjunta (Adjoint Matching)

La recompensa estimada $\hat{r}$ se utiliza para ajustar el modelo base mediante RL regularizado por KL.

• El objetivo es maximizar la recompensa esperada mientras se minimiza la divergencia KL entre las distribuciones de trayectoria del modelo ajustado y el modelo base.
• Los autores utilizan Coincidencia Adjunta (Adjoint Matching), un algoritmo de vanguardia para modelos de flujo que optimiza el objetivo a nivel de trayectoria mediante la propagación de gradientes a través del proceso de muestreo.
• Para asegurar la estabilidad del entrenamiento con el SDE sin memoria requerido para el sesgo de extremo (endpoint tilting) no sesgado, los autores introducen un Integrador Lineal Local. Este integrador maneja la rigidez numérica cerca de $t=0$ (donde los coeficientes de deriva y difusión divergen) de manera más efectiva que las discretizaciones estándar de Euler-Maruyama con desplazos (offsets).

Contribuciones Clave

1. Análisis Teórico de las Limitaciones de la Coincidencia: El artículo proporciona un análisis formal que muestra que las pérdidas de flujo/coincidencia de puntuación no garantizan el control sobre las brechas de recompensa (propiedades de interés) debido a los cambios de distribución entre entrenamiento y prueba, y la desalineación entre la geometría del espacio de velocidad y la geometría de recompensa del espacio de muestra.
2. Aprendizaje por Refuerzo Guiado por Discriminador (DRL): Un marco novedoso que utiliza un discriminador entrenado en un espacio de representación SSL congelado para generar una señal de recompensa. Esto permite al modelo recuperar propiedades de los datos sin etiquetas de preferencia humana.
3. Integrador Lineal Local: Un nuevo integrador numérico para el SDE sin memoria utilizado en el entrenamiento de RL, que evita la explosión de la norma en tiempos tempranos y mejora tanto la calidad del muestreo como la estabilidad del entrenamiento.
4. Validación Empírica: Demostración de que el DRL mejora significativamente la alineación de la distribución y la calidad de la imagen a través de diversas arquitecturas (SiT, JiT, REPA, RAE) y espacios de características (DINOv2, DINOv3, SigLIP, Inception).

Resultados

Los autores evaluaron el DRL en cuatro modelos de flujo preentrenados en ImageNet. Los hallazgos clave incluyen:

• Alineación Distribuida: El DRL reduce significativamente la Distancia de Fréchet (FD). Por ejemplo, en SiT, la FD sin guía en el espacio DINOv3 cayó de 88.2 a 19.3. En todos los modelos y espacios de características, el DRL mejoró consistentemente la alineación sin utilizar CFG (Classifier-Free Guidance).
• Calidad de Imagen sin Datos de Preferencia: El DRL mejoró las puntuaciones en modelos de recompensa de preferencia humana mantenidos fuera del entrenamiento (ImageReward, HPSv2, PickScore, Aesthetics), a pesar de no haber sido entrenado en ellos. Esto sugiere que el DRL recupera con éxito las propiedades de "realismo" que las pérdidas de coincidencia estándar omiten.
• Base Mejorada para la Alineación de Preferencias: Cuando el DRL se utiliza como punto de partida para el posterior RL basado en preferencias (PRL), produce una frontera de Pareto superior entre la recompensa y la fidelidad de la imagen. El DRL absorbe las correcciones estructurales recuperables de los datos, permitiendo que el PRL se concentre en preferencias subjetivas genuinas en lugar de corregir artefactos de bajo nivel (por ejemplo, sobresaturación, brillo excesivo).
• Brecha de Destilación: Los experimentos mostraron que un modelo estudiante entrenado mediante coincidencia de flujo estándar sobre muestras de un profesor DRL no pudo replicar el rendimiento del profesor, incluso tras ver 50 millones de muestras. Esto confirma que el propio objetivo de RL, y no solo las muestras, es necesario para recuperar estas propiedades.

Significado y Reivindicaciones

El artículo sostiene que el post-entrenamiento de RL para modelos de flujo debe verse no meramente como una herramienta para optimizar preferencias externas, sino como un mecanismo complementario para recuperar la estructura en los datos de entrenamiento que es imperfectamente capturada por los objetivos de coincidencia estándar.

Al utilizar un discriminador en un espacio de representación preentrenado, el DRL proporciona una señal de recompensa que está:

1. Alineada con las Propiedades de los Datos: Apunta directamente a la distribución de los datos en lugar de a las preferencias humanas subjetivas.
2. Geométricamente Sólida: Evita las obstrucciones de la geometría del espacio de velocidad al optimizar directamente en el paisaje de recompensa.
3. Eficiente: Requiere significativamente menos cómputo que el preentrenamiento y evita el costo de recolectar datos de preferencia.

Los autores concluyen que los objetivos de coincidencia y el RL no son métodos con metas conflictivas, sino dos herramientas complementarias para maximizar la utilidad de los datos de entrenamiento. El trabajo sugiere que la "recompensa" para la generación de alta calidad a menudo "siempre ha estado en los datos", pero requiere un mecanismo de guía de recompensa on-policy para ser plenamente realizada.