Flow Matching/es

Article
Topic area	generative-models
Prerequisites	Diffusion Models, Optimal Transport, Neural Ordinary Differential Equations

Resumen

Flow Matching es un marco de entrenamiento sin simulación para Continuous Normalizing Flows en el que una red neuronal se ajusta por regresión directamente sobre un campo vectorial dependiente del tiempo objetivo que transporta una distribución previa sencilla hacia la distribución de los datos. Introducido por Lipman, Chen, Ben-Hamu, Nickel y Le en 2022, generaliza y unifica varios enfoques anteriores, incluidos Score Matching y Rectified Flow, y se ha convertido en un paradigma destacado para el modelado generativo de imágenes, vídeo, audio y moléculas. En comparación con el entrenamiento clásico por máxima verosimilitud de los flujos, que requiere una costosa simulación de trayectorias, y con Denoising Diffusion Probabilistic Models, que requieren una formulación basada en ecuaciones diferenciales estocásticas, Flow Matching ofrece una alternativa determinista y conceptualmente sencilla: se elige una trayectoria de probabilidad entre el ruido y los datos, se deriva el campo vectorial que la genera y se aprende ese campo vectorial con una pérdida de error cuadrático medio.

Intuición

Un flujo normalizador continuo describe una curva en el espacio de distribuciones, transportando muestras desde una densidad de probabilidad inicial $ p_0 $ en el tiempo $ t=0 $ hacia una densidad objetivo $ p_1 $ en el tiempo $ t=1 $ mediante una ecuación diferencial ordinaria (EDO). Para cada tiempo $ t \in [0,1] $, un Vector Field $ u_t(x) $ especifica la velocidad instantánea en la posición $ x $; integrar este campo vectorial a lo largo del tiempo empuja las muestras desde la previa hasta la distribución de los datos.

El reto central es que, en general, no podemos observar directamente un campo vectorial que transporte una previa tratable hasta la distribución empírica de los datos. Flow Matching sortea este obstáculo construyendo la trayectoria por partes, condicionada a puntos de datos individuales. Para una muestra de datos fija $ x_1 $, es sencillo escribir una trayectoria suave desde una muestra de ruido hasta $ x_1 $ y leer la velocidad que la genera. Promediar estas velocidades por muestra bajo el muestreo conjunto de ruido y datos produce el campo de velocidades incondicional que conduce a toda la población desde la previa hasta los datos. La idea destacable es que ajustar por regresión una red neuronal a las velocidades condicionales recupera la incondicional en esperanza, eliminando la necesidad de evaluar nunca la densidad marginal.

Trayectorias de probabilidad y campos vectoriales

Una trayectoria de probabilidad es una familia de densidades indexada por el tiempo $ \{p_t\}_{t \in [0,1]} $, donde $ p_0 $ es una previa elegida (típicamente una Gaussian Distribution estándar) y $ p_1 $ es la distribución de los datos. Un campo vectorial $ u_t $ genera la trayectoria cuando se cumple la Continuity Equation:

$ {\displaystyle \frac{\partial p_t(x)}{\partial t} + \nabla \cdot (p_t(x)\, u_t(x)) = 0.} $

De manera equivalente, las muestras extraídas de $ p_0 $ y evolucionadas por la EDO $ dx/dt = u_t(x) $ se distribuyen según $ p_t $ en cada tiempo intermedio. Múltiples campos vectoriales pueden generar la misma trayectoria, por lo que se necesita estructura adicional (como rectitud u optimalidad respecto a un coste de transporte) para singularizar uno preferido.

Flow Matching condicional

La regresión directa sobre $ u_t $ no es viable porque $ u_t $ depende de la densidad marginal desconocida. El objetivo de Flow Matching condicional (CFM) resuelve esto condicionando a una muestra objetivo $ x_1 $. Para una trayectoria condicional elegida $ p_t(x \mid x_1) $ (por ejemplo, una gaussiana cuya media interpola linealmente desde $ 0 $ en $ t=0 $ hasta $ x_1 $ en $ t=1 $) y su campo vectorial condicional generador $ u_t(x \mid x_1) $, la pérdida es

$ {\displaystyle \mathcal{L}_{\mathrm{CFM}}(\theta) = \mathbb{E}_{t,\, x_1,\, x \sim p_t(\cdot \mid x_1)}\!\left[\, \lVert v_\theta(t, x) - u_t(x \mid x_1) \rVert^2 \right],} $

donde $ v_\theta $ es el campo vectorial aprendido, $ t $ se muestrea uniformemente en $ [0,1] $ y $ x_1 $ se muestrea de los datos. Lipman et al. demostraron que este objetivo tiene el mismo gradiente respecto a $ \theta $ que la regresión sobre la marginal $ u_t $, aun cuando la marginal es intratable. La elección de diseño crucial es la trayectoria condicional; entre las opciones populares se encuentran las trayectorias gaussianas que preservan la varianza, las trayectorias que la hacen explotar y la interpolante lineal de desplazamiento de Optimal Transport $ x_t = (1-t)\, x_0 + t\, x_1 $, que produce el objetivo de regresión sorprendentemente simple $ u_t(x \mid x_0, x_1) = x_1 - x_0 $.

Entrenamiento e inferencia

El entrenamiento sólo requiere muestrear un tiempo, un vector de ruido y un punto de datos; calcular la velocidad condicional en forma cerrada; y minimizar el error cuadrático. No es necesario simular la EDO durante el entrenamiento, no hay red de puntuación auxiliar, ni una cota inferior variacional que rastrear. Los mini-lotes consisten en tripletas independientes $ (t, x_0, x_1) $ con $ x_0 $ extraído de la previa y $ x_1 $ extraído del conjunto de datos.

En la inferencia, las muestras se generan integrando la EDO aprendida $ dx/dt = v_\theta(t, x) $ desde $ t=0 $ hasta $ t=1 $ con una condición inicial extraída de la previa. Puede usarse cualquier solucionador de EDO de caja negra; entre las opciones habituales se encuentran los métodos adaptativos de Runge-Kutta y los integradores de paso fijo de Euler o el Heun's Method. Dado que Flow Matching entrenado con interpolantes lineales (de transporte óptimo) tiende a producir trayectorias casi rectas, la generación de muestras suele requerir sólo unos pocos pasos del solucionador, en contraste con los modelos de difusión, que pueden necesitar decenas o cientos.

Variantes

Varias variantes de Flow Matching ajustan la trayectoria condicional, el acoplamiento entre $ x_0 $ y $ x_1 $ o el procedimiento de entrenamiento:

• Rectified Flow (Liu et al., 2022) entrena con la misma interpolante lineal que OT-CFM y luego reentrena iterativamente el modelo sobre sus propias trayectorias rectificadas, produciendo flujos cada vez más rectos que admiten el muestreo en uno o pocos pasos.
• Stochastic Interpolants (Albergo y Vanden-Eijnden, 2023) generalizan el marco para permitir dinámicas estocásticas, unificando el modelado generativo basado en flujos y en difusión bajo un único formalismo de interpolantes.
• Optimal Transport Conditional Flow Matching (Tong et al., 2023) reemplaza el acoplamiento independiente de $ x_0 $ y $ x_1 $ por un acoplamiento de Optimal Transport sobre mini-lote, afinando la alineación entre ruido y datos y reduciendo la curvatura de las trayectorias.
• Multisample Flow Matching (Pooladian et al., 2023) desarrolla una perspectiva afín de acoplamiento por lotes y proporciona análisis teórico de los estimadores resultantes.
• Riemannian Flow Matching extiende la construcción a datos sobre variedades, sustituyendo la interpolación euclídea por interpolación geodésica y empleando integradores de EDO conscientes de la variedad.
• Discrete Flow Matching adapta el marco a datos categóricos mediante cadenas de Markov en tiempo continuo en lugar de EDOs.

Comparación con los modelos de difusión

Los modelos de difusión y Flow Matching están estrechamente relacionados: ambos aprenden una transformación dependiente del tiempo de ruido a datos y ambos pueden plantearse como problemas de regresión contra un campo objetivo. Las diferencias residen en la elección del proceso y de la parametrización. Los modelos de difusión se formulan a través de procesos directo y reverso estocásticos y aprenden la Score Function $ \nabla \log p_t(x) $; su entrenamiento corresponde a una trayectoria gaussiana específica que preserva la varianza dentro de la familia de Flow Matching. Flow Matching es puramente determinista a nivel de la EDO, trata la trayectoria como una elección de diseño libre y parametriza la velocidad en lugar de la puntuación. Empíricamente, Flow Matching de tipo OT produce trayectorias más rectas y permite un muestreo más rápido, mientras que la estocasticidad de la difusión puede mejorar la diversidad de las muestras en ciertos regímenes. Los muestreadores de difusión basados en puntuación pueden reinterpretarse como integradores de EDO de una EDO de flujo de probabilidad, exponiendo un puente matemático preciso entre ambas familias.

Limitaciones

Flow Matching hereda las dificultades habituales de los Continuous Normalizing Flows: la integración de la EDO en la inferencia puede ser costosa cuando las trayectorias son curvas o rígidas, el cálculo exacto del log-verosimilitud requiere el Hutchinson Trace Estimator o costosas evaluaciones jacobianas, y las variedades de alta dimensión pueden requerir una elección cuidadosa de la previa para no desperdiciar capacidad de modelado. El marco también supone que se dispone de una trayectoria condicional tratable, lo que es directo en el espacio euclídeo pero más delicado en variedades, grafos o espacios discretos. La condicionamiento, Classifier-Free Guidance y la evaluación libre de verosimilitud se trasladan de la difusión a Flow Matching, pero a veces se requiere una adaptación cuidadosa porque el objeto subyacente es un campo vectorial y no una puntuación.

Aplicaciones

Flow Matching se ha aplicado a la generación de imágenes de alta resolución, incluidos modelos de texto a imagen que escalan OT-CFM a miles de millones de parámetros; a la síntesis de habla y audio, donde las trayectorias rectas permiten la generación en tiempo real; a la generación de estructuras de proteínas y moléculas en variedades $ \mathrm{SE}(3) $; y a la generación de trayectorias en robótica. Muchos sistemas generativos recientes a gran escala adoptan el entrenamiento por flujo rectificado o por OT-CFM debido a su sencillez y a su perfil de inferencia en pocos pasos.