Resumen Técnico: Modelado Generativo con Ajuste de Flujo

Enunciado del Problema

El estado del arte actual en modelado generativo, particularmente en dominios de alta dimensión como la síntesis de imágenes, el diseño de proteínas y la robótica, depende en gran medida de los modelos basados en puntuación. Estos modelos aprenden la función de puntuación de Stein, $\nabla \log p_{\text{data}}(x)$, típicamente mediante ajuste de puntuación. Una vez aprendida, esta puntuación se utiliza para impulsar procesos de difusión (por ejemplo, dinámica de Langevin) para muestrear desde la distribución objetivo $p_{\text{data}}$.

Sin embargo, la función de puntuación no es el único campo vectorial capaz de generar $p_{\text{data}}$. Según la ecuación de Fokker–Planck (FPE), cualquier campo vectorial $f_\theta$ que satisfaga la condición $\nabla \cdot (p_{\text{data}} f_\theta) = \nabla \cdot (p_{\text{data}} \nabla \log p_{\text{data}})$ tendrá a $p_{\text{data}}$ como su distribución estacionaria. Esto implica que existe una familia infinita de campos vectoriales generativos (campos no conservativos) que pueden generar la misma distribución objetivo pero inducen dinámicas diferentes (por ejemplo, diferentes tasas de mezcla, patrones de circulación o dependencias dirigidas).

Los objetivos existentes de ajuste de puntuación fuerzan al modelo a converger a un único punto en este espacio ($\nabla \log p_{\text{data}}$), penalizando cualquier desviación incluso si dicha desviación preserva la distribución objetivo. Esto restringe la capacidad de codificar sesgos inductivos útiles, priores estructurales o propiedades dinámicas (como una mezcla más rápida o dependencias causales) directamente en el campo generativo.

Metodología: Ajuste de Flujo

Los autores proponen el Ajuste de Flujo, un nuevo paradigma que generaliza el ajuste de puntuación al aprender cualquier campo vectorial dentro de la familia de campos que generan la distribución objetivo, en lugar de forzar al modelo a igualar la puntuación punto por punto.

Objetivo Central

En lugar de minimizar la divergencia de Fisher entre el modelo $f_\theta$ y la puntuación $\nabla \log p_{\text{data}}$, el Ajuste de Flujo minimiza la divergencia del flujo de probabilidad. La condición para que un campo sea generativo es:
$$\nabla \cdot (p_{\text{data}} f_\theta) = \nabla \cdot (p_{\text{data}} \nabla \log p_{\text{data}})$$
Esta condición es más débil que la igualdad punto a punto, admitiendo un "nulo" de infinitos campos vectoriales válidos.

Derivación Teórica

Para optimizar esta condición, los autores definen una Divergencia de Fisher Proyectada:
$$\tilde{J}(\theta) := \mathbb{E}_{x \sim p_{\text{data}}} \left[ \| \Pi_{\text{flux}} f_\theta(x) - \nabla \log p_{\text{data}}(x) \|^2 \right]$$
donde $\Pi_{\text{flux}}$ proyecta un campo vectorial sobre el espacio de campos gradientes que comparten la misma divergencia de flujo. Aunque $\tilde{J}(\theta)$ es teóricamente sólido, es intratable en altas dimensiones.

Los autores derivan una Pérdida de Ajuste de Flujo ($L_{\text{flux}}$) tratable mediante:

1. Definir el campo residual $u_\theta = f_\theta - \nabla \log p_{\text{data}}$ y el operador de Langevin-Stein $r_\theta = \nabla \cdot u_\theta + u_\theta \cdot \nabla \log p_{\text{data}}$.
2. Utilizar simulaciones de difusión $d x_t = \nabla \log p_{\text{data}}(x_t) dt + \sqrt{2} dW_t$ para propagar $r_\theta$ a través de la distribución de datos.
3. Aplicar integración por partes y un truco de stop-gradient para convertir el objetivo en una forma que evite la retropropagación a través de derivadas costosas por trayectoria.

La pérdida final es:
$$L_{\text{flux}}(\theta) := -\mathbb{E}_{t \sim q, x_0 \sim p_{\text{data}}, x_t | x_0} \left[ \frac{1}{q(t)} u_\theta(x_0)^\top \text{sg}\left( \frac{\partial x_t}{\partial x_0}^\top \nabla_{x_t} r_\theta(x_t) \right) \right]$$
donde $\text{sg}$ denota el operador stop-gradient. Los autores demuestran que $\nabla_\theta \tilde{J}(\theta) = 2 \nabla_\theta L_{\text{flux}}(\theta)$.

Annealing de Ruido e Implementación Práctica

El método se extiende a entornos con annealing de ruido (similar a los modelos de difusión) aprendiendo un continuo de campos para niveles de ruido $\sigma$.

• Aproximación de Puntuación: Dado que $\nabla \log p_{\text{data}}$ es desconocido, los autores utilizan una aproximación no paramétrica de Estimación de Densidad de Kernel (KDE) en mini-lotes para estimar la puntuación.
• Reducción de Varianza: Para estimar el término por trayectoria $\frac{\partial x_t}{\partial x_0}^\top \nabla_{x_t} r_\theta(x_t)$, los autores emplean un estimador de cadena cruzada utilizando todas las muestras en un mini-lote, ponderadas por probabilidades de transición gaussianas, en lugar de depender de una sola cadena.

Selección de Campos Específicos

El Ajuste de Flujo aprende una familia de campos. Para seleccionar un campo específico con propiedades deseadas, los autores proponen dos estrategias:

1. Pérdida Específica de la Aplicación: Añadir pérdidas auxiliares (por ejemplo, para velocidad de mezcla o restricciones estructurales) al objetivo de Ajuste de Flujo.
2. Parametrización del Modelo: Restricciones arquitectónicas (por ejemplo, máscaras de atención causal) que imponen estructuras específicas directamente en $f_\theta$.

Contribuciones Clave

1. Paradigma de Ajuste de Flujo: Introducción de un marco de modelado generativo que aprende campos vectoriales más allá de la puntuación al igualar la divergencia del flujo de probabilidad.
2. Derivación Eficiente de Pérdida: Derivación de una pérdida de Ajuste de Flujo escalable que preserva la geometría de la divergencia de Fisher mientras permite campos vectoriales no conservativos.
3. Extensión con Annealing de Ruido: Extensión del método a distribuciones con annealing de ruido, habilitando la aplicación a conjuntos de datos de imágenes complejos y de alta dimensión.
4. Flexibilidad Demostrada: Evidencia empírica que muestra que el Ajuste de Flujo habilita aplicaciones imposibles bajo el ajuste de puntuación estándar, incluyendo:
   • Mezcla Más Rápida: Aprender campos vectoriales que aceleran la convergencia del muestreo.
   • Interpretabilidad: Ajustar modelos estructurados y biológicamente significativos (por ejemplo, velocidad de ARN) sin aproximaciones neuronales de caja negra.
   • Dinámicas Estructuradas: Incorporar dependencias dirigidas (por ejemplo, autorregresión temporal) directamente en el campo generativo.

Resultados

El artículo evalúa el Ajuste de Flujo en cinco configuraciones:

1. Campos Controlables (Datos de Juguetes): En una mezcla gaussiana 2D, el Ajuste de Flujo asigna pérdida cero a campos que preservan la distribución objetivo pero tienen dinámicas diferentes (por ejemplo, rotación, circulación), whereas el ajuste de puntuación penaliza estas desviaciones.
2. Campos Interpretables (Velocidad de ARN): Aplicado a genómica de células individuales, el Ajuste de Flujo ajusta los parámetros del modelo dinámico scVelo mediante descenso de gradiente. Supera al ajuste estándar basado en EM de scVelo en consistencia y corrección a través de límites en cinco conjuntos de datos, demostrando que las ganancias provienen del procedimiento de optimización y no de una mayor expresividad del modelo.
3. Generación Sin Restricciones (Imágenes): En CIFAR-10 y CelebA, el Ajuste de Flujo sirve como objetivo independiente. Si bien logra puntuaciones FID competitivas, actualmente es 3–4$\times$ más lento y utiliza 2–3$\times$ más memoria durante el entrenamiento en comparación con el Ajuste de Puntuación de Eliminación de Ruido (DSM). Sin embargo, escala a distribuciones de alta dimensión.
4. Mezcla Rápida: Al optimizar un proxy de mezcla, el Ajuste de Flujo aprende campos que alcanzan puntuaciones FID razonables con significativamente menos pasos de muestreo que el Ajuste de Flujo estándar o DSM.
5. Incorporación de Estructura: En una tarea de generación de trayectorias (resortes no lineales), el Ajuste de Flujo incorpora con éxito una máscara de atención causal para imponer autorregresión temporal. En contraste, aplicar la misma máscara a DSM degrada el rendimiento porque el Jacobiano simétrico de la función de puntuación entra en conflicto con las dependencias dirigidas.

Significado y Afirmaciones

Los autores afirman que el Ajuste de Flujo abre una "nueva dimensión" en el modelado generativo al tratar el propio campo vectorial como una elección de diseño en lugar de un objetivo fijo.

• Más Allá del Muestreo: El valor principal no es necesariamente reemplazar DSM para la generación estándar de imágenes (donde DSM está altamente optimizado), sino habilitar dinámicas con propiedades útiles.
• Flexibilidad: Permite a los practicantes imponer sesgos inductivos, priores estructurales y restricciones mecánicas directamente en el proceso generativo.
• Aplicaciones: El artículo destaca casos de uso específicos desbloqueados por esta flexibilidad: muestreo más rápido mediante dinámicas no reversibles, modelos científicos interpretables (como la velocidad de ARN) y modelos generativos con estructuras causales o temporales prescritas.

El artículo concluye que, aunque el Ajuste de Flujo incurre actualmente en mayores costos computacionales, su capacidad para aprender campos generativos estructurados y no conservativos ofrece una herramienta poderosa para aplicaciones donde las dinámicas de la generación son tan importantes como la distribución de la salida.