Resumen Técnico: Guía Explícita de Críticos para Alinear Modelos de Difusión

Declaración del Problema

Alinear modelos de difusión post-entrenamiento con objetivos no diferenciables (por ejemplo, preferencia humana, corrección composicional, renderizado de texto) es cada vez más crítico, ya que el preentrenamiento por sí solo a menudo produce resultados subóptimos en propiedades centradas en el usuario. Un desafío central en este dominio es la asignación de crédito. Los modelos de difusión generan muestras mediante trayectorias de eliminación de ruido multi-paso, pero las recompensas suelen definirse solo en la imagen decodificada final. Las soluciones existentes enfrentan limitaciones significativas:

• Métodos Relativos a Grupos (por ejemplo, GRPO): Estos evitan críticos explícitos comparando muestras dentro de un grupo. Aunque evitan entrenar una función de valor, la señal de recompensa resultante es gruesa y a nivel de trayectoria, careciendo de asignación de crédito detallada para decisiones intermedias de eliminación de ruido.
• Críticos en Espacio de Píxeles: Los métodos existentes de actor-crítico explícito a menudo definen críticos en el espacio de píxeles, evaluando imágenes "limpias" reconstruidas a partir de latentes ruidosos intermedios. Esto crea una discrepancia de estado: la política opera sobre estados latentes ruidosos, mientras que el crítico evalúa estados proxy que a menudo están fuera de distribución para los codificadores de visión estándar (por ejemplo, CLIP, BLIP). Esta discrepancia conduce a una estimación de valor frágil, inestabilidad en la optimización y una sobrecarga sustancial de decodificación.
• Hackeo de Recompensas: Optimizar para recompensas verificables específicas a menudo conduce a soluciones atajo (por ejemplo, fondos simplificados, patrones repetitivos) que maximizan la métrica objetivo mientras degradan aspectos no medidos de la calidad de la imagen.

Metodología

Los autores proponen un Marco Actor-Crítico Latente Alineado al Estado que aborda estos problemas integrando el crítico directamente en el proceso de generación de difusión.

1. Crítico Latente Alineado al Estado

En lugar de operar en el espacio de píxeles, el marco aprende una función de valor $V_\phi(z_t, t, y)$ directamente sobre los estados latentes ruidosos $z_t$ visitados por la política.

• Arquitectura: El crítico se inicializa a partir del backbone de difusión preentrenado (UNet o DiT) y se modifica con cambios arquitectónicos ligeros para predecir un valor escalar.
  • UNet: El mapa de características se divide en parches, y un token de resumen aprendible agrega información mediante una cabeza de atención y una MLP.
  • DiT: Se insertan ramas de crítico ligeras en capas intermedias (por ejemplo, capas 7, 15, 23) utilizando tokens de consulta aprendibles para agregar información visual.
• Alineación del Estado: Al operar en el mismo espacio latente y compartir el prior generativo, el crítico es naturalmente consciente del paso de tiempo y del nivel de ruido, evitando los problemas de desplazamiento de distribución de los críticos en espacio de píxeles.

2. Técnicas de Estabilización del Entrenamiento

Para garantizar una optimización estable de actor-crítico, se introducen dos técnicas específicas:

• Condicionamiento Consciente del Paso de Tiempo: La cabeza de valor se condiciona explícitamente al paso de tiempo $t$ (mediante modulación estilo AdaLN), proporcionando acceso directo al nivel de ruido para mejorar la predicción de valor a lo largo de la trayectoria.
• Preentrenamiento de Valor: Antes de la optimización conjunta actor-crítico, se realiza una breve etapa de calentamiento donde solo se actualiza el crítico utilizando una pérdida de regresión contra objetivos de retorno de Monte Carlo. Esto produce estimaciones de valor iniciales mejores, estabilizando las actualizaciones PPO subsiguientes.

3. Dirección en Tiempo de Inferencia

Dado que el crítico se define sobre los mismos estados latentes que la política, puede reutilizarse en tiempo de inferencia como un módulo de guía. El gradiente de la función de valor, $\nabla_{z_t} V_\phi$, apunta hacia regiones latentes predichas para producir recompensas terminales más altas. Este gradiente se inyecta en la actualización del muestreador (modificando la predicción de ruido para DDIM o la velocidad para la coincidencia de flujo), dirigiendo efectivamente la trayectoria de eliminación de ruido sin requerir entrenamiento adicional.

4. Optimización Multi-Recompensa

Para mitigar el hackeo de recompensas, el marco se extiende a la optimización multi-recompensa. Se utiliza un backbone latente compartido con cabezas de valor separadas y ligeras para cada tipo de recompensa. Este diseño permite que recompensas heterogéneas (por ejemplo, composicional vs. estética) se regularicen mutuamente mientras comparten la extracción de características costosa, evitando el colapso en soluciones atajo comunes en el entrenamiento de recompensa única.

Contribuciones Clave

1. Crítico Latente Alineado al Estado: Un marco que convierte un modelo de difusión preentrenado en su propia red de valor condicionada al paso de tiempo, habilitando la predicción de valor en latentes ruidosos y la dirección en tiempo de inferencia.
2. Estrategias de Estabilización: Identificación del condicionamiento consciente del paso de tiempo y el preentrenamiento de valor como técnicas esenciales para estabilizar el entrenamiento explícito de actor-crítico en difusión.
3. Extensión Multi-Recompensa: Un mecanismo que utiliza cabezas de valor separadas en un backbone compartido para optimizar conjuntamente recompensas complementarias, mitigando la sobreoptimización de proxy y el hackeo de recompensas.
4. Superioridad Empírica: Rendimiento consistente superior sobre métodos relativos a grupos (GRPO) y baselines actor-crítico anteriores tanto en benchmarks de recompensa única como multi-recompensa en backbones UNet y DiT.

Resultados Experimentales

El método fue evaluado en Stable Diffusion 1.5 (UNet) y Stable Diffusion 3.5M (DiT) a través de diversos tipos de recompensa:

• Benchmarks de Recompensa Única: El método logró resultados state-of-the-art en métricas CLIP, HPSv2.1, PickScore, GenEval y OCR. Notablemente, mostró ganancias significativas en recompensas verificables (GenEval, OCR), donde señales dispersas y discretas se benefician más de estimaciones de valor por paso de tiempo. También demostró estabilidad donde DDPO falló en converger en SD3.5M.
• Benchmarks Multi-Recompensa: En tareas de optimización conjunta (por ejemplo, CLIP + HPS + GenEval), el método superó a GRPO, particularmente en GenEval, manteniendo un rendimiento comparable en métricas de preferencia.
• Dirección en Tiempo de Inferencia: Aplicar el crítico aprendido como señal de guía durante el muestreo mejoró aún más la calidad de generación. En SD1.5, la guía del crítico por sí sola superó al muestreo Best-of-N (BoN); en SD3.5M, combinar la guía con BoN produjo los mejores resultados.
• Eficiencia: El crítico latente fue significativamente más rápido que los críticos en espacio de píxeles (CLIP/BLIP/HPS) tanto en el cálculo de valor (10s vs. 28–35s) como en los pasos de entrenamiento (20s vs. 52–86s).
• Ablaciones:
  • Eliminar el preentrenamiento de valor degradó el rendimiento, especialmente en recompensas dispersas.
  • Cabezas de valor separadas para tareas multi-recompensa superaron a diseños de cabeza compartida.
  • Entrenar sin Guía Libre de Clasificador (CFG) produjo la optimización más estable y las recompensas finales más altas, evitando la inestabilidad y los artefactos asociados con escalas CFG altas.

Significado y Afirmaciones

El artículo afirma establecer una receta práctica para aplicar aprendizaje por refuerzo a modelos de difusión modernos. Su importancia radica en:

1. Resolver la Discrepancia de Estado: Al alinear el crítico con el espacio de estados latentes de la política, elimina la fragilidad y la sobrecarga de la evaluación en espacio de píxeles.
2. Habilitar el Entrenamiento Estable de Actor-Crítico: La combinación de alineación latente, condicionamiento por paso de tiempo y preentrenamiento de valor hace viable y estable el entrenamiento explícito de actor-crítico para modelos de difusión, ofreciendo una alternativa más eficiente en muestras a los métodos relativos a grupos.
3. Utilidad Dual: El crítico aprendido sirve a un doble propósito: optimizar la política durante el entrenamiento y dirigir el proceso de generación en tiempo de inferencia, proporcionando ganancias de rendimiento más allá de la fase de entrenamiento.
4. Mitigar el Hackeo de Recompensas: La extensión multi-recompensa ofrece un mecanismo práctico para equilibrar objetivos competitivos, reduciendo la prevalencia de soluciones atajo en RL de difusión.

Los autores reconocen limitaciones, señalando el enfoque actual en difusión de imágenes (no video), la posible degradación de eficiencia en configuraciones multi-recompensa en comparación con el entrenamiento de recompensa única, y los errores de discretización inherentes al muestreo SDE.