Resumen Técnico: Modelo de Valor Cosido para la Alineación de Difusión (StitchVM)

1. Declaración del Problema

Los modelos generativos basados en difusión y flujo destacan en la captura de distribuciones de datos, pero a menudo requieren una alineación con recompensas específicas de la tarea (por ejemplo, fidelidad al prompt, preferencia estética, preferencia humana) para ser útiles en la práctica. El desafío central en esta alineación es que las recompensas suelen definirse para imágenes de salida limpias ($x_0$), mientras que los procedimientos de alineación (tanto en tiempo de inferencia como de entrenamiento) requieren estimar la recompensa esperada de muestras limpias inducidas por latentes intermedios ruidosos ($z_t$). Esta estimación se conoce como función de valor.

Los métodos existentes para aproximar esta función de valor sufren de compensaciones significativas:

• Aproximación de Tweedie: Estima la media posterior de la muestra limpia y calcula la recompensa sobre ese proxy. Aunque es eficiente, introduce un sesgo sustancial en regímenes de alto ruido y requiere evaluaciones adicionales costosas del desruidizador y del decodificador VAE.
• Aproximación Monte Carlo (MC): Despliega múltiples trayectorias de desruido desde $z_t$ y promedia las recompensas. Aunque es más precisa, incurre en costos computacionales prohibitivos debido a la necesidad de múltiples despliegues completos.
• Modelos de Valor Aprendidos: Entrenar directamente una función de valor para latentes ruidosos es un problema abierto. Los intentos existentes a menudo dependen de entrenamientos a pequeña escala o conjuntos de datos de preferencia estrechos, resultando en una precisión y generalización inferiores en comparación con los modelos de recompensa en el espacio de píxeles a escala fundacional. Además, entrenar un nuevo modelo de valor para cada nuevo backbone de difusión o modelo de recompensa es computacionalmente insostenible.

2. Metodología: StitchVM

Los autores proponen StitchVM (Modelo de Valor Cosido), un marco que transfiere las capacidades de modelos de recompensa en el espacio de píxeles preentrenados y de alta calidad al régimen de latentes ruidosos con un costo computacional mínimo. El enfoque aprovecha el ensamblaje de modelos (model stitching), una técnica que recombina capas de diferentes redes neuronales.

Arquitectura Central

StitchVM construye un modelo híbrido al adjuntar un backbone de difusión truncado y congelado (la "cabeza") a un modelo de recompensa preentrenado y cortado (la "cola") mediante una capa de ensamblaje ligera.

• Cabeza: Un backbone de difusión (por ejemplo, SD 3.5, FLUX) truncado en la capa $i$. Este componente procesa nativamente latentes ruidosos y extrae características conscientes del ruido.
• Cola: Un modelo de recompensa preentrenado (por ejemplo, CLIP, HPSv2, Aesthetic Predictor) que comienza en la capa $j$. Este componente retiene la capacidad robusta y específica de la tarea de recompensa aprendida durante el preentrenamiento a gran escala.
• Capa de Ensamblaje ($s_\psi$): Un módulo ligero que alinea el espacio de características de la cabeza de difusión con la cola de recompensa.

Procedimiento de Entrenamiento

El entrenamiento de StitchVM ocurre en dos etapas:

1. Selección de Interfaz (Forma Cerrada): Los autores buscan los índices óptimos de ensamblaje $(i, j)$ donde las características de difusión (de latentes ruidosos) y las características de recompensa (de imágenes limpias) son más compatibles. Esto se logra ajustando un mapeo lineal $W$ entre características emparejadas extraídas de un pequeño conjunto de sondeo de imágenes limpias y sus latentes ruidosos correspondientes. Se selecciona el par que minimiza la pérdida de coincidencia de características.
2. Ajuste Fino Ligero: Una vez fijada la interfaz, la capa de ensamblaje y la cola de recompensa truncada se ajustan finamente utilizando imágenes limpias sin etiquetas. El objetivo es regredir la función de valor $V(z_t)$ hacia la puntuación de recompensa original $r(z_0)$ de la imagen limpia. Este objetivo fuera de política permite que el modelo aprenda la recompensa esperada de muestras limpias condicionadas a latentes ruidosos sin requerir despliegues completos.

Aplicación a la Alineación

StitchVM sirve como un estimador directo de la función de valor, reemplazando las aproximaciones tanto en la alineación en tiempo de inferencia como en tiempo de entrenamiento:

• Tiempo de Inferencia:
  • DPS (Muestreo Posterior de Difusión): Reemplaza el gradiente basado en Tweedie con el gradiente directo de StitchVM ($\nabla_{z_t} V(z_t)$), eliminando la necesidad de retropropagación a través del desruidizador y el VAE.
  • FK Steering: Habilita la "escalación de propuestas", donde cada partícula genera múltiples propuestas locales. Estas se califican de manera económica por StitchVM para seleccionar el mejor candidato, un proceso que sería demasiado costoso con aproximaciones de Tweedie.
• Tiempo de Entrenamiento:
  • Ajuste Fino Directo de Recompensa (por ejemplo, AlignProp, DRaFT): Permite detener el entrenamiento en latentes ruidosos intermedios $z_\tau$ y utilizar la predicción de StitchVM como objetivo de supervisión, evitando despliegues de trayectoria completos y cadenas largas de retropropagación.
  • DiffusionNFT: Reemplaza las recompensas terminales limpias con estimaciones de valor intermedias, permitiendo RL en línea eficiente en el proceso hacia adelante.

3. Contribuciones Clave

• Un Nuevo Paradigma para el Modelado de Valor: StitchVM demuestra que los potentes modelos de recompensa en el espacio de píxeles pueden transferirse efectivamente al régimen de latentes ruidosos sin reentrenamiento a gran escala, simplemente cosiéndolos a un backbone de difusión.
• Eficiencia y Precisión: El método logra un rendimiento de vanguardia en la estimación de valor mientras es excepcionalmente ligero. Por ejemplo, coser CLIP ViT-L en un modelo de valor SD 3.5 Medium toma solo $\approx 10$ horas de GPU.
• Mejoras de Rendimiento en la Alineación:
  • Inferencia: DPS se vuelve 3.2$\times$ más rápido con una reducción del 50% en la memoria pico de GPU. El FK steering con StitchVM supera significativamente al FK steering estándar y al muestreo Best-of-$N$, particularmente al habilitar una selección eficiente de propuestas locales.
  • Entrenamiento: DiffusionNFT se vuelve 2.3$\times$ más rápido, y el ajuste fino directo de recompensa (DRaFT) se vuelve 1.3$\times$ más rápido, logrando al mismo tiempo una mayor calidad de generación (por ejemplo, +30% en GenEval) debido a una supervisión efectiva en pasos de alto ruido.
• Generalización: El enfoque funciona en diversos backbones de difusión (SD 3.5 Medium/Large, FLUX) y modelos de recompensa (CLIP, DFN-CLIP, HPSv2, Aesthetic Predictor).

4. Resultados Experimentales

El artículo evalúa StitchVM en múltiples benchmarks:

• Calidad de la Función de Valor: En latentes ruidosos con niveles de ruido variables ($\sigma \in [0.1, 0.9]$), StitchVM coincide estrechamente con el rendimiento de los modelos de recompensa limpios originales en recuperación zero-shot (MSCOCO, Flickr30K), predicción de preferencias (HPDv2, ImageReward) y puntuación estética (AVA). Supera significativamente a las líneas base como el ensamblaje VAE (que colapsa con alto ruido) y los modelos entrenados desde cero con datos ruidosos (por ejemplo, NoisyCLIP, DiNa-LRM).
• Alineación en Inferencia: En los experimentos de DPS y FK steering en DrawBench, los métodos mejorados con StitchVM mejoran consistentemente métricas como ImageReward, HPSv2 y GenEval mientras reducen drásticamente el costo computacional. Cabe destacar que aumentar el número de propuestas locales ($M$) en el FK steering usando StitchVM resulta más eficiente que aumentar el número de partículas ($N$).
• Alineación en Entrenamiento: Las curvas de entrenamiento muestran que las variantes de StitchVM alcanzan puntuaciones finales más altas con menos horas de GPU en comparación con DRaFT y DiffusionNFT estándar. Los resultados cualitativos indican que los modelos aumentados con StitchVM convergen más rápido y producen imágenes más nítidas y de mayor calidad en etapas tempranas del entrenamiento.

5. Significado y Afirmaciones

El artículo afirma que StitchVM proporciona una plantilla genérica para combinar backbones de difusión latente con modelos feedforward en el espacio de píxeles sin sacrificar el extenso preentrenamiento de ninguno de los dos. Al elevar potentes modelos de recompensa al espacio latente, habilita un nuevo estilo de alineación de difusión donde la función de valor correcta para latentes ruidosos se construye una vez y se amortiza sobre muchas muestras e iteraciones.

Los autores posicionan a StitchVM no como un reemplazo para el preentrenamiento de modelos fundacionales, sino como un mecanismo práctico y de bajo costo para adaptar estos modelos a tareas específicas. Destacan que este enfoque hace que la alineación basada en modelos de valor sea accesible y eficiente, acelerando potencialmente la investigación hacia un modelado generativo más seguro y fiable. El trabajo sugiere que el "ensamblaje" de componentes preentrenados es una alternativa viable a los costos prohibitivos de entrenar funciones de valor desde cero o depender de aproximaciones sesgadas.