Fine-Tuning Diffusion Models via Intermediate Distribution Shaping

Each language version is independently generated for its own context, not a direct translation.

¡Claro que sí! Imagina que los modelos de difusión (como los que crean imágenes a partir de texto) son como un artista muy talentoso pero un poco distraído. Este artista ha visto millones de cuadros y sabe pintar cosas bonitas, pero a veces comete errores: no sigue bien tus instrucciones, pinta cosas inestables o se repite demasiado.

El objetivo de este paper es enseñarle a este artista a mejorar sin tener que volver a estudiar desde cero (lo cual sería muy costoso y lento). Lo hacen con dos trucos principales que llaman P-GRAFT y Corrección de Ruido Inverso.

Aquí te lo explico con analogías sencillas:

1. El Problema: El "Ruido" y la "Búsqueda de la Perfección"

Imagina que el proceso de crear una imagen es como limpiar una ventana muy sucia.

Empiezas con una ventana totalmente cubierta de barro (ruido puro).
Paso a paso, el artista quita un poco de barro para revelar la imagen debajo.
Al final, la ventana está limpia y tienes tu imagen.

El problema es que, a veces, el artista se equivoca en los pasos finales o no sabe qué quitar exactamente para que la imagen coincida con lo que le pediste (por ejemplo, "un gato rojo").

2. La Primera Solución: P-GRAFT (El "Entrenador Intermedio")

La mayoría de los métodos anteriores intentaban corregir al artista mirando solo el resultado final. Si el gato no era rojo, decían: "¡Eso está mal, inténtalo de nuevo!". Pero esto es ineficiente porque el artista ya ha gastado mucha energía limpiando la ventana y es difícil saber en qué paso exacto se equivocó.

¿Qué hace P-GRAFT?
Imagina que en lugar de esperar a que termine de limpiar la ventana, un entrenador se para a mitad del camino (cuando la ventana está semilimpia).

El entrenador mira la imagen semilimpia y dice: "Oye, si sigues así, el gato saldrá rojo. ¡Bien hecho! Pero si sigues así, saldrá azul. ¡Cuidado!".
El artista aprende a reconocer el camino correcto antes de llegar al final.

La analogía del "Bias-Variance" (Sesgo-Variación):
El paper explica que hay un equilibrio mágico:

Si miras muy al principio (cuando todo es barro), no sabes qué imagen saldrá (mucha incertidumbre).
Si miras muy al final, el artista ya está cansado y es difícil corregirlo (el error ya está hecho).
P-GRAFT encuentra el punto dulce (el momento intermedio) donde el artista todavía tiene tiempo de corregir su rumbo, pero ya tiene suficiente información para saber hacia dónde va. Es como enseñar a un nadador a corregir su brazada cuando ya está en el agua, no cuando está en la orilla.

Resultado: El artista aprende más rápido, gasta menos energía y pinta cosas que se parecen mucho más a lo que le pediste.

3. La Segunda Solución: Corrección de Ruido Inverso (El "Mapa del Tesoro")

Esta parte es para un tipo de modelo llamado "Flow Models" (Modelos de Flujo), que son como un tobogán.

La idea es que si lanzas una pelota desde arriba (ruido), debe llegar al fondo (imagen) siguiendo una línea recta y perfecta.
Pero a veces, el tobogán está un poco torcido o la pelota no empieza exactamente donde debería.

¿Qué hace la Corrección de Ruido Inverso?
Imagina que tienes un mapa del tobogán que está un poco mal dibujado. En lugar de intentar arreglar todo el tobogán (que es enorme y costoso), haces algo inteligente:

Tomas una imagen perfecta que ya existe (el tesoro).
La subes de nuevo por el tobogán (hacia atrás) para ver dónde debería haber empezado la pelota.
Te das cuenta: "¡Ah! La pelota no debería haber empezado en el punto A, sino en el punto B".
Creas un nuevo punto de partida (un "ruido corregido") que, cuando baja por el tobogán viejo, llega exactamente a la imagen perfecta.

La magia:
No necesitas cambiar el tobogán (el modelo grande). Solo necesitas enseñarle a la pelota dónde empezar.

Es como si le dieras a un conductor un GPS que le dice: "No empieces en la calle X, empieza en la calle Y para llegar al destino sin tráfico".
Esto mejora la calidad de la imagen y, además, es más rápido porque el tobogán sigue siendo el mismo, solo que la trayectoria es más eficiente.

¿Por qué es importante todo esto?

Ahorro de dinero y energía: En lugar de entrenar a un artista gigante desde cero (que cuesta millones de dólares en electricidad), solo le das unos "tutoriales intermedios" o le corriges el punto de partida.
Mejores resultados: Las imágenes generadas son más fieles a lo que pides (mejor alineación texto-imagen) y tienen menos errores (menos moléculas inestables, mejores diseños).
Versatilidad: Funciona para dibujar cuadros, diseñar planos de casas, crear moléculas para medicina y mucho más.

En resumen

El paper dice: "No intentes arreglar todo el proceso de creación desde el principio ni desde el final. Intervén en el momento justo (P-GRAFT) o corrige el punto de partida (Corrección de Ruido Inverso) para obtener resultados increíbles con mucho menos esfuerzo."

Es como si en lugar de reescribir todo un libro para corregir un error, solo cambiaras la primera frase para que el resto de la historia fluya perfectamente.

Each language version is independently generated for its own context, not a direct translation.

Aquí tienes un resumen técnico detallado del artículo "Fine-Tuning Diffusion Models via Intermediate Distribution Shaping" (Ajuste Fino de Modelos de Difusión mediante Moldeado de Distribuciones Intermedias), publicado en ICLR 2026.

1. El Problema

Los modelos generativos preentrenados, especialmente los modelos de difusión y flujo (flow models), a menudo requieren un ajuste fino (fine-tuning) para corregir errores de aprendizaje o alinearse con aplicaciones específicas mediante retroalimentación de recompensas.

Sin embargo, existen desafíos significativos en la aplicación de métodos de optimización de políticas (como PPO) a modelos de difusión:

Intratabilidad de la KL: A diferencia de los modelos autoregresivos, la verosimilitud marginal necesaria para la regularización KL (divergencia de Kullback-Leibler) en PPO es intratable en modelos de difusión.
Inestabilidad: Ignorar la regularización KL conduce a inestabilidad en entornos a gran escala, mientras que usar restricciones de KL en la trayectoria (trajectory KL) puede resultar en sesgos en la función de valor inicial y resultados subóptimos.
Ineficiencia: Los métodos basados en muestreo por rechazo (como RAFT) existentes a menudo no aprovechan la estructura temporal de los modelos de difusión, tratando el proceso como una caja negra hasta la salida final.

2. Metodología Propuesta

Los autores proponen un marco unificado que conecta el muestreo por rechazo con la maximización de recompensas regularizada por KL, y luego extiende esto para moldear distribuciones en niveles intermedios de ruido.

A. Unificación bajo GRAFT (Generalized Rejection sAmpling Fine-Tuning)

Los autores unifican variantes existentes de ajuste fino por muestreo por rechazo (como RAFT, RSO, Best-of-N) bajo un marco llamado GRAFT.

Fundamento Teórico: Demuestran que GRAFT realiza implícitamente la maximización de recompensas regularizada por KL, pero con recompensas "reconfiguradas" (reshaped rewards).
Ventaja: Esto permite lograr la restricción de KL marginal para modelos de difusión, a pesar de que la verosimilitud marginal es intratable, evitando los problemas de los métodos basados en gradientes de política.

B. P-GRAFT (Partial-GRAFT): Moldeado de Distribuciones Intermedias

La contribución central es P-GRAFT, que explota la estructura temporal de los modelos de difusión. En lugar de ajustar el modelo desde el ruido puro hasta la imagen final, P-GRAFT:

Genera trayectorias completas y calcula recompensas basadas en la imagen final ( $X_0$ ).
Realiza el muestreo por rechazo (aceptación) en un estado intermedio de la trayectoria ( $X_t$ ), no en la imagen final.
Entrena el modelo solo para los pasos de desruido desde el ruido inicial hasta el tiempo intermedio $t$ .
Utiliza el modelo de referencia (preentrenado) para completar el desruido desde $t$ hasta 0.

Justificación Matemática (Compensación Sesgo-Variabilidad):

Variabilidad: A medida que $t$ aumenta (más cerca del ruido), la varianza condicional de la recompensa dada el estado intermedio aumenta (la información sobre la recompensa final es más ruidosa).
Sesgo: Sin embargo, aprender la función de puntuación (score function) es más fácil en tiempos intermedios ( $t$ altos) porque la distribución se acerca a una Gaussiana simple (sesgo reducido).
Equilibrio: P-GRAFT encuentra un punto óptimo intermedio donde la facilidad de aprendizaje (bajo sesgo) compensa el aumento en la varianza de la recompensa.

C. Inverse Noise Correction (Corrección de Ruido Inverso)

Para modelos de flujo (flow models), que son deterministas y reversibles, los autores proponen un método para corregir errores de modelos preentrenados sin necesidad de recompensas explícitas:

Se utiliza el modelo preentrenado invertido (mediante integración hacia atrás de Euler) para mapear datos reales ( $X_0$ ) de vuelta a un espacio de "ruido inverso" ( $X_{rev}$ ).
Se observa que la distribución de este ruido inverso ( $p_{1}^{rev}$ ) difiere de la distribución de ruido estándar ( $N(0, I)$ ) debido a errores de aprendizaje o discretización.
Se entrena un pequeño adaptador ("Noise Corrector") para mapear $N(0, I)$ a esta distribución de ruido inverso corregida.
Durante la inferencia, se genera ruido corregido y se pasa por el modelo original, mejorando la calidad de salida sin reentrenar el modelo principal.

3. Contribuciones Clave

Marco GRAFT: Una unificación teórica que demuestra que el muestreo por rechazo es equivalente a la maximización de recompensas con regularización KL implícita para modelos de difusión.
P-GRAFT: Un nuevo algoritmo que ajusta el modelo solo hasta un paso intermedio, logrando mejoras empíricas significativas al equilibrar el sesgo y la varianza en el aprendizaje.
Corrección de Ruido Inverso: Un método eficiente en parámetros para mejorar modelos de flujo preentrenados corrigiendo la distribución de ruido inicial, incluso sin acceso a funciones de recompensa.
Validación Empirica: Evaluación exhaustiva en múltiples dominios: generación de texto-a-imagen, diseño de maquetas (layout), generación de moléculas y generación de imágenes incondicionales.

4. Resultados Experimentales

Generación Texto-a-Imagen (T2I):
- Se evaluó en Stable Diffusion v2 utilizando benchmarks como GenAI-Bench, T2ICompBench++ y GenEval.
- P-GRAFT superó consistentemente a los métodos de gradiente de política (como DDPO) y al modelo base.
- En GenAI-Bench, P-GRAFT (con $N_I = 0.25N$ ) logró un puntaje VQAScore de 71.94, superando a SDv2 base (66.87) y a DDPO (65.70). Esto representa una mejora relativa del 8.81% sobre el modelo base.
- P-GRAFT también mostró una mejor generalización a prompts no vistos.
Generación de Maquetas y Moléculas:
- En PubLayNet (maquetas), P-GRAFT mejoró la alineación de elementos y redujo el FID en comparación con GRAFT estándar.
- En QM9 (moléculas), P-GRAFT logró un aumento en la estabilidad de las moléculas generadas (92.61% vs 90.50% del base) utilizando 9 veces menos rondas de muestreo que GRAFT para lograr mejoras similares, demostrando una mayor eficiencia de aprendizaje.
Corrección de Ruido Inverso (Imágenes Incondicionales):
- En CelebA-HQ y LSUN-Church, el método mejoró significativamente el FID (Fréchet Inception Distance) comparado con el modelo preentrenado.
- Logró mejoras de calidad con menores FLOPs por imagen (menor costo computacional) al permitir una generación más rápida y eficiente.

5. Significado e Impacto

Este trabajo es significativo por varias razones:

Superación de PPO: Proporciona una alternativa robusta y estable a los métodos de gradiente de política (PPO) para el ajuste fino de difusión, evitando la inestabilidad y los problemas de sesgo asociados con la regularización KL en trayectorias completas.
Eficiencia Computacional: P-GRAFT demuestra que ajustar el modelo en etapas intermedias es más eficiente en términos de datos y cómputo que ajustar todo el proceso o usar métodos de política que requieren retropropagación a través de toda la trayectoria.
Nueva Perspectiva en Modelos de Flujo: La "Corrección de Ruido Inverso" introduce una nueva vía para mejorar modelos de flujo existentes sin necesidad de recompensas externas, aprovechando la reversibilidad matemática de estos modelos.
Aplicabilidad General: El marco no se limita a imágenes, sino que se extiende exitosamente a dominios complejos como la generación de moléculas y diseños, validando la teoría de compensación sesgo-varianza en contextos diversos.

En resumen, el artículo propone un cambio de paradigma en el ajuste fino de modelos generativos: en lugar de optimizar la salida final directamente, se optimiza la distribución en un punto intermedio del proceso de generación, logrando un equilibrio óptimo entre la dificultad de aprendizaje y la calidad de la recompensa.