Resumen Técnico: Destilación de Mapas de Transición Estocástica para Inferencia Probabilística Rápida

Enunciado del Problema

Los modelos de difusión han establecido un rendimiento de vanguardia en calidad de generación, diversidad y cobertura de distribución en comparación con predecesores como GANs y VAEs. Sin embargo, este rendimiento suele ir acompañado de costos de inferencia elevados debido a la necesidad de integrar ecuaciones diferenciales ordinarias (ODEs) de tiempo continuo o ecuaciones diferenciales estocásticas (SDEs) a lo largo de muchos pasos.

Aunque existen técnicas de aceleración actuales como la destilación (por ejemplo, Modelos de Consistencia) y enfoques basados en Transporte Óptimo, a menudo enfrentan una compensación entre la calidad/diversidad de la generación y la complejidad de la inferencia. Una limitación crítica de los métodos de destilación actuales es su dependencia de la inferencia determinista (por ejemplo, a través de la ODE de flujo de probabilidad). Dada una muestra de ruido fija, la salida es única. Esto restringe su utilidad en tareas aguas abajo que inherentemente requieren inferencia estocástica, como el muestreo posterior de difusión para problemas inversos, el ajuste fino basado en energía o la mejora de la alineación con consultas de usuario. Además, muchos marcos de destilación existentes requieren un modelo profesor preentrenado, optimización de dos niveles o la simulación y almacenamiento en caché de trayectorias, lo que complica el entrenamiento y la escalabilidad.

Metodología: Destilación de Mapas de Transición Estocástica (STMD)

Los autores proponen la Destilación de Mapas de Transición Estocástica (STMD), un marco libre de profesores diseñado para acelerar la inferencia de modelos de difusión mientras preserva la estructura de transición estocástica del proceso subyacente.

Concepto Central

A diferencia de los modelos de difusión basados en puntuación que típicamente modelan la media de la distribución posterior, STMD destila el mapa de transición completo asociado con la SDE de muestreo. El método aprende un generador condicional que predice la probabilidad de transición $p(x_{t'} | x_t)$ para $t' < t$ en uno o pocos pasos, evitando la costosa integración numérica de la SDE inversa.

Formulación Técnica

1. Flujo Medio Condicional (MF): El marco parametriza las transiciones de la SDE utilizando un modelo de Flujo Medio Condicional. En lugar de aprender un mapa de flujo determinista, aprende la velocidad promedio a lo largo de la trayectoria del proceso de difusión.
2. Parametrización x0: Para modelar la transición $p(x_{t'} | x_t)$, los autores utilizan una parametrización $x_0$. Expresan la transición como una integral sobre la posterior $p(x_0 | x_t)$ y la medida de puente conocida $p(x_{t'} | x_0, x_t)$. Dado que la medida de puente para las SDEs de preservación de varianza (VP) está disponible en forma cerrada, el problema se reduce a aprender un modelo condicional $p_\theta(x_0 | x_t)$.
3. Objetivo de Entrenamiento: El modelo $u_\theta$ se entrena para aproximar la velocidad media condicional. Los autores definen una función de pérdida de Flujo Medio Condicional (CMF):
   $$L_{dCMF} = \mathbb{E} \left[ \left\| u_\theta(z_s, r, s, x_t, t) - \text{sg}\left( (z_1 - x_0) - (s-r) \frac{d}{ds}u_\theta(z_s, r, s, x_t, t) \right) \right\|^2 \right]$$
   donde $z_s$ es un interpolante de línea recta entre los datos $x_0$ y una muestra previa $z_1$, condicionado al estado de difusión $x_t$ y al tiempo $t$. El operador $\text{sg}$ denota el stop-gradient.
4. Entrenamiento Libre de Profesores: El proceso de entrenamiento no requiere un modelo profesor preentrenado. Se basa en muestras de la distribución conjunta $p(x_0, x_t)$, que pueden generarse muestreando $x_0$ desde la distribución de datos y aplicando el kernel de la SDE directa (que es conocido en forma cerrada).

Inferencia

La inferencia se realiza integrando el Flujo Medio condicional aprendido hacia atrás en el tiempo. El proceso implica:

1. Muestrear un ruido inicial $x_1$ (o $z_1$).
2. Aplicar iterativamente el campo de velocidad aprendido $u_\theta$ para avanzar desde el tiempo $s$ hasta $s-\Delta s$.
3. En cada paso de tiempo de difusión $t$, muestrear desde la distribución de puente $p(x_{t-\Delta t} | x_0, x_t)$ para actualizar el estado, simulando efectivamente la transición estocástica en grandes pasos.

Contribuciones Clave

1. Marco Novel: La propuesta de STMD, un marco de destilación basado en Flujos Medios condicionales que permite aprender mapas de transición de SDEs de difusión de alta dimensión utilizando un generador de uno o pocos pasos mediante un objetivo de regresión, sin ningún profesor preentrenado.
2. Convergencia Teórica: La derivación de límites de convergencia para el algoritmo de Flujo Medio en la distancia de Wasserstein-2. Hasta donde saben los autores, este es el primer análisis de convergencia del algoritmo de Flujo Medio en esta métrica. Extienden esto para demostrar que su mapa de transición condicional aprendido converge al núcleo de transición de la SDE inversa subyacente.
3. Validación Empírica: Demostración del método en tareas de generación de imágenes y restauración (inpainting) en los conjuntos de datos MNIST, CIFAR-10 y CelebA, mostrando un rendimiento competitivo con las líneas base mientras retiene la estocasticidad.

Resultados Experimentales

Los autores evaluaron STMD en tres conjuntos de datos:

• MNIST: Utilizando una arquitectura base U-Net, el modelo se comparó contra el algoritmo original de Flujo Medio y un Modelo Probabilístico de Difusión con Eliminación de Ruido (DDPM). Los autores utilizaron la Distancia Fréchet (FD) en el espacio latente de un clasificador (ya que FID es poco fiable para MNIST). Los resultados mostraron que STMD se desempeña de manera comparable a las líneas base mientras soporta muestreo estocástico.
• CIFAR-10: Se realizaron experimentos similares utilizando una U-Net. El modelo logró puntuaciones competitivas de Distancia Fréchet Inception (FID) en comparación con las líneas base DDPM y Flujo Medio para un número fijo de iteraciones de entrenamiento, manteniendo al mismo tiempo la estocasticidad completa del muestreador de difusión.
• CelebA: Los experimentos utilizaron una arquitectura DiT latente. El modelo generó muestras incondicionales con un FID de 8.28 (utilizando 4 pasos de inferencia y 2 pasos de Flujo Medio) y realizó con éxito tareas de restauración de imágenes.

En todos los casos, los modelos se entrenaron desde cero con las mismas arquitecturas de red base y conteos de iteraciones que las líneas base para garantizar una comparación justa.

Significado y Afirmaciones

El artículo afirma que STMD aborda la brecha entre la destilación determinista rápida y la necesidad de inferencia estocástica en tareas aguas abajo. Al destilar el mapa de transición en lugar de solo el mapa de flujo, STMD preserva la estructura estocástica del proceso de difusión.

Los autores enfatizan que su enfoque es libre de profesores, eliminando la necesidad de optimización de dos niveles o almacenamiento en caché de trayectorias, lo que hace que el proceso de entrenamiento sea más eficiente y escalable. La contribución teórica de establecer límites de convergencia de Wasserstein-2 proporciona una base rigurosa para el enfoque de Flujo Medio, cerrando una brecha entre las garantías teóricas y la aplicación práctica.

El trabajo se posiciona como una solución para tareas que requieren inferencia probabilística, como problemas inversos y muestreo posterior, donde los métodos de destilación determinista son insuficientes. Los autores se mantienen modestos respecto al rendimiento actual, señalando que el número limitado de iteraciones de entrenamiento impide resultados de vanguardia, pero el estudio comparativo valida la eficacia del método en relación con las líneas base existentes.