Resumen Técnico: Modelado Generativo en Un Paso mediante Flujos de Gradiente de Wasserstein

Enunciado del Problema

Los modelos generativos modernos, en particular los modelos de difusión y los métodos basados en flujos, han alcanzado una calidad de muestra state-of-the-art, pero sufren altos costos computacionales y latencia debido a su dependencia del muestreo iterativo de múltiples pasos. Aunque existen métodos de destilación de uno o pocos pasos, a menudo heredan señales de entrenamiento de trayectorias originalmente diseñadas para la generación iterativa o dependen de dinámicas heurísticas (por ejemplo, Modelos Derivantes) que carecen de garantías de convergencia fundamentadas y son propensos al colapso de modos. El desafío central es diseñar una dinámica de entrenamiento que permita a un generador aprender una mapeo directo y de un solo paso desde una distribución de referencia hacia una distribución de datos objetivo, asegurando al mismo tiempo una convergencia estable y una cobertura global de modos.

Metodología: W-Flow

Los autores proponen W-Flow, un marco que desacopla el diseño de las dinámicas de entrenamiento de la parametrización del generador. La idea central es prescribir la evolución de la distribución del modelo hacia la distribución objetivo utilizando Flujos de Gradiente de Wasserstein (WGF) y luego entrenar una red neuronal estática para comprimir esta evolución de múltiples pasos en un único paso de inferencia.

1. Dinámicas de Entrenamiento mediante Flujo de Gradiente de Wasserstein

En lugar de aprender una secuencia de pasos de eliminación de ruido o campos de velocidad directamente, W-Flow define la trayectoria de entrenamiento como una evolución en tiempo continuo en el espacio de medidas de probabilidad.

• Funcional de Energía: La evolución está impulsada por el descenso más pronunciado de un funcional de energía $F(q) = D(q \| p)$, donde $q$ es la distribución actual del modelo y $p$ es la distribución de datos objetivo.
• Campo de Velocidad: El campo de velocidad en tiempo continuo $V_t$ se deriva de la primera variación del funcional de energía: $V_t(x) = -\nabla \frac{\delta F}{\delta q}(q_t)(x)$.
• Discretización: El flujo continuo se discretiza utilizando un esquema de Euler explícito. En cada paso de entrenamiento $k$, la distribución del modelo $q^{(k)}$ se actualiza mediante un mapa de transporte local $g^{(k)}(x) = x + \eta V^{(k)}(x)$, donde $\eta$ es el tamaño del paso.

2. Instanciación con Divergencia de Sinkhorn

Para garantizar que el campo de velocidad sea computacionalmente tratable y geométricamente bien comportado en altas dimensiones, los autores instan el funcional de energía utilizando la divergencia de Sinkhorn ($S_\epsilon$), una métrica de transporte óptimo (OT) regularizada entrópicamente.

• Formulación del Campo de Velocidad: El campo de velocidad inducido viene dado por la diferencia entre dos proyecciones baricéntricas:
  $$V^\epsilon_{q,p}(x) = T^\epsilon_{q,p}(x) - T^\epsilon_{q,q}(x)$$
  Aquí, $T^\epsilon_{q,p}$ transporta masa desde la distribución generada $q$ hacia el objetivo $p$, mientras que $T^\epsilon_{q,q}$ representa un término de auto-transporte que impone consistencia con la distribución actual.
• Coordinación Global: A diferencia de las interacciones locales de kernel (por ejemplo, en MMD o Modelos Derivantes), el campo de velocidad de Sinkhorn se deriva de un plan de OT global, asegurando que los movimientos de las partículas estén coordinados en todo el lote para prevenir el colapso de modos.
• Estimador de Dos Lotes: Para evitar artefactos de auto-emparejamiento donde las partículas se emparejan consigo mismas con costo cero, los autores emplean una estrategia de dos lotes. Calculan el término de auto-transporte $T^\epsilon_{q,q}$ utilizando un segundo lote independiente de muestras de la distribución actual, en lugar de enmascarar la diagonal de la matriz de costos de un solo lote.

3. Entrenamiento del Generador

El generador $f_\theta$ se entrena para imitar las actualizaciones discretas de partículas definidas por el WGF.

• Objetivo: La función de pérdida minimiza la distancia cuadrada entre la salida del generador y la posición de la partícula "objetivo" derivada de la actualización de Sinkhorn:
  $$\mathcal{L}(\theta) = \mathbb{E}_{z \sim p_{ref}} \| f_\theta(z) - \text{sg}(f_\theta(z) + \eta V^\epsilon(\hat{q}_\theta, \hat{p})) \|^2$$
  donde $\text{sg}(\cdot)$ denota el operador de stop-gradient, tratando la posición actualizada de la partícula como un objetivo de regresión fijo.
• Guía sin Clasificador (CFG): El marco admite CFG modificando directamente el campo de velocidad. En lugar de mezclar distribuciones linealmente, los autores proponen guía de velocidad, que añade un término proporcional a la diferencia entre los campos de velocidad condicional y incondicional. Se demuestra que este enfoque recupera la distribución objetivo inclinada exponencialmente estándar al usar la divergencia KL y proporciona una extensión geométricamente consistente para la divergencia de Sinkhorn.

Contribuciones Clave

1. Marco Fundamentado de Un Paso: W-Flow establece un nuevo paradigma para el modelado generativo de un paso al fundamentar las dinámicas de entrenamiento en flujos de gradiente de Wasserstein, ofreciendo una interpretación variacional clara ausente en enfoques heurísticos como los Modelos Derivantes.
2. Dinámicas Basadas en Sinkhorn: La instanciación de WGF con divergencia de Sinkhorn proporciona una regla de actualización eficiente basada en transporte óptimo que captura discrepancias distribucionales globales y mejora la cobertura de modos en comparación con métodos basados en kernel o de interacción local.
3. Convergencia Teórica: Los autores demuestran que las dinámicas de partículas de muestra finita de W-Flow convergen a las dinámicas distribucionales en tiempo continuo bajo supuestos de regularidad adecuados, proporcionando una justificación teórica para el proceso de aprendizaje.
4. Estimador de Dos Lotes: Se introduce una novedosa estrategia de dos lotes para estimar términos de auto-transporte, la cual elimina el sesgo de auto-emparejamiento inherente a los estimadores de un solo lote y a las heurísticas de enmascaramiento de la diagonal.

Resultados Experimentales

Los autores evalúan W-Flow en ImageNet 256×256 y FFHQ.

• Rendimiento State-of-the-Art: W-Flow logra un nuevo FID state-of-the-art de 1.29 (escala XL) y 1.35 (escala L) para la generación condicional por clase de un solo paso en ImageNet, superando significativamente a métodos anteriores de un solo paso (por ejemplo, Modelos Derivantes, MeanFlow) y superando a muchos modelos de difusión de múltiples pasos (por ejemplo, LightningDiT, SiT) a pesar de utilizar solo una evaluación de función (NFE).
• Eficiencia de Muestreo: En comparación con modelos de difusión de múltiples pasos con puntuaciones FID similares, W-Flow produce un muestreo aproximadamente 100 veces más rápido (por ejemplo, 77.88 imágenes/segundo frente a 0.93 imágenes/segundo para SiT-XL/2 en una sola GPU H100).
• Cobertura de Modos: En experimentos con distribuciones desequilibradas (por ejemplo, 95% rostros de personas mayores, 5% rostros de niños), W-Flow captura con éxito los modos minoritarios, mientras que los Modelos Derivantes exhiben colapso de modos, ignorando por completo la clase minoritaria.
• Velocidad de Convergencia: W-Flow converge sustancialmente más rápido durante el entrenamiento, alcanzando el rendimiento de un Modelo Derivante entrenado durante 1280 épocas en solo 384 épocas.
• Transferencia de Dominio: El método demuestra una fuerte capacidad en tareas de transferencia de dominio (por ejemplo, traducción de edad en FFHQ), produciendo trayectorias de transporte más cortas y coherentes en comparación con las líneas base heurísticas.

Significado y Afirmaciones

El artículo afirma que los flujos de gradiente de Wasserstein proporcionan una base fundamentada y efectiva para el modelado generativo rápido y de alta fidelidad. Al separar el diseño de las dinámicas de entrenamiento de la parametrización del generador, W-Flow demuestra que un generador de un solo paso puede ser entrenado para lograr un rendimiento comparable a procesos iterativos complejos de múltiples pasos. Los autores enfatizan que la coordinación global proporcionada por las dinámicas de transporte óptimo es clave para estabilizar el entrenamiento y asegurar una cobertura completa de modos, abordando limitaciones críticas de los enfoques existentes de un solo paso. Los resultados sugieren que W-Flow puede servir como una alternativa escalable para aplicaciones que requieren generación de alta calidad y baja latencia, como herramientas creativas en tiempo real y simulación.