Resumen Técnico: Emparejamiento de Trayectorias de Primer Orden (FTM)

1. Planteamiento del Problema

Muchos sistemas dinámicos científicos son efectivamente estocásticos debido al caos, el granularismo (coarse-graining), la observación parcial o grados de libertad no resueltos. Aunque existen simuladores numéricos de alta fidelidad, la simulación tradicional repetida es computacionalmente prohibitiva para tareas de bucle externo como el diseño, el control y la cuantificación de la incertidumbre. Esto hace necesario el desarrollo de modelos sustitutos (surrogates) rápidos entrenados con datos de trayectoria.

Los enfoques existentes enfrentan limitaciones específicas en este contexto:

• Aprendizaje de Operadores (Operator Learning): Entrenados típicamente con error cuadrático medio, estos modelos predicen la media condicional. Esto colapsa la distribución de los estados futuros posibles en un único punto, fallando en capturar el transporte de la masa de probabilidad requerido para las estadísticas de conjunto (ensemble statistics).
• Dinámica de Población (Emparejamiento de Marginales): Los métodos que emparejan los marginales temporales aseguran que la masa de probabilidad correcta exista en las regiones correctas en cada paso de tiempo. Sin embargo, no especifican cómo fluye la masa a lo largo del tiempo. En consecuencia, omiten fenómenos críticos dependientes de la trayectoria, tales como circulaciones, flujos y corrientes de cruce de barreras en sistemas fuera del equilibrio estadísticamente estacionarios.
• Modelos Generativos Condicionales (Difusión/Flujo): Estos modelos capturan leyes de transición condicionales completas y variabilidad estocástica. No obstante, típicamente requieren múltiples evaluaciones de la red neuronal o pasos del solver por cada paso de tiempo físico para generar el siguiente estado, lo que conlleva altos costos de inferencia durante los despliegues (rollouts) autorregresivos.

El desafío central es desarrollar un modelo sustituto que propague conjuntos rápidamente, preserve los marginales temporales, capture el transporte de trayectoria de primer orden (corrientes) y mantenga costos de despliegue determinista bajos (idealmente una sola evaluación de la red neuronal por paso de tiempo).

2. Metodología: Emparejamiento de Trayectorias de Primer Orden (FTM)

FTM aborda estos desafíos aprendiendo directamente la velocidad de la corriente de probabilidad a partir de trayectorias estocásticas, en lugar de aprender la deriva (drift), la difusión o las funciones de puntuación (score functions) por separado.

2.1 Fundamento Teórico

FTM busca un campo de velocidad determinista $u(t, x)$ tal que la Ecuación Diferencial Ordinaria (ODE) inducida $\dot{\hat{x}}(t) = u(t, \hat{x}(t))$ preserve los marginales temporales de la Ecuación Diferencial Estocástica (SDE) subyacente y coincida con el transporte local de la masa de probabilidad.

El método se basa en la velocidad de la corriente de probabilidad $v(t, x)$, definida como:
$$v(t, x) = b(t, x) - \frac{1}{2}\Sigma(t)\nabla \log \rho(t, x)$$
donde $b$ es la deriva, $\Sigma$ es la matriz de difusión y $\rho$ es la densidad de probabilidad. Esta velocidad satisface dos propiedades clave:

1. Consistencia Global: Su flujo preserva los marginales temporales $\rho(t)$ (resolviendo la ecuación de Fokker-Planck).
2. Consistencia Local: Representa el movimiento simétrico de primer orden de las trayectorias. Al combinar incrementos hacia adelante (hacia donde va la masa) e incrementos hacia atrás (de donde vino la masa), FTM captura el flujo neto de masa a través de un punto espacio-temporal.

2.2 Función de Pérdida Escalable

La regresión directa de $v$ es imposible porque la velocidad objetivo es desconocida en los datos. FTM deriva una función de pérdida libre de simulación (simulation-free) y basada en trayectorias utilizando la identidad de Stratonovich:
$$E\left[\int_0^T \langle v_\theta(t, X(t)), v(t, X(t)) \rangle dt\right] = E\left[\int_0^T v_\theta(t, X(t)) \circ dX(t)\right]$$
Esta identidad permite que el término lineal en el objetivo de regresión se estime utilizando únicamente los incrementos de trayectoria observados y el campo de velocidad candidato $v_\theta$, sin necesidad de conocer la verdadera velocidad de corriente o de estimar la función de puntuación.

2.3 Pérdida de un Solo Paso y Análisis de Estabilidad

Para hacer la pérdida escalable para trayectorias largas, los autores introducen un enfoque de pérdida por fragmentos (chunked loss). Analizan la compensación entre el tamaño del paso de tiempo $h$ (error de discretización) y la longitud del fragmento $\tau$ (reducción de la varianza).

• Límite de Estabilidad: El análisis teórico (Proposición 1) muestra que la varianza de la pérdida está gobernada por el producto del tamaño de la muestra $N$ y el paso de tiempo $h$.
• FTM de un Solo Paso: Cuando $N$ y $h$ están equilibrados, los autores proponen una pérdida de un solo paso donde $\tau \approx h$. Esta pérdida depende solo de estados adyacentes ($X(t-h), X(t), X(t+h)$) y requiere una sola evaluación de la red neuronal por punto de datos.
• Resultado: Este enfoque evita el costo computacional del promedio de la trayectoria completa, manteniendo al mismo tiempo una señal de entrenamiento estable, siempre que la resolución temporal y el tamaño de la muestra estén suficientemente equilibrados.

3. Contribuciones Clave

1. Formulación: El artículo formula la predicción de conjuntos como el aprendizaje del transporte local de la masa de probabilidad (la velocidad de la corriente de probabilidad) en lugar de medias condicionales o flujos de emparejamiento de marginales. Esto identifica la velocidad de corriente como el sustituto determinista que preserva tanto los marginales como el transporte de trayectoria de primer orden.
2. Algoritmo: Deriva una función de pérdida escalable, de un solo paso y libre de simulación que aprende la velocidad de corriente directamente de las trayectorias. Esto evita la difícil y a menudo poco fiable estimación separada de la deriva, la difusión y las funciones de puntuación.
3. Análisis Teórico: Los autores proporcionan un análisis de estabilidad que separa el error de discretización de la varianza de muestreo, justificando el uso de la pérdida simple de un solo paso cuando la resolución temporal y los tamaños de muestra están equilibrados.
4. Desempeño Empírico: En diversos sistemas dinámicos estocásticos y ejemplos de Ecuaciones en Derivadas Parciales (PDE), FTM demuestra:
   • Un emparejamiento preciso de las estadísticas del conjunto (marginales temporales).
   • Una captura superior de las cantidades de interés basadas en la trayectoria (QoIs), tales como flujos, circulaciones y corrientes de cruce de barreras, en comparación con los métodos de solo marginales.
   • Bajo Costo de Inferencia: Logra estos resultados con solo una evaluación de la red neuronal (NFE) por paso de tiempo físico, evitando el colapso hacia la media observado en el aprendizaje de operadores y los altos costos de inferencia de los modelos generativos autorregresivos.

4. Resultados Experimentales

Los autores evaluaron FTM en cuatro ejemplos:

1. Oscilador de Duffing: Un sistema de potencial de doble pozo con forzamiento browniano.
2. Convección de Rayleigh–Bénard Caótica: Un modelo de nueve modos forzado estocásticamente.
3. Burgers Estocásticamente Forzado: Un modelo de turbulencia simplificado.
4. Turbulencia Forzada Estocásticamente: Ecuaciones de Navier–Stokes incompresibles 2D.

Hallazgos Clave:

• Precisión de la Distribución: FTM empareja la evolución de la distribución del conjunto (medida por la distancia de Wasserstein) de manera competitiva con los flujos de solo marginales y los métodos explícitos de emparejamiento de marginales (ej. DICE).
• QoIs de Trayectoria: FTM supera significativamente a los métodos de solo marginales en la captura de las QoIs basadas en la trayectoria. Para el oscilador de Duffing, minimizó los errores de la corriente de cruce de barreras; para Rayleigh–Bénard, minimizó los errores de la corriente rotacional.
• Eficiencia vs. Precisión: En los ejemplos de turbulencia, FTM emparejó las estadísticas de energía y entropía con un NFE por paso de tiempo. En contraste, los modelos de difusión autorregresivos (ARDM) requirieron hasta 100 NFEs por paso de tiempo para lograr una precisión similar. Los modelos de emparejamiento de flujo condicional destilados (AR-CFM) ofrecieron costos menores, pero resultaron en errores un orden de magnitud superiores a los de FTM.
• Robustez: Se demostró que la pérdida de un solo paso es estable a través de 100 inicializaciones aleatorias de la red neuronal.

5. Significado y Limitaciones

Significancia:
El artículo sostiene que FTM ofrece una ruta práctica para predicciones de conjuntos rápidas en sistemas estocásticos. Al aprender la velocidad de la corriente de probabilidad, cierra la brecha entre el bajo costo del aprendizaje de operadores deterministas y la alta fidelidad de los modelos generativos. Permite predicciones "conscientes de la trayectoria" (capturando flujos y circulaciones) con el costo computacional de un despliegue determinista, evitando el "colapso hacia la media" inherente al aprendizaje de operadores estándar.

Limitaciones:

• Alcance de las Estadísticas: FTM está diseñado para preservar los marginales temporales y el transporte de trayectoria de primer orden. No tiene la intención de reproducir las estadísticas de trayectoria dominadas por martingalas, tales como las distribuciones de tiempo de llegada (hitting-time) o las autocorrelaciones temporales.
• Alcance Teórico: Las garantías teóricas se establecen para SDEs de dimensión finita con ruido aditivo. Aunque el método se aplica empíricamente a PDEs de alta dimensión (Burgers, Navier-Stokes), el análisis teórico no se extiende completamente a estos entornos de dimensión infinita.
• Compensación (Trade-off): El método intercambia la capacidad de modelar las leyes de trayectoria estocásticas completas por el beneficio de la eficiencia en el despliegio determinista. Es más efectivo cuando las cantidades de interés son estadísticas de conjunto u observables de transporte tipo corriente.