Resumen Técnico: Un Marco de Sinkhorn Local para la Reconstrucción de Distribuciones Condicionales de Campos Aleatorios Multidimensionales

1. Planteamiento del Problema

La identificación de campos aleatorios a partir de datos observacionales es un desafío crítico en la cuantificación de la incertidumbre (UQ) y el aprendizaje automático científico. Muchos sistemas físicos, como el flujo en medios porosos y el transporte turbulento, están gobernados por parámetros estocásticos que resultan en respuestas estocásticas que los modelos deterministas no pueden caracterizar adecuadamente. El objetivo principal no es meramente predecir la esperanza condicional, sino reconstruir la distribución de probabilidad condicional completa de la solución estocástica.

Aunque los modelos generativos profundos (por ejemplo, CVAEs, GANs, modelos de difusión) han mostrado resultados prometedores en el aprendizaje de distribuciones complejas, a menudo fallan en preservar la estructura geométrica de las medidas de probabilidad asociadas con estados físicos cercanos. El Transporte Óptimo (OT), específicamente la métrica de Wasserstein, ofrece una distancia físicamente significativa que permanece informativa incluso cuando las distribuciones tienen soportes disjuntos. Sin embargo, el cálculo exacto de las distancias de Wasserstein implica resolver problemas de programación lineal a gran escala, lo que se vuelve computacionalmente prohibitivo para el entrenamiento de redes neuronales, particularmente en entornos de alta dimensión. El trabajo previo de los autores introdujo un marco de $W_2$ cuadrado local para abordar la localidad, pero la dependencia de los cálculos de OT exactos seguía siendo un cuello de botella para la escalabilidad.

2. Metodología

Los autores proponen un Marco de Divergencia de Sinkhorn Local para entrenar Redes Neuronales Estocásticas (SNN) para la reconstrucción de campos aleatorios multidimensionales. La metodología integra tres componentes principales:

A. Arquitectura de la Red Neuronal Estocástica (SNN)

El modelo emplea una SNN donde los parámetros aleatorios (que representan la variable de incertidumbre $\omega$) se muestrean durante la propagación hacia adelante. Esto permite que la red genere múltiples realizaciones para una única ubicación de entrada, aproximando la distribución condicional $\mu_{\mathbf{x}}$ del campo aleatorio objetivo $\mathbf{y}(\mathbf{x}, \omega)$.

B. Emparejamiento de Distribución Local mediante Vecindarios

Para manejar la naturaleza condicional del problema sin ecuaciones gobernantes explícitas, el marco utiliza una técnica de vecindad. Para una entrada dada $\mathbf{x}_i$, se define un vecindario $B(\mathbf{x}_i, \delta)$ basado en un radio $\delta$. Las distribuciones condicionales empíricas se construyen a partir de muestras dentro de este vecindario. La función de pérdida minimiza la discrepancia entre las distribuciones empíricas de la verdad de campo (ground truth) y las predicciones de la SNN en todo el dominio de entrada.

C. Divergencia de Sinkhorn Desesgada

En lugar de la distancia $W_2$ cuadrada exacta, los autores utilizan la divergencia de Sinkhorn desesgada ($S_\varepsilon$).

• Regularización Entrópica: El algoritmo de Sinkhorn introduce un término de regularización entrópica ($\varepsilon$) al costo de OT, haciendo que el problema sea diferenciable y soluble mediante el escalado iterativo de matrices en lugar de programación lineal.
• Desesgo (Debiasing): Para eliminar el sesgo entrópico inherente al transporte regularizado (donde $S_\varepsilon(\mu, \mu) \neq 0$), los autores utilizan la definición:
  $$S_\varepsilon(\mu, \hat{\mu}) = W_\varepsilon^2(\mu, \hat{\mu}) - \frac{1}{2}W_\varepsilon^2(\mu, \mu) - \frac{1}{2}W_\varepsilon^2(\hat{\mu}, \hat{\mu})$$
• Función de Pérdida: La pérdida propuesta es la divergencia de Sinkhorn promediada sobre el dominio de entrada:
  $$\mathcal{S}_{\varepsilon, \delta}^e(\mathbf{y}_{\mathbf{x}}, \hat{\mathbf{y}}_{\mathbf{x}}) = \int_{\mathcal{D}} S_\varepsilon(\mu_{\mathbf{x}, \delta}^e, \hat{\mu}_{\mathbf{x}, \delta}^e) \nu^e(d\mathbf{x})$$
  donde $\mu_{\mathbf{x}, \delta}^e$ y $\hat{\mu}_{\mathbf{x}, \delta}^e$ son medidas empíricas construidas a partir de muestras de vecindad.

D. Análisis Teórico

El artículo establece límites de error de generalización para el marco propuesto. El análisis revela un compromiso controlado por el parámetro de regularización $\varepsilon$ y el radio del vecindario $\delta$:

• Compromiso Sesgo-Varianza: Un $\varepsilon$ pequeño se aproxima a la distancia $W_2$ exacta (alta fidelidad geométrica) pero sufre de una convergencia estadística lenta en altas dimensiones. Un $\varepsilon$ mayor mejora la eficiencia estadística y la velocidad computacional, pero introduce un sesgo de regularización.
• Maldición de la Dimensionalidad: Los límites sugieren que el término de regularización entrópica puede mitigar parcialmente la maldición de la dimensionalidad en comparación con las distancias de Wasserstein empíricas, particularmente cuando las distribuciones condicionales varían suavemente.

3. Contribuciones Clave

1. Extensión del Marco: Los autores extienden su marco anterior de transporte óptimo local de la distancia de Wasserstein exacta a la divergencia de Sinkhorn desesgada. Esto resulta en un método totalmente diferenciable, escalable y computacionalmente eficiente para entrenar SNNs.
2. Garantías Teóricas: El artículo proporciona límites de error de generalización teóricos que caracterizan explícitamente el compromiso entre el sesgo de aproximación y la eficiencia estadística. Estos límites demuestran cómo el parámetro de regularización influye en la tasa de convergencia y en la capacidad de mitigar la maldición de la dimensionalidad.
3. Validación Empírica: El marco se valida a través de tres ejemplos numéricos distintos:
   • Distribución Condicional 1D: Reconstrucción de una mezcla de Gaussianas bimodal.
   • Flujo de Darcy Estocástico: Un problema multidimensional que involucra campos de permeabilidad y correlaciones espaciales.
   • Sistemas de FitzHugh–Nagumo (FHN) Estocásticos: Una red de osciladores estocásticos no lineales acoplados.

4. Resultados

Los experimentos numéricos demuestran que el marco de Sinkhorn Local logra un equilibrio superior entre la precisión de la reconstrucción y la eficiencia computacional:

• Precisión: En el ejemplo 1D, la pérdida de Sinkhorn Local superó a las pérdidas de regresión puntual (MSE, MAE) y a otras pérdidas basadas en distribuciones (Distancia de Energía, MMD, $W_2$ local) en la reconstrucción tanto de las medias condicionales como de las varianzas.
• Eficiencia: En el benchmark de flujo de Darcy estocástico, la SNN basada en Sinkhorn logró los errores de media y varianza más bajos entre todos los métodos probados (incluyendo Regresión Gaussiana Heterocedástica, MDN, CVAE y CNF). Crucialmente, redujo significativamente el tiempo de entrenamiento en comparación con el enfoque de $W_2$ cuadrado local (308s vs. 500s) manteniendo una precisión comparable o mejor.
• Sistemas Dinámicos: Para el sistema FHN estocástico, el método reconstruyó con éxito tanto los componentes de deriva (drift) determinista como los de difusión estocástica de la dinámica. El enfoque de Sinkhorn mostró errores más bajos en las funciones de deriva y difusión aprendidas en comparación con la línea base de $W_2$ local, con una reducción modesta en el tiempo de entrenamiento.
• Robustez: Los análisis de sensibilidad indicaron que el método permanece estable bajo diferentes niveles de ruido y que un radio de vecindad y un parámetro de regularización intermedios proporcionan el compromiso óptimo entre el sesgo de aproximación y el error estadístico.

5. Significado y Reivindicaciones

El artículo afirma que el marco de Sinkhorn Local propuesto ofrece un compromiso práctico entre la fidelidad geométrica, la eficiencia estadística y la escalabilidad computacional para la cuantificación de la incertidumbre.

• Escalabilidad: Al reemplazar los cálculos de OT exactos con la divergencia de Sinkhorn, el método supera el cuello de botella computacional que anteriormente limitaba la aplicación del transporte óptimo local a sistemas estocásticos multidimensionales.
• Fidelidad Geométrica: A diferencia de los métodos basados en kernels (por ejemplo, MMD) o modelos basados en verosimilitud que pueden tener dificultades con soportes disjuntos o geometrías complejas, la divergencia de Sinkhorn preserva la estructura geométrica de las medidas de probabilidad subyacentes.
• Perspectiva Teórica: Los límites de error derivados proporcionan una justificación teórica para el uso de la regularización entrópica en el aprendizaje de campos aleatorios de alta dimensión, sugiriendo que el ajuste adecuado de $\varepsilon$ puede mitigar la maldición de la dimensionalidad.
• Aplicabilidad General: El marco se presenta como una herramienta versátil para el aprendizaje automático científico probabilístico, capaz de manejar ecuaciones diferenciales parciales estocásticas y sistemas dinámicos complejos donde solo se dispone de observaciones dispersas.

Los autores conclizan que explotar vecindarios locales con pérdidas basadas en OT es más crítico para el rendimiento que aumentar la complejidad del generador condicional, y que el enfoque de Sinkhorn Local aprovecha esto efectivamente para superar los puntos de referencia existentes de aprendizaje automático en UQ.