Resumen Técnico: Recompensas Latentes Adaptativas a la Muestra para el Post-entrenamiento de Difusión Guiado por Incertidumbre

1. Planteamiento del Problema

El post-entrenamiento de modelos de difusión visual para alinearlos con las preferencias humanas depende típicamente de la evaluación de las salidas RGB decodificadas, lo que requiere una decodificación VAE repetitiva y costosa dentro del bucle de retroalimentación. Aunque han surgido modelos de recompensa latente ("latent reward models" o LRM) "nativos de la difusión" que califican directamente los latentes ruidosos —evitando los costos de decodificación y permitiendo la supervisión de estados intermedios—, los enfoques existentes sufren de dos limitaciones críticas:

1. Las salidas escalares carecen de incertidumbre: Los modelos de recompensa latente actuales solo producen una puntuación escalar. No logran estimar la incertidumbre de cada predicción. En consecuencia, el generador no puede distinguir entre retroalimentación fiable e no fiable.
2. Suposiciones de incertidumbre estáticas: Trabajos previos como DiNa-LRM asignan la incertidumbre basándose únicamente en el nivel de ruido (paso de tiempo o timestep) del latente. Asumen que todas las muestras en un mismo paso de tiempo comparten la misma fiabilidad. Sin embargo, el artículo argumenta que la dificultad de la recompensa también depende del contenido; una muestra "ruidosa" con un sujeto claro puede ser más predecible que una muestra "limpia" con contenido ambiguo.
3. Riesgos de optimización: Sin la estimación de la incertidumbre, el generador puede optimizar basándose en retroalimentación local no fiable, lo que conduce al "reward hacking" (donde el modelo explota la función de recompensa en lugar de mejorar la calidad) y a una dinámica de entrenamiento inestable.

2. Metodología

Los autores proponen SURE (Recompensa Latente Adaptativa a la Muestra para el Post-entrenamiento de Difusión Guiado por Incertidumbre), un marco unificado que consta de dos componentes: SURE-LRM (el modelo de recompensa) y SURE-REFL (el algoritmo de post-entrenamiento).

2.1 SURE-LRM: Modelo de Recompensa Latente Adaptativo a la Muestra

SURE-LRM está diseñado para predecir no solo una puntuación de recompensa, sino una distribución de utilidad gaussiana completa para cada latente ruidoso, condicionado al prompt y al paso de tiempo.

• Arquitectura: Reutiliza la arquitectura Diffusion Transformer (DiT) del generador objetivo (tanto para imágenes como para videos) para procesar latentes ruidos, prompts y pasos de tiempo sin necesidad de decodificación VAE.
• Objetivos de Predicción: Para un latente ruidoso $z_\sigma$, el modelo predice:
  • Una media de recompensa $\mu_\theta$.
  2. Un parámetro de varianza residual adaptativo a la muestra $a_\theta$.
• Modelado de la Varianza: La varianza predictiva total $v_\eta(\sigma)$ es la suma de una prioridad dependiente del ruido $v_0(\sigma)$ (que captura la pérdida de información promedio a medida que aumenta el ruido) y una varianza residual aprendida y específica de la muestra $v_{\eta, res} = \text{softplus}(a_\theta) + \epsilon_v$.
• Objetivo de Entrenamiento: El modelo se entrena con preferencias por pares estándar ($z^+ \succ z^-$) utilizando un objetivo de heterocedasticidad:
  • Pérdida de Verosimilitud ($L_{pref}$): Basada en el modelo de utilidad aleatoria de Thurstone, optimiza la probabilidad de que la muestra preferida tenga una utilidad mayor. Los gradientes fomentan naturalmente que el modelo aumente la varianza (incertidumbre) cuando el orden es incorrecto o ambiguo, y la disminuya cuando el orden es confiable y correcto.
  • Regularización:
    • Regularización de la Varianza ($L_{var}$): Evita que el modelo explique todos los errores de clasificación simplemente inflando la incertidumbre (inflación de la varianza).
    • Regularización de Margen ($L_{margin}$): Asegura una separación mínima entre las puntuaciones preferidas y las rechazadas.
    • Pérdida de Anclaje ($L_{anchor}$): Estabiliza la escala general de la puntuación.

2.2 SURE-REFL: Aprendizaje de Retroalimentación de Recompensa Densa Guiado por Incertidumbre

Una vez entrenado el SURE-LRM, sus parámetros se congelan. SURE-REFL utiliza este modelo para guiar el post-entrenamiento del generador en el espacio latente.

• Mecanismo de Retroalimentación Local: En lugar de retropropagar a través de todo el grafo de eliminación de ruido, SURE-REFL consulta el modelo de recompensa congelado en transiciones seleccionadas ($j$) a lo largo de la trayectoria de eliminación de ruido.
• Ponderación Guiada por la Incertidumbre:
  1. El modelo predice la media de la recompensa $\mu$ y la varianza residual para el estado en la transición $j$.
  2. La varianza total se convierte en un peso de fiabilidad $\rho = V^{-1/2}$ (inverso de la desviación estándar).
  3. Este peso se desvincula (stop-gradient) de la recompada, lo que significa que el generador aprende del peso pero no actualiza las estimaciones de incertidumbre del modelo de recompensa durante este paso.
  4. Los pesos se normalizan entre las muestras evaluadas en la misma transición para asegurar una comparación justa.
• Optimización: El generador se actualiza utilizando un objetivo de recompensa ponderado donde las muestras con menor varianza predicha (mayor fiabilidad) contribuyen más fuertemente al gradiente. Esta retroalimentación densa se aplica localmente a cada transición, evitando la necesidad de retención de todo el grafo o decodificación al espacio de píxeles.

3. Contribuciones Clave

1. Formulación de la Incertidumbre Adaptativa a la Muestra: El artículo formula la incertidumbre de la recompensa en el espacio latente como una distribución aprendible derivada de las preferencias por pares estándar, sin requerir anotaciones de confianza humanas adicionales.
2. SURE-LRM: Un modelo de recompensa latente que aprende la incertidumbre dependiente de la muestra, demostrando que las muestras en un mismo paso de tiempo pueden tener diferentes estimaciones de fiabilidad basadas en su contenido.
   el mismo proceso de lógica de ponderación de incertidumbre tanto para la generación de imágenes como de video.
3. SURE-REFL: Un método de post-entrenamiento unificado tanto para imagen como para video que aprovecha la incertidumbre desvinculada para ponderar la retroalimentación local densa, eliminando la necesidad de decodificación VAE y de la retención de todo el grafo de eliminación de ruido.
4. Perspectiva Teórica: Los autores derivan gradientes de varianza para demostrar cómo la pérdida de preferencia proporciona una señal intrínseca para el aprendizaje de la incertidumbre.

4. Resultados Experimentales

4.1 Precisión del Modelo de Recompensa (SURE-LRM)

• Benchmarks de Imagen: SURE-LRM alcanzó la precisión de preferencia por pares de vanguardia (SOTA) en cuatro benchmarks (ImageReward, HPD v2, HPD v3, GenAI-Bench) en tres backbones diferentes (SD3.5, Z-Image-Turbo, FLUX.1-dev), superando a bases sólidas incluyendo modelos de píxeles basados en VLM y DiNa-LRM.
• Benchmarks de Video: En VisionRewardDB y T2V Ranking, SURE-LRM alcanzó la mayor precisión, superando a los evaluadores de video decodificados y a los modelos de recompensa específicos de video.
• Diagnóstico de Incertidumbre: Los experimentos confirmaron que la varianza residual aprendida se correlaciona con los errores de predicción. Las muestras con clasificaciones incorrectas tenían una varianza residual significativamente mayor (un 24.26% más alta en promedio). Filtrar las muestras por baja varianza predicha mejoró la precisión por pares del 79.40% al 90.10%.

4.2 Rendimiento de Post-entrenamiento (SURE-REFL)

• Generación de Imágenes: SURE-REFL logró el mejor rendimiento en las métricas HPSv2.1, ImageReward y PickScore en SD3.5, Z-Image-Turbo y FLUX.1-dev, superando a bases como Flow-GRPO, ViPO y DiNa-LRM.
• Generación de Video: En VBench, SURE-REFL alcanzó las puntuaciones más altas en las métricas de Calidad, Semántica y Total para el generador de video Wan2.1-14B, superando a DanceGRPO, ViPO, Video-DPO y HY-PRFL.
• Estabilidad y Reward Hacking: La dinámica de optimización mostró que, mientras los modelos base (por ejemplo, DiNa-LRM-REFL) exhibieron reward hacking (las puntuaciones de recompensa subían mientras que métricas de calidad independientes como HPSv3 caían), SURE-REFL mantuvo la alineación entre la recompensa optimizada y las métricas de calidad independientes durante más tiempo, alcanzando puntuaciones finales más altas.

4.3 Estudios de Ablación

La ablación de los componentes de pérdida del modelo de recompensa confirmó que todos los términos son necesarios. Eliminar la regularización de la varianza ($L_{var}$) causó la mayor caída en la precisión, validando la necesidad de controlar la inflación de la varianza.

5. Significado y Reivindicaciones

El artículo afirma que SURE proporciona un marco unificado y eficiente para alinear modelos de difusión visual abordando la "brecha de fiabilidad" en el modelado de recompensa latente. Su importancia radica en:

• Eficiencia: Al operar enteramente en el espacio latente sin decodificación VAE, reduce la carga computacional durante el bucle de retroalimentación.
• Robustez: Al modelar y utilizar explícitamente la incertidumbre adaptativa a la muestra, evita que el generador se sobreajuste a retroalimentación ruidosa o ambigua, reduciendo así el reward hacking y mejorando la estabilidad del entrenamiento.
• Generalidad: El marco es agnóstico a la modalidad, aplicando con éxito la misma lógica de ponderación de incertidumbre tanto a la generación de imágenes como a la de video.

Los autores concluyen que aprender la incertidumbre dependiente de la muestra a partir de las preferencias por pares estándar es una estrategia viable y efectiva para mejorar la fiabilidad de la retroalimentación de recompensa densa en el post-entrenamiento de la difusión.