Mapas de Flujo: Aprendiendo la Integral de un Modelo de Difusión para un Muestreo más Rápido

Los mapas de flujo son un nuevo enfoque para acelerar el muestreo de modelos de difusión entrenando redes neuronales para predecir toda la trayectoria desde el ruido hasta los datos de una sola vez, en lugar de integrar iterativamente direcciones tangentes. La entrada de blog de Sander Dieleman desglosa la teoría y la taxonomía.

Cómo funciona hoy el muestreo por difusión

Los modelos de difusión estándar predicen la dirección tangente (estimación denoizada) en cada nivel de ruido y toman pequeños pasos para integrar la trayectoria desde el ruido hasta los datos. Esto requiere muchas iteraciones (por ejemplo, 50–1000 pasos en DDPM). Los muestreadores deterministas como DDIM tratan el proceso como una EDO, pero aún requieren integración numérica.

Qué hacen los mapas de flujo de manera diferente

Un mapa de flujo aprende a predecir directamente cualquier punto en la trayectoria entre ruido y datos desde cualquier otro punto en la misma trayectoria. En lugar de predecir la tangente en un punto, la red genera la integral completa — la muestra limpia final dada una entrada ruidosa y un nivel de ruido objetivo. Esto significa que el muestreo se puede realizar en uno o muy pocos pasos.

Enfoques de entrenamiento

Dieleman categoriza el entrenamiento de mapas de flujo en varias estrategias basadas en la taxonomía de Boffi et al.:

• Modelos de consistencia: Imponen que la salida de la red sea la misma para todos los niveles de ruido a lo largo de una trayectoria (autoconsistencia).
• Regresión directa: Entrenan la red para predecir los datos limpios directamente a partir de una muestra ruidosa en cualquier nivel de ruido, minimizando la pérdida L2 con respecto a la verdad básica.
• Destilación: Usan un modelo de difusión preentrenado como maestro para generar pares (ruidoso, limpio), luego entrenan un estudiante de mapa de flujo para imitar la integral de la EDO del maestro.

La elección depende de si se quiere retropropagar a través del solucionador de EDO (costoso pero preciso) o evitarlo (más barato pero potencialmente menos estable).

Consideraciones prácticas

Dieleman señala que los mapas de flujo se pueden entrenar desde cero (sin maestro) aprovechando el mapeo conocido de ruido a datos en el caso de difusión gaussiana. En la práctica, los mapas de flujo ofrecen:

• Muestreo más rápido: 1–4 pasos (vs. 10–50 para difusión destilada).
• Aprendizaje basado en recompensas: La formulación integral permite la optimización basada en gradientes directos de recompensas posteriores (por ejemplo, métricas de calidad de imagen).
• Direccionabilidad: El condicionamiento en variables latentes se vuelve más simple porque toda la trayectoria de generación es una sola función.

Aplicaciones y extensiones

Los mapas de flujo se están aplicando a generación de imágenes (por ejemplo, modelos de consistencia como LCM, SDXL-Turbo), video y datos 3D. La publicación referencia una monografía reciente de Lai et al. para un tratamiento matemático más profundo.