Flow-based Models

Las GAN definen un modelo generativo implícito y no permiten evaluar directamente $p(x)$. Los VAE sí definen una verosimilitud, pero normalmente optimizan la ELBO porque marginalizar las variables latentes es intratable.

Como ya vimos, un modelo generativo típico con variables latentes viene definido por:

$$p(x) = \int p(x \mid z)\, p(z)\, dz,$$

y computar $p(x)$ de forma explícita es intratable porque es prácticamente imposible recorrer todos los valores de $z$.

Flow-based models

Los flow-based models modelan transformaciones entre una distribución base sencilla y la distribución de datos. Los normalizing flows discretos permiten evaluar densidades mediante el teorema del cambio de variable; Flow Matching aprende un campo vectorial continuo y solo proporciona densidad cuando se interpreta como un continuous normalizing flow y se integra también el cambio de volumen.

Una buena estimación de $p(x)$ permite generar muestras, inferir variables latentes o completar datos. La verosimilitud también puede aportar información para detectar datos fuera de distribución, aunque por sí sola no es un detector fiable: un modelo puede asignar mayor densidad a datos OOD que a los datos de entrenamiento.

Normalizing flows

La idea de un normalizing flow no surge de exigir que la posterior sea gaussiana. Se parte de una distribución base sencilla —por ejemplo, una gaussiana— y se aprende una transformación invertible y diferenciable que la convierte en una distribución compleja. La invertibilidad permite tanto generar muestras como evaluar la densidad del modelo.

Un normalizing flow transforma una distribución simple $p_0(z_0)$ en una compleja $p_K(z_K)$ aplicando una secuencia de funciones de transformación invertibles. "Fluyendo" a través de una cadena de transformaciones, se sustituye repetidamente la variable $z_i$ por la nueva $z_{i+1}$ para finalmente obtener una distribución de probabilidad acorde a la variable objetivo final.

Tanto la transformación directa como su inversa se pueden calcular exactamente. Esto permite realizar la estimación de densidad. Para ello, se debe tener en cuenta dos cosas:

• La densidad de la muestra transformada inversamente: para obtener la muestra transformada, se aplica la secuencia de transformaciones inversas a la muestra original. Luego, se evalúa la densidad de esta muestra transformada bajo la distribución simple original. Es decir, se comprueba si la transformación resultante es una Normal.
• El cambio de volumen debido a las transformaciones: a medida que se realizan las transformaciones, el espacio de la muestra se distorsiona. El cambio de volumen se calcula multiplicando los valores absolutos de los determinantes de las matrices jacobianas de cada transformación.

Multiplicando estos dos valores (la densidad de la muestra transformada y el cambio de volumen), se obtiene la densidad de la muestra original bajo la distribución compleja.

Aplicaciones

Las aplicaciones más directas de los normalizing flows son:

• Estimación de densidad: para calcular la densidad exacta de los datos. Se pueden aplicar para ajustar densidades multimodales a los datos observados. También pueden utilizarse como modelos híbridos que modelan la densidad conjunta de entradas y objetivos $p(x, y)$, a diferencia de los modelos de clasificación que solo modelan $p(y \mid x)$ y los modelos de densidad que solo modelan $p(x)$. Esto es útil para tareas como la detección de anomalías.
• Generación de datos: para diferentes modalidades de datos, incluyendo imágenes, vídeo, audio, texto y objetos estructurados como grafos y nubes de puntos.
• Inferencia: para modelar distribuciones posteriores variacionales en modelos de variables latentes. También se pueden utilizar para guiar simulaciones con el fin de hacer la inferencia más eficiente. Este enfoque se ha utilizado para la inferencia de modelos de simulación en cosmología^{[alsing2019fast]} y neurociencia computacional^{[goncalves2019training]}.

Conceptos básicos de álgebra lineal

Matriz jacobiana

Dada una función de mapeo de un vector de entrada $n$-dimensional a un vector de salida $m$-dimensional, la matriz de todas las derivadas parciales de primer orden de esta función se denomina matriz jacobiana:

$$\mathbf{J} = \begin{bmatrix} \dfrac{\partial f_1}{\partial x_1} & \cdots & \dfrac{\partial f_1}{\partial x_n} \\ \vdots & \ddots & \vdots \\ \dfrac{\partial f_m}{\partial x_1} & \cdots & \dfrac{\partial f_m}{\partial x_n} \end{bmatrix}$$

Se puede entender como un traductor de un espacio vectorial a otro.

Determinante

El valor absoluto del determinante (solo existe para matrices cuadradas) puede considerarse como una medida de "cuánto expande o contrae el espacio la multiplicación por la matriz".

El determinante de una matriz cuadrada $M$ detecta si es invertible:

• si $\det(M) = 0$ entonces no es invertible (una matriz singular con filas o columnas linealmente dependientes; o cualquier fila o columna toda 0);

• si $\det(M) \neq 0$, entonces $M$ es invertible.

El determinante del producto es equivalente al producto de los determinantes: $\det(AB) = \det(A)\det(B)$.

Teorema del cambio de variable

Cuando transformamos una variable aleatoria, su probabilidad no cambia, pero el espacio donde vive sí puede estirarse o comprimirse.

Imagina una variable sencilla $z$ con densidad conocida $p_Z(z)$, por ejemplo una gaussiana. Ahora aplicamos una función invertible:

$$x = f(z)$$

Como $f$ es invertible, si conocemos un valor de $x$ también podemos recuperar el valor de $z$ que lo produjo:

$$z = f^{-1}(x)$$

La pregunta es: si conocemos la densidad en el espacio original, $p_Z(z)$, ¿cómo calculamos la densidad en el nuevo espacio, $p_X(x)$?

Así de primeras se podría pensar que es:

$$p_X(x) = p_Z(f^{-1}(x))$$

Esto nos dice de dónde viene $x$, pero se deja una cosa fuera: la transformación puede haber cambiado el tamaño de las regiones del espacio.

En una dimensión, podemos pensarlo con intervalos pequeños. Un trocito alrededor de $z$ tiene longitud $dz$. Después de aplicar la transformación, ese trocito se convierte en otro alrededor de $x$ con longitud $dx$.

La probabilidad dentro de ambos trocitos debe ser la misma:

$$p_X(x)\, dx = p_Z(z)\, dz$$

Reordenando:

$$p_X(x) = p_Z(z) \left| \frac{dz}{dx} \right|$$

Y como $z = f^{-1}(x)$:

$$p_X(x) = p_Z(f^{-1}(x)) \left| \frac{df^{-1}}{dx} \right|$$

El valor absoluto aparece porque una densidad no puede ser negativa. Si la función invierte el eje, por ejemplo de izquierda a derecha, la derivada puede ser negativa, pero el factor de cambio de tamaño sigue siendo positivo.

En varias dimensiones ocurre lo mismo, pero los "trocitos" ya no son intervalos, sino pequeñas áreas, volúmenes o hipervolúmenes. Ahí entra el determinante jacobiano:

$$p_X(x) = p_Z(f^{-1}(x)) \left| \det J_{f^{-1}}(x) \right|$$

Donde $J_{f^{-1}}(x)$ es la matriz jacobiana de la función inversa. Su determinante mide cuánto se expande o contrae localmente el volumen al pasar de $x$ de vuelta a $z$.

También se suele escribir usando la transformación directa $f$:

$$p_X(x) = p_Z(z) \left| \det J_f(z) \right|^{-1}$$

Esta forma es equivalente. Si $f$ estira el volumen por un factor $3$, entonces la densidad se divide por $3$. Si lo comprime por un factor $3$, entonces la densidad se multiplica por $3$.

Aplicación a normalizing flows

Un normalizing flow no hace una única transformación, sino una cadena de transformaciones invertibles:

$$z_0 \xrightarrow{f_1} z_1 \xrightarrow{f_2} z_2 \xrightarrow{f_3} \cdots \xrightarrow{f_K} z_K = x$$

Normalmente empezamos con una distribución fácil:

$$z_0 \sim p_0(z_0)$$

Después aplicamos transformaciones hasta llegar a una muestra con aspecto de dato real:

$$x = z_K = f_K \circ f_{K-1} \circ \cdots \circ f_1(z_0)$$

Esto sirve para generar datos: muestreamos un $z_0$ sencillo y lo vamos transformando hasta obtener $x$.

Pero lo interesante de los flows es que también podemos hacer el camino inverso para calcular la densidad exacta de un dato $x$:

$$x = z_K \xrightarrow{f_K^{-1}} z_{K-1} \xrightarrow{f_{K-1}^{-1}} \cdots \xrightarrow{f_1^{-1}} z_0$$

Una vez tenemos $z_0$, su densidad es fácil de calcular porque pertenece a la distribución simple. Lo único que falta es corregir todos los cambios de volumen que se han producido por el camino.

Para una sola transformación:

$$z_i = f_i(z_{i-1})$$

Aplicando el teorema del cambio de variable:

$$p_i(z_i) = p_{i-1}(z_{i-1}) \left| \det J_{f_i}(z_{i-1}) \right|^{-1}$$

Esta ecuación dice:

• la densidad después de la transformación depende de la densidad antes de la transformación;
• si la transformación expande el espacio, la densidad baja;
• si la transformación comprime el espacio, la densidad sube.

Como comentamos en el capítulo de las GAN, en machine learning solemos trabajar con logaritmos de probabilidades debido a que:

• los productos de muchos números pequeños tienden a cero;
• los productos se convierten en sumas, que son más fáciles de optimizar.

Tomando logaritmos:

$$\log p_i(z_i) = \log p_{i-1}(z_{i-1}) - \log \left| \det J_{f_i}(z_{i-1}) \right|$$

Si repetimos esto para las $K$ transformaciones del flow, obtenemos:

$$\log p_X(x) = \log p_0(z_0) - \sum_{i=1}^{K} \log \left| \det J_{f_i}(z_{i-1}) \right|$$

Esta es la fórmula central de los normalizing flows.

Se lee así:

• $\log p_X(x)$: la log-densidad del dato real que queremos evaluar;
• $\log p_0(z_0)$: la log-densidad del punto correspondiente en la distribución simple;
• $\sum_i \log |\det J_{f_i}|$: la corrección acumulada por todos los cambios de volumen introducidos por las transformaciones.

El camino recorrido por las variables aleatorias $z_i = f_i(z_{i-1})$ es el flujo. La cadena completa de distribuciones sucesivas es lo que llamamos normalizing flow.

Modelos basados en Normalizing flows

Algunas arquitecturas populares son Real NVP^{[dinh2016density]}, Masked Autoregressive Flows^{[papamakarios2017masked]} y Glow^{[kingma2018glow]}. SurVAE^{[nielsen2020survae]} generaliza el marco para incluir transformaciones sobreyectivas y estocásticas; no es simplemente una mezcla de VAE y normalizing flows.

Las diferencias entre estas arquitecturas residen en las transformaciones que aplican y en diseños de red específicos, pero todas comparten el objetivo común de transformar una distribución simple en una más compleja.

Ventajas y desventajas

Flow Matching

Como hemos visto, los normalizing flows construyen el modelo como una cadena de transformaciones invertibles:

$$z_0 \xrightarrow{f_1} z_1 \xrightarrow{f_2} \cdots \xrightarrow{f_K} x$$

Flow Matching parte de una idea parecida, pero cambia la pregunta.

En vez de preguntar:

¿Qué transformaciones invertibles puedo diseñar para convertir ruido en datos?

pregunta:

¿Qué dirección debería seguir cada punto para moverse desde una distribución simple hasta la distribución objetivo?

Es decir, Flow Matching aprende un campo de velocidades.

Imagina que tenemos muchos puntos de ruido, por ejemplo muestras de una gaussiana, y muchos puntos reales, por ejemplo imágenes. Queremos mover poco a poco los puntos de ruido hasta que acaben pareciéndose a los datos reales.

Para describir ese movimiento introducimos una variable de tiempo:

$$t \in [0, 1]$$

Cuando $t = 0$, estamos en la distribución simple:

$$x_0 \sim p_0$$

Cuando $t = 1$, queremos estar en la distribución de datos:

$$x_1 \sim p_{\text{data}}$$

Entre medias tenemos puntos intermedios $x_t$. El modelo aprende una función:

$$v_\theta(x_t, t)$$

Esta función recibe dos cosas:

• el punto actual $x_t$;
• el tiempo actual $t$.

Y devuelve un vector que indica hacia dónde debería moverse ese punto en ese instante.

Entrenamiento

Para entrenar el modelo necesitamos saber cuál sería una buena velocidad en distintos puntos intermedios. Una forma sencilla de verlo es emparejar un punto de ruido $z$ con un dato real $x$ y trazar una línea entre ambos:

$$x_t = (1 - t)z + tx$$

Esta ecuación solo dice que $x_t$ es una interpolación:

• si $t = 0$, entonces $x_t = z$;
• si $t = 1$, entonces $x_t = x$;
• si $t$ está entre 0 y 1, entonces $x_t$ está entre el ruido y el dato.

Si el camino es una línea recta, la velocidad que lleva de $z$ a $x$ es:

$$u_t = x - z$$

Por tanto, durante el entrenamiento podemos hacer lo siguiente:

1. Tomamos una muestra de ruido $z$.
2. Tomamos un dato real $x$.
3. Elegimos un tiempo aleatorio $t$ entre 0 y 1.
4. Construimos el punto intermedio $x_t = (1 - t)z + tx$.
5. Pedimos a la red que prediga la velocidad correcta $u_t = x - z$.

La función de pérdida puede escribirse como:

$$\mathcal{L}(\theta) = \mathbb{E}_{z, x, t} \left[ \left\| v_\theta(x_t, t) - u_t \right\|^2 \right]$$

Aunque la fórmula pueda parecer densa, la idea es bastante directa: la red predice una flecha, y la penalizamos si esa flecha apunta en una dirección distinta de la que debería.

Generación de muestras

Una vez entrenado el modelo, generar una muestra consiste en resolver una ecuación diferencial ordinaria (ODE):

$$\frac{dx_t}{dt} = v_\theta(x_t, t)$$

Esto significa: "actualiza $x_t$ siguiendo la velocidad que predice la red en cada instante".

En la práctica:

1. Muestreamos un punto inicial de ruido $x_0 \sim p_0$.
2. Evaluamos la red para saber hacia dónde moverlo.
3. Damos un pequeño paso en esa dirección.
4. Repetimos el proceso desde $t = 0$ hasta $t = 1$.

Al final obtenemos $x_1$, que debería parecer una muestra de la distribución de datos.

Diferencia con Normalizing Flows

Normalizing flows y Flow Matching comparten una intuición: ambos transforman una distribución simple en una distribución compleja. La diferencia está en cómo lo hacen.

En normalizing flows discretos:

• usamos una secuencia finita de transformaciones invertibles;
• calculamos cómo cambia la densidad con determinantes jacobianos;
• podemos evaluar la densidad asignada por el modelo mediante determinantes jacobianos.

En Flow Matching:

• aprendemos un campo de velocidades continuo;
• generamos datos siguiendo una trayectoria desde ruido hasta datos;
• no necesitamos diseñar manualmente capas invertibles con determinantes fáciles.

Si el campo de velocidades define una ODE invertible, obtenemos un continuous normalizing flow. Evaluar su densidad requiere integrar la divergencia del campo junto con la trayectoria; no es un resultado gratuito del objetivo de Flow Matching.

Por eso Flow Matching resulta atractivo: permite usar redes neuronales más flexibles y entrenarlas con una pérdida de regresión relativamente simple.

Transporte Óptimo

Una vez entendida la idea central de Flow Matching, tiene sentido hablar de Optimal Transport como una forma de elegir mejores caminos entre la distribución inicial y la distribución de datos.

La intuición de Optimal Transport es transformar una distribución en otra pagando el menor coste posible.

En Flow Matching esto aparece cuando construimos caminos entre puntos de ruido y puntos reales. Si emparejamos puntos al azar, las trayectorias pueden cruzarse mucho o dar rodeos innecesarios. Si usamos una idea de transporte óptimo, intentamos emparejar puntos de forma más coherente.

Por eso en la literatura aparecen variantes como Optimal Transport Flow Matching u OT-CFM. El paper original de Flow Matching for Generative Modeling^{[lipman2022flow]} ya destaca el uso de caminos basados en transporte óptimo como una opción especialmente interesante, y trabajos posteriores como Improving and Generalizing Flow-Based Generative Models with Minibatch Optimal Transport^{[tong2023improving]} usan transporte óptimo en mini-batches para construir emparejamientos más útiles durante el entrenamiento.

Lo importante es no confundir los niveles:

• Flow Matching: aprende un campo de velocidades que mueve muestras desde ruido hacia datos.
• Optimal Transport: puede ayudar a definir caminos o emparejamientos más eficientes para entrenar ese campo de velocidades.

Relación con los modelos de difusión

Flow Matching está muy relacionado con los modelos de difusión, porque ambos describen un proceso que conecta ruido con datos.

En difusión se aprende habitualmente a invertir un proceso estocástico que añade ruido. El mismo modelo puede parametrizar ruido, score, $x_0$ o velocidad, y estas parametrizaciones se pueden transformar entre sí para una ruta de ruido dada. Además, la probability-flow ODE proporciona una descripción determinista con las mismas marginales.

Flow Matching parte directamente de un camino de probabilidad y regresa su campo de velocidades. Por tanto, ambos marcos están estrechamente relacionados: difieren en el objetivo y en cómo se especifica el camino, no en una oposición absoluta entre “ruido” y “velocidad”.

Lo veremos en el siguiente capítulo.