Post-Hoc Guidance for Consistency Models by Joint Flow Distribution Learning

1. El Problema

Los Modelos de Difusión (DMs) han demostrado un éxito notable en la generación de imágenes de alta fidelidad, pero sufren de un alto costo computacional durante la inferencia debido a su proceso iterativo de eliminación de ruido. Para controlar la generación (equilibrando fidelidad y diversidad), se utiliza la Guía sin Clasificador (Classifier-Free Guidance, CFG). Sin embargo, CFG en DMs requiere múltiples pasos de inferencia, lo que es costoso.

Por otro lado, los Modelos de Consistencia (CMs) permiten generar imágenes en uno o pocos pasos, resolviendo el problema de la velocidad. No obstante, los métodos de guía existentes para CMs tienen limitaciones críticas:

• Dependencia de un "Maestro": Las técnicas actuales de guía para CMs (como la Distilación de Consistencia, CD) requieren la destilación de conocimiento de un DM pre-entrenado que actúe como maestro. Esto vincula la capacidad de guía a la fase de entrenamiento y limita su aplicación a modelos entrenados mediante distilación.
• Falta de flexibilidad Post-Hoc: No existe un método efectivo para añadir capacidades de guía a un CM pre-entrenado (especialmente aquellos entrenados solo con Entrenamiento de Consistencia - CT, que son puramente impulsados por datos) sin necesidad de reentrenar desde cero o tener un modelo DM maestro.

El objetivo del artículo es cerrar esta brecha: ¿Cómo habilitar la guía CFG en un CM pre-entrenado de manera post-hoc (después del entrenamiento), sin depender de un DM maestro?

2. Metodología: Joint Flow Distribution Learning (JFDL)

Los autores proponen JFDL, un método ligero de alineación que permite la guía en CMs pre-entrenados sin necesidad de un modelo maestro. La metodología se basa en los siguientes pilares teóricos y prácticos:

A. Conexión entre Modelos de Flujo (FMs) y CMs

El trabajo establece una correspondencia teórica entre los Modelos de Difusión, los Modelos de Flujo (Flow-based Models) y los CMs.

• Bajo el esquema de Explosión de Varianza (VE), tanto los CMs como los FMs aprenden un mapeo entre una distribución de datos y una distribución gaussiana.
• Los autores demuestran que el CM pre-entrenado puede actuar como un solver de Ecuaciones Diferenciales Ordinarias (ODE) para la trayectoria de difusión.

B. El Mecanismo de JFDL

JFDL funciona interpolando las direcciones de las distribuciones condicionales y no condicionales (o "nulas") durante la fase de ajuste fino (fine-tuning):

1. Solver ODE: Se utiliza el CM pre-entrenado como solver ODE para generar soluciones de trayectoria.
2. Anclaje Híbrido: Se define una solución de flujo híbrido $y^{\emptyset,t}_0$ que conecta la trayectoria condicional con la no condicional.
3. Direcciones de Desruido: Se calculan dos vectores de desruido:
   • $u_{cls}$: La dirección hacia el dato limpio condicional.
   • $u_{\emptyset}$: La dirección hacia el dato limpio no condicional (o "nulo").
4. Extrapolación: Se combina estos vectores mediante un factor de guía $\omega$:
   $$u_{\omega} = \omega \cdot u_{cls} + (1 - \omega) \cdot u_{\emptyset}$$
5. Pérdida Conjunta: Se entrena el modelo minimizando una pérdida combinada que incluye la pérdida original de Consistencia (ECT) y una nueva pérdida JFDL que fuerza al modelo a seguir la trayectoria extrapolada.

C. Análisis de "Ruido Pseudo" y Normalidad

Un hallazgo teórico crucial es que el término de "señal mixta" resultante de la interpolación de flujos, cuando se combina con el ruido gaussiano de alta varianza ($\sigma_{max}$), se comporta efectivamente como una distribución Gaussiana.

• Los autores realizan pruebas de normalidad (Shapiro-Wilk, Anderson-Darling, Kolmogorov-Smirnov) en conjuntos de datos sintéticos y CIFAR-10.
• Resultado: Confirman que la distribución del "ruido pseudo" es gaussiana en casi todos los pasos de tiempo $t$. Esto valida teóricamente que el vector $u_{\emptyset}$ actúa correctamente como el campo vectorial condicional necesario para la guía, incluso sin un solver ODE explícito para la clase "nula" ($\emptyset$).

D. Variantes: Naive vs. Random JFDL

• Naive JFDL: Requiere que el CM pre-entrenado haya aprendido explícitamente la ODE de la clase "nula" ($\emptyset$).
• Random JFDL: Inspirado en trabajos recientes que eliminan la necesidad de la clase $\emptyset$, esta variante reemplaza la etiqueta "nula" con una clase aleatoria extraída de la distribución de datos. Esto permite aplicar JFDL a modelos CT puros que nunca vieron una clase "nula" durante su entrenamiento original.

3. Contribuciones Clave

1. Propuesta de JFDL: Un nuevo método de guía post-hoc para CMs que es independiente de los Modelos de Difusión (DMs), eliminando la necesidad de un modelo maestro para la destilación.
2. Fundamentación Teórica: Proporciona una justificación teórica desde la perspectiva de los Modelos de Flujo, demostrando que la interpolación de flujos preserva la distribución marginal de datos y que el ruido resultante es gaussiano.
3. Versatilidad del Modelo: Demuestra que un CM pre-entrenado sin diseño explícito para muestreo no condicional (como los modelos CT) sigue siendo efectivo para el ajuste de guía mediante JFDL (especialmente con la variante Random).
4. Mejoras Empíricas: Logra mejoras significativas en la métrica FID (Fréchet Inception Distance) en datasets estándar, permitiendo el control de la fidelidad-diversidad en modelos de un solo paso.

4. Resultados Experimentales

Los autores evaluaron JFDL en CIFAR-10 e ImageNet 64x64, comparando modelos base (ECT) con modelos ajustados mediante JFDL.

• CIFAR-10:

  • El modelo base (ECT) obtuvo un FID de 3.40 (1 paso).
  • Con Naive JFDL, el FID mejoró a 3.29 (con $\omega=1.05$).
  • Con Random JFDL, el FID mejoró a 3.24 (con $\omega=1.15$).
  • Se observó una curva de FID vs. $\omega$ similar a la de CFG en DMs: mejora en escalas bajas, degradación en escalas altas.

• ImageNet 64x64:

  • El modelo base obtuvo un FID de 5.84 (1 paso).
  • Con Naive JFDL, el FID bajó a 4.38 (mejora significativa).
  • Con Random JFDL, el FID bajó a 4.68.
  • En muestreo de 2 pasos, JFDL también mostró mejoras en ImageNet, aunque en CIFAR-10 la ganancia perceptual fue superada por la pérdida de diversidad debido a la fuerte línea base.

• Eficiencia: El ajuste fino de JFDL consumió solo el 7.5% de los datos de entrenamiento y el 30% de las horas de GPU en comparación con el entrenamiento completo del modelo base.

• Visualización: Las imágenes generadas muestran una mayor nitidez y contraste (fidelidad) a medida que aumenta $\omega$, con una reducción predecible en la diversidad de la muestra (ej. pérdida de toppings en hot-dogs generados), replicando el comportamiento clásico de CFG.

5. Significado e Impacto

Este trabajo es el primero en demostrar que los Modelos de Consistencia pueden recibir una guía efectiva post-hoc sin depender de un maestro DM.

• Puente Tecnológico: Conecta la velocidad de inferencia de los CMs con el control de calidad de la CFG, haciendo que los modelos de un solo paso sean más prácticos para aplicaciones que requieren alta fidelidad.
• Accesibilidad: Al no requerir un DM maestro ni un proceso de destilación complejo, JFDL democratiza el uso de CMs guiados, permitiendo que modelos entrenados puramente con datos (CT) ganen capacidades de control.
• Eficiencia: Reduce drásticamente el costo computacional para obtener modelos guiados, lo cual es crucial para la adopción de modelos generativos en entornos con recursos limitados.

En resumen, JFDL resuelve una limitación fundamental en la familia de Modelos de Consistencia, permitiendo un control de generación flexible y de alta calidad sin sacrificar la velocidad de inferencia que los hace atractivos.