Joint Distillation for Fast Likelihood Evaluation and Sampling in Flow-based Models

1. El Problema: La Paradoja de la Evaluación de Verosimilitud

Los modelos generativos modernos basados en difusión y flujo (Flow Matching) han demostrado un rendimiento excepcional en la generación de imágenes de alta calidad. Sin embargo, presentan una limitación crítica: la evaluación de la verosimilitud (log-likelihood) es computacionalmente prohibitiva.

• Costo de Muestreo vs. Verosimilitud: Mientras que métodos recientes de destilación han logrado reducir el muestreo a pocos pasos (ej. 2-4 pasos), la evaluación precisa de la verosimilitud sigue requiriendo la integración numérica de la divergencia a lo largo de toda la trayectoria del flujo, lo que implica cientos o miles de evaluaciones de funciones neuronales (NFEs).
• Consecuencia: Esto impide el uso de objetivos de ajuste fino basados en verosimilitud, la comparación de modelos y aplicaciones posteriores (como optimización de preferencias o RL) en estos modelos rápidos. Los métodos actuales de muestreo rápido o bien abandonan por completo el cálculo de la verosimilitud o requieren una integración costosa que anula la ventaja de la velocidad.

2. Metodología: F2D2 (Fast Flow Joint Distillation)

Los autores proponen F2D2, un marco modular que permite realizar simultáneamente muestreo rápido y evaluación rápida de verosimilitud mediante la reducción de NFEs en dos órdenes de magnitud.

Insight Fundamental

En los Flujos Normalizantes Continuos (CNF) y Flow Matching, tanto la trayectoria de muestreo como la evolución de la densidad de probabilidad dependen del mismo campo de velocidad subyacente $v_\theta$:

1. Muestreo: $\frac{d}{dt}x_t = v_\theta(x_t, t)$
2. Verosimilitud: $\frac{d}{dt}\log p_t(x_t) = -\text{div}(v_\theta(x_t, t))$

Dado que la divergencia es una derivada del mismo campo de velocidad, los autores proponen destilar conjuntamente la trayectoria de muestreo y la divergencia acumulada en un único mapa de flujo.

Arquitectura y Entrenamiento

• Mapa de Flujo Conjunto: Se parametriza un mapa de flujo unificado $\Phi_Y$ que predice simultáneamente la nueva posición $x_{t \to s}$ y el cambio en el log-likelihood $z_{t \to s}$.
• Red Neuronal: Utiliza una red base (backbone) compartida con dos cabezas de predicción separadas: una para la velocidad promedio ($u_\theta$) y otra para la divergencia promedio ($D_\theta$).
• Objetivo de Pérdida: Se optimiza una función de pérdida compuesta que combina:
  1. Pérdidas para el subsistema de muestreo (ajuste de velocidad instantánea y consistencia del mapa de flujo, ej. propiedad de semigrupo o ecuación de Euler).
  2. Pérdidas para el subsistema de verosimilitud (ajuste de la divergencia instantánea y consistencia de la divergencia acumulada).
• Estrategias de Implementación: El marco es compatible con métodos existentes como Shortcut Models (basados en la propiedad de semigrupo) y MeanFlow (basado en la ecuación de Euler). Se introduce un enfoque de entrenamiento en etapas (warm-start) donde primero se entrena la velocidad y luego se añade la cabeza de divergencia.
• Estimación de Trazas: Para evitar el costo $O(d)$ del cálculo exacto de la divergencia, se utiliza el estimador de traza de Hutchinson.

3. Aplicación: Auto-Guía de Máxima Verosimilitud

Una aplicación novedosa propuesta es la Auto-Guía de Máxima Verosimilitud. Dado que F2D2 permite evaluar la verosimilitud en pocos pasos, se puede optimizar el ruido inicial ($x_0$) antes del muestreo para maximizar la verosimilitud predicha por el modelo mismo.

• Esto actúa como una guía de inferencia ligera que no requiere modelos de recompensa externos.
• Resultado destacado: Un modelo MeanFlow de 2 pasos con esta guía supera en calidad (FID) a un modelo de Flow Matching de 1024 pasos en CIFAR-10, utilizando solo una pasada adicional hacia adelante y hacia atrás.

4. Resultados Experimentales

Los experimentos se realizaron en CIFAR-10, ImageNet 64x64 y CelebA-64, comparando F2D2 contra Flow Matching, Shortcut Models y MeanFlow.

• Precisión de Verosimilitud (NLL/BPD):
  • Los modelos base de muestreo rápido (Shortcut, MeanFlow) producen estimaciones de NLL inválidas o catastróficamente erróneas en pocos pasos.
  • F2D2 logra estimaciones de NLL calibradas y cercanas al modelo "maestro" (entrenado con 1024 pasos) incluso con 1 a 2 pasos.
  • En CIFAR-10, F2D2 reduce el error absoluto medio por muestra en la verosimilitud en un factor de 7 a 171 veces en comparación con las líneas base.
• Calidad de Muestreo (FID):
  • F2D2 mantiene o mejora la calidad de las muestras (FID) en comparación con sus contrapartes sin destilación conjunta.
  • En ImageNet 64x64, F2D2 logra FID competitivo y verosimilitud significativa en configuraciones de 2 y 4 pasos, donde los métodos anteriores fallan.
• Eficiencia: El método reduce el costo computacional de la evaluación de verosimilitud de miles de NFEs a menos de 10, haciendo viable el uso de estos objetivos en entornos de producción.

5. Contribuciones Clave y Significancia

1. Primera solución unificada: F2D2 es, según los autores, el primer método que habilita la evaluación precisa de verosimilitud en pocos pasos para modelos basados en difusión/flujo sin sacrificar la calidad de muestreo.
2. Desbloqueo de nuevos objetivos: Al hacer que la verosimilitud sea computable y rápida, permite aplicar técnicas de ajuste fino basadas en verosimilitud (como DPO, PPO, optimización de preferencias) a modelos de difusión modernos, algo que antes era prohibitivo.
3. Marco Modular: La metodología es plug-and-play y se puede aplicar a arquitecturas de muestreo rápido existentes (Shortcut, MeanFlow, LSD) con un mínimo cambio (una cabeza adicional de predicción).
4. Superación de límites: Demuestra que un modelo de 2 pasos con auto-guía puede superar a un modelo de 1024 pasos, redefiniendo el equilibrio entre velocidad y calidad en la generación de imágenes.

En resumen, F2D2 resuelve una de las principales cuellos de botella en la teoría de modelos generativos modernos, permitiendo que la velocidad de inferencia y la capacidad de inferencia estadística (verosimilitud) coexistan, lo cual es fundamental para el desarrollo futuro de modelos generativos más robustos y aplicables.