Resumen Técnico: Modelo de Difusión Sin Simulación y de Tiempo Finito

Planteamiento del Problema

Los modelos de difusión generativos han alcanzado un rendimiento de vanguardia en diversos dominios, pero enfrentan un compromiso fundamental entre el entrenamiento sin simulación y la generación en tiempo finito.

• Entrenamiento sin simulación: Permite que el objetivo de entrenamiento se evalúe sin simular la ecuación diferencial estocástica (SDE) de referencia, lo que reduce significativamente los costos computacionales. Los Modelos Basados en Puntuación (SBMs) convencionales logran esto utilizando un proceso de Ornstein–Uhlenbeck (OU) como referencia. Sin embargo, el proceso OU solo se aproxima a la distribución a priori de forma asintótica ($T \to \infty$), lo que hace que la generación en tiempo finito sea ineficiente o inexacta para datos de alta dimensión.
• Generación en tiempo finito: Requiere que el proceso de referencia conecte la distribución de datos empíricos y la distribución a priori dentro de un horizonte de tiempo finito prescrito. Enfoques como el puente de Schrödinger logran esto, pero generalmente requieren la simulación de la SDE de referencia durante el entrenamiento (por ejemplo, mediante el Emparejamiento de Puntuación Implícito) debido a que la correspondiente ecuación de Fokker–Planck es analíticamente intratable.

El desafío central abordado por este trabajo es la construcción de un proceso de referencia que satisfaga simultáneamente ambas propiedades: permitir el entrenamiento sin simular la SDE de referencia y asegurar que el proceso conecte las distribuciones a priori y empírica dentro de un intervalo de tiempo fijo y finito.

Metodología

Los autores proponen un marco que revierte el procedimiento de diseño convencional. En lugar de especificar primero una SDE de referencia y luego analizar sus distribuciones inducidas, los autores prescriben primero una familia de distribuciones condicionales dependientes del tiempo tratables $\{\rho_t\}_{t \in [0,1]}$ y luego construyen una SDE de referencia que realice estas distribuciones como sus marginales.

1. Construcción del Proceso de Referencia

Sea $\mu$ la distribución empírica y $\pi$ la distribución a priori. Los autores definen una familia de distribuciones condicionales $\rho_t(z|x)$ que cumplen con:

1. Regularidad: $\rho_t$ es $C^{1,2}$ con respecto al tiempo y al espacio.
2. Tractabilidad: Se pueden extraer muestras de $\rho_t(\cdot|x)$ de forma directa.
3. Condiciones de Contorno: Las marginales inducidas $p_t(z) = \int \rho_t(z|x)\mu(x)dx$ satisfacen $p_0 = \pi$ y $p_1 = \mu$.

Dados estos supuestos, los autores derivan los coeficientes $\alpha(t, z, x)$ (deriva condicional) y $b(t, z)$ (coeficiente de difusión) tales que $\rho_t$ satisface la ecuación de Fokker–Planck:
$$\partial_t \rho_t = -\nabla_z \cdot [\alpha(t, z, x)\rho_t] + \frac{1}{2} \sum_{i,j} \partial_{z_i}\partial_{z_j} [\Gamma_{ij}(t, z)\rho_t]$$
donde $\Gamma = bb^\top$. El proceso de referencia $P$ se define entonces por la SDE:
$$dZ_t = a(t, Z_t)dt + b(t, Z_t)dW_t, \quad Z_0 \sim p_0$$
donde la deriva $a(t, z)$ es la esperanza de $\alpha$ con respecto a $\rho_t$.

2. Objetivo Sin Simulación

Al aplicar el teorema de Girsanov, la divergencia de Kullback–Leibler (KL) en el espacio de trayectorias entre el proceso de referencia $P$ y el proceso de generación $Q_\theta$ se reformula. Crucialmente, debido a que $\rho_t$ está prescrita y es tratable, el objetivo puede evaluarse sin simular la SDE de referencia:
$$L_{SF}(\theta) = \frac{1}{2n} \sum_{i=1}^n \int_0^1 dt \, \mathbb{E}_{Z \sim \rho_t(\cdot|x_i)} \left[ \| \alpha(t, Z, x_i) - s_\theta(t, Z) \|_{\Gamma(t, Z)^{-1}}^2 \right]$$
Este objetivo depende únicamente del muestreo directo de $\rho_t$ y de la deriva condicional $\alpha$, evitando la necesidad de la puntuación de la marginal $p_t$ (que típicamente requiere reversión temporal y es intratable para puentes de tiempo finito).

3. Construcciones Prácticas

El artículo proporciona construcciones específicas para dos casos:

• Prior Gaussiano: $\rho_t$ se elige como una distribución Gaussiana con media y varianza dependientes del tiempo. Esto recupera una forma similar al emparejamiento de puntuación de eliminación de ruido (denoising score matching) pero con un horizonte de tiempo finito.
• Priors No Gaussianos: Se utiliza una construcción de "push-forward". Un proceso de difusión base con distribución estacionaria $\pi$ es transformado mediante una biyección suave dependiente del tiempo $\phi^x_t$. Esto permite que el marco maneje priors de cola pesada o no gaussianos (por ejemplo, la distribución SU de Johnson) manteniendo la propiedad de no simulación.

Contribuciones Clave

1. Marco Unificado: El artículo introduce un marco general que logra simultáneamente el entrenamiento sin simulación y la generación en tiempo finito, resolviendo el compromiso inherente a los modelos de difusión convencionales.
2. Reinterpretación del Emparejamiento de Puntuación: Los autores demuestran que el emparejamiento de puntuación no es fundamental para el entrenamiento de modelos de difusión. En su lugar, emerge naturalmente como una consecuencia de revertir el tiempo de un proceso de referencia definido en la dirección de los datos hacia el prior. Al construir el proceso de referencia directamente en la dirección de la generación (del prior a los datos), el objetivo de entrenamiento se vuelve independiente de la función de puntuación derivada de la reversión temporal.
3. Conexión con el Emparejamiento de Flujo (Flow Matching): El artículo muestra que el Emparejamiento de Flujo Condicional (CFM) surge como el límite de ruido pequeño ($\epsilon \to 0$) del marco estocástico propuesto. Esto proporciona un vínculo teórico entre los modelos de difusión estocásticos y los modelos de flujo deterministas, aclarando que el CFM puede verse como un límite determinista de un objetivo de KL en el espacio de trayectorias.
4. Priors No Gaussianos: La metodología se extiende a priors no gaussianos mediante una construcción de push-forward, demostrando flexibilidad más allá del supuesto gaussiano estándar.

Resultados

Los autores realizaron experimentos numéricos en conjuntos de datos de juguete en dos dimensiones (mezcla gaussiana, espiral, tablero de ajedrez y dos lunas) utilizando tanto priors gaussianos estándar como la distribución SU de Johnson (de cola pesada).

• Rendimiento: El método propuesto aprendió con éxito la estructura cualitativa de todos los conjuntos de datos con ambos tipos de prior.
• Eficiencia: A diferencia de los SBMs, que requieren ajustar el horizonte de tiempo $T$ para equilibrar la convergencia al prior con los errores de discretización, el método propuesto opera en un intervalo fijo $[0, 1]$. Los experimentos mostraron que el método propuesto logra una generación de alta calidad sin la necesidad de ajustar el horizonte de tiempo, mientras que las elecciones inapropiadas de $T$ en los SBMs degradaron la calidad de la muestra.
• Comparación: Los resultados fueron comparables a un SBM de preservación de varianza (VP-SBM) estándar con un horizonte de tiempo grande ($T=10$), pero sin la sobrecarga computacional asociada de tiempos de simulación largos o la sensibilidad a la selección de $T$.

Significación y Reivindicaciones

El artículo afirma proporcionar una reinterpretación estructural de los modelos de difusión existentes. Al cambiar el paradigma de diseño de "especificar SDE $\to$ analizar distribuciones" a "especificar distribuciones $\to$ derivar SDE", los autores revelan que:

• La dependencia del emparejamiento de puntuación en los SBMs es un artefacto del procedimiento de reversión temporal requerido para alinear el proceso de referencia con la dirección de generación.
• La distinción entre modelos de difusión estocásticos y el emparejamiento de flujo determinista es una cuestión de nivel de ruido, siendo el CFM el límite de ruido pequeño del objetivo estocástico basado en KL propuesto.

El trabajo sugiere que el diseño de las distribuciones condicionales $\{\rho_t\}$ es un aspecto crítico, aunque poco explorado, del modelado de difusión que impacta directamente en la calidad de la generación. Los autores concluyen que su principio de construcción ofrece un camino robusto para diseñar modelos de difusión de tiempo finito eficientes sin los cuellos de botella computacionales de los enfoques convencionales.