PG-MAP: Joint MAP Optimization for Inference-Time Alignment of Diffusion and Flow-Matching Models

Resumen Técnico: PG-MAP: Optimización de MAP Conjunta para la Alineación en Tiempo de Inferencia de Modelos de Difusión y Flow-Matching

Declaración del Problema
Los métodos actuales de alineación en tiempo de inferencia para modelos preentrenados de texto a imagen (Difusión y Flow-Matching) operan típicamente a lo largo de un único eje de control. Los métodos del lado de la condicionante (por ejemplo, edición de atención, aprendizaje de embeddings de texto) modifican la codificación de texto $c$, mientras que los métodos del lado del latente (por ejemplo, perturbaciones de latentes guiadas por recompensa) modifican el estado latente $z_t$. Los marcos existentes no logran modelar las dependencias conjuntas entre las variables de condicionamiento y latentes, tratando las actualizaciones como aditivas en lugar de coordinadas. Además, estos métodos suelen ser estáticos, fijando las estrategias de control al inicio de la trayectoria ($z_T$) o de forma offline, ignorando las necesidades de información dinámica del denoiser a medida que la trayectoria evoluciona desde regímenes de alto ruido (diseño global) hacia regímenes de bajo ruido (detalle local). Esta limitación dificulta la transferencia de técnicas de alineación a través de diferentes transportes generativos (Difusión vs. Flow-Matching).

Metodología: PG-MAP
Los autores proponen PG-MAP (Preference-Guided Adaptive MAP), un marco libre de entrenamiento que formula la alineación en tiempo de inferencia como una optimización de energía proximal / Gibbs-MAP a nivel de trayectoria. La innovación central es un acoplamiento de consistencia hacia adelante que optimiza conjuntamente la condicionante $c$ y el estado latente $z_t$ en cada paso de denoisado.

• Función Objetivo: El marco define una función de energía dependiente del paso $J_t(c, z_t)$ que consta de tres componentes:

  1. Residuo de Consistencia hacia Adelante: Un término $\ell_t(c, z_t)$ que penaliza la desviación entre el latente actual $z_t$ y la predicción del modelo del estado del paso anterior, asegurando que la trayectoria sea consistente con el kernel hacia adelante del denoiser.
  2. Prios de Anclaje Gaussiano: Términos $R_c(c)$ y $R_z(z_t)$ que anclan las variables optimizadas a sus valores no perturbados ($c_0$ y el punto de la trayectoria DDIM $z^{ddim}_t$). Crucialmente, la escala del prior latente $\sigma_z(t)$ es adaptativa al esquema, siguiendo la escala de ruido marginal del proceso de difusión ($\sigma_z(t) = \gamma\sqrt{1-\bar{\alpha}_t}$) para proporcionar una región de confianza invariante a la escala que se estrecha a medida que la trayectoria se acerca a la variedad de datos.
  3. Inclinación de Recompensa de Preferencia: Un término $\lambda Q(\hat{x}_0, y)$ que guía la optimización hacia la preferencia humana (por ejemplo, PickScore, HPS) basada en la vista previa por paso de la imagen limpia del modelo.

• Conjunto Activo Adaptativo: El marco emplea un conjunto activo dependiente del paso $A_t \subseteq \{c, z_t\}$ y una puerta de recompensa $\lambda_t$.

  • Modelos de Difusión (SD 1.5, SDXL): El marco refina conjuntamente tanto $c$ como $z_t$, principalmente en la ventana de alto ruido, permitiendo que la condicionante y las variables latentes se informen mutuamente a medida que la trayectoria se desarrolla.
  • Modelos de Flow-Matching (SD3.5): Debido a la naturaleza determinista de la ODE y la alta capacidad de los codificadores de texto concatenados, el marco se reduce a una variante de solo latente (UG-FM). Aquí, el conjunto activo colapsa a $\{z_t\}$ restringido a la ventana del lado de los datos (bajo ruido), ya que las actualizaciones de condicionamiento resultan ineficaces y las perturbaciones del lado del ruido son excesivamente amplificadas por la dinámica de la ODE.

• Optimización: El par refinado $(c^\star_t, z^\star_t)$ se obtiene mediante un pequeño número de pasos de ascenso conjunto (típicamente $K=2$) utilizando gradientes derivados de la función de energía. Este proceso es composable con métodos existentes como Classifier-Free Guidance (CFG), ya que CFG modifica el campo vectorial mientras que PG-MAP mueve el punto de consulta bajo un campo vectorial fijo.

Contribuciones Clave

1. Marco MAP Conjunto ($c, z_t$): El primer marco de tiempo de inferencia que acopla las variables de condicionamiento y latentes a través del kernel hacia adelante del denoiser, abordando simultáneamente los fallos de composición (lado de $c$) y los fallos de textura (lado de $z$).
2. Objetivo Unificado: El objetivo $J_t$ recupera analíticamente métodos previos de un solo eje (MAP-c, Reward-z) como casos especiales, mientras que CFG permanece como un mecanismo externo composable en lugar de un caso especial.
3. Optimización Adaptativa al Esquema: Introduce una región de confianza dependiente del tiempo y un conjunto activo que se adapta al esquema de ruido, relajando las restricciones en alto ruido y estrechándolas cerca del final de los datos.
4. Análisis Dependiente del Transporte: Demuestra que el conjunto activo óptimo depende del transporte generativo. La difusión se beneficia del refinamiento conjunto en alto ruido, mientras que el Flow-Matching se beneficia del refinamiento de solo latente al final de los datos, una hipótesis respaldada por el análisis de perturbación local y la validación empírica.

Resultados

• Backbones de Difusión (SD 1.5, SDXL): PG-MAP mejora consistentemente las métricas de alineación. En SDXL, logra una tasa de victoria de 56.4% en PickScore y 56.2% en Aesthetic contra un baseline estático. Combinar PG-MAP con un CFG ajustado produce el rendimiento general más fuerte, alcanzando un 64.6% de HPS y 56.5% de Aesthetic.
• Backbone de Flow-Matching (SD3.5-medium): La variante de solo latente (UG-FM) logra tasas de victoria del 91.9% en PickScore y 75.7% en HPS contra el baseline estático. Experimentos controlados descartan artefactos relacionados con el ruido, confirmando que las ganancias se deben a perturbaciones latentes estructuradas de baja frecuencia.
• Evaluación Humana: En SDXL, PG-MAP es preferido sobre los baselines estáticos, el CFG ajustado y el Universal Guidance con cómputo equivalente en evaluaciones humanas pareadas (tasas de victoria del 60–67%).
• Enrutamiento Oráculo: Un análisis de oráculo revela que diferentes tipos de prompts (por ejemplo, unión de atributos vs. escenas atmosféricas) se benefician de diferentes variantes de optimización. Un selector por prompt podría teóricamente mejorar las métricas en un 5–14 de puntos porcentuales adicional, indicando un margen significativo para futuros mecanismos de enrutamiento.

Significancia y Reivindicaciones
El artículo afirma que PG-MAP representa un cambio de las heurísticas de guía estáticas de un solo eje hacia una optimización dinámica consciente de la trayectoria como principio de diseño por defecto para la alineación en tiempo de inferencia. Al formular la alineación como una optimización de energía proximal con acoplamiento de consistencia hacia adelante, el marco permite actualizaciones coordinadas a través de las modalidades que son compatibles tanto con los transportes de difusión como de flow-matching. Los autores enfatizan que los ingredientes "de carga" son el acoplamiento conjunto y la programación no estacionaria, en lugar de simplemente pasos más grandes o señales de recompensa más fuertes. El trabajo proporciona una base teórica y empírica unificada para comprender cómo interactúan la condicionante y las variables latentes durante el proceso de generación, ofreciendo un enfoque flexible y libre de entrenamiento para mejorar la alineación de los modelos generativos.