Leveraging Gauge Freedom for Learning Non-Gradient Population Dynamics of Stochastic Systems

Resumen Técnico: Aprovechamiento de la Libertad de Gauge para el Aprendizaje de Dinámicas de Población No Gradientes de Sistemas Estocásticos

1. Enunciado del Problema

El artículo aborda el desafío de inferir dinámicas a nivel de población para sistemas estocásticos donde las trayectorias individuales no se observan o son intrínsecamente impredecibles (por ejemplo, sistemas caóticos o turbulentos). La entrada consiste en muestras independientes extraídas de distribuciones marginales en el tiempo $\rho_t$ en puntos discretos de tiempo $t \in [0, T]$. El objetivo es aprender un campo de velocidad dependiente del tiempo $u_t: X \to \mathbb{R}^d$ tal que el flujo inducido por $u_t$ transporte la distribución inicial $\rho_0$ a la secuencia observada de marginales $(\rho_t)_{t \in [0, T]}$.

Una dificultad fundamental en este escenario es que la ecuación de continuidad, $\partial_t \rho_t + \nabla \cdot (\rho_t u_t) = 0$, proporciona solo una restricción subdeterminada sobre el campo de velocidad. La ecuación restringe únicamente la divergencia ponderada por $\rho_t$ del flujo, dejando las componentes libres de divergencia (solenoide) del campo de velocidad sin restricciones. Esta "libertad de gauge" implica que infinitos campos de velocidad pueden producir la misma evolución de las marginales en el tiempo.

Las metodologías existentes a menudo resuelven esta no unicidad codificando en el código una elección específica de gauge: minimizar la energía cinética. Este principio de selección, arraigado en la teoría del transporte óptimo, fuerza al campo de velocidad a ser un gradiente de un potencial escalar ($u_t = \nabla s_t$). Aunque óptimo en términos de energía cinética, esta restricción de gradiente impone un sesgo inductivo que puede ser perjudicial cuando las dinámicas subyacentes involucran componentes rotacionales significativas (por ejemplo, sistemas hamiltonianos o fluidos incompresibles). Imponer una estructura de gradiente sobre dinámicas rotacionales puede resultar en potenciales altamente oscilatorios y patológicos que son difíciles de representar y aprender, lo que conduce a una pobre precisión distribucional.

2. Metodología: Flujos de Inferencia No Gradientes (NGIF)

Los autores proponen Flujos de Inferencia No Gradientes (NGIF), un algoritmo que desacopla la tarea de ajustar las marginales en el tiempo de la tarea de seleccionar un campo de velocidad específico.

Formulación Débil de la Ecuación de Continuidad:
En lugar de imponer la ecuación de continuidad punto por punto, NGIF utiliza una formulación débil (basada en momentos). Para un conjunto de funciones de prueba $\Phi$, el método requiere que la tasa de cambio del valor esperado de cualquier función de prueba $\phi \in \Phi$ bajo la marginal $\rho_t$ coincida con el transporte predicho por el campo de velocidad candidato:
$$\frac{d}{dt} \mathbb{E}_{x \sim \rho_t}[\phi(x)] = \mathbb{E}_{x \sim \rho_t}[\nabla \phi(x) \cdot u_t(x)]$$
Esta identidad se deriva integrando la ecuación de continuidad contra funciones de prueba y aplicando integración por partes.

Pérdida Empírica y Características de Fourier Aleatorias:
Para operacionalizar esta restricción, los autores emplean funciones de prueba de características de Fourier aleatorias. Las frecuencias $\omega$ se muestrean desde una distribución $\nu$ (utilizando anchos de banda multiescala para capturar estructuras en diferentes resoluciones), generando funciones de prueba de la forma $\sin(\omega^\top x)$ y $\cos(\omega^\top x)$.
La pérdida empírica minimiza el residuo de la ecuación de continuidad débil sobre estas funciones de prueba. Crucialmente, las derivadas temporales de los momentos se estiman mediante splines de suavizado ajustados a los momentos empíricos de los datos de entrenamiento, lo que permite el pre-cálculo y la reducción de la varianza.

Regularización Explícita de Gauge:
En lugar de restringir la clase de hipótesis a campos de gradiente, NGIF introduce un regularizador de gauge explícito $G(u)$ en el objetivo de optimización:
$$\min_{u} \mathcal{L}(u) + \lambda G(u)$$
donde $\mathcal{L}(u)$ es la pérdida de residuo de la continuidad débil. Esto permite al usuario seleccionar entre una amplia familia de sesgos estructurales. El artículo propone tres regularizadores específicos:

1. Energía Cinética ($G_{KIN}$): Minimiza $\int \|u_t\|^2 \rho_t$, recuperando el enfoque estándar basado en gradiente pero como una penalización suave en lugar de una restricción rígida.
2. Penalización de Rotacional ($G_{CURL}$): Penaliza la parte antisimétrica del Jacobiano, desalentando la rotación y sesgando hacia dinámicas tipo gradiente sin imponer una parametrización de potencial.
3. Penalización de Divergencia ($G_{DIV}$): Penaliza $\nabla \cdot u_t$, promoviendo flujos aproximadamente incompresibles. Esto es particularmente útil para el transporte tipo fluido donde las dinámicas reales son libres de divergencia.

Al utilizar una red neuronal para parametrizar $u_t$ directamente (en lugar de un potencial $s_t$), el método evita las complejidades computacionales de la retropropagación a través de las derivadas espaciales de un potencial y permite el aprendizaje de campos no gradientes.

3. Contribuciones Clave

El artículo describe tres contribuciones principales:

1. Desacoplamiento del Ajuste de Marginales y la Fijación de Gauge: Los autores formulan la inferencia de dinámicas de población utilizando restricciones de continuidad débil, separando el ajuste de las marginales en el tiempo de la selección de la estructura del campo de velocidad. La no unicidad se aborda mediante un regularizador de gauge explícito elegido por el usuario.
2. Restricciones Débiles Escalables: El método instancia las restricciones débiles utilizando funciones de prueba de características de Fourier aleatorias, generando una función de pérdida computacionalmente eficiente y escalable.
3. Aprendizaje No Gradiente: El enfoque permite el aprendizaje de campos de velocidad no gradientes, lo que mejora la representabilidad y la estabilidad. Demuestra que relajar la restricción de gradiente conduce a una inferencia de dinámicas de población más precisa en comparación con las líneas base restringidas a gradientes.

4. Resultados Experimentales

Los autores evalúan NGIF en una variedad de problemas de física de baja y alta dimensión:

• Ejemplo de Gigli (Dinámicas Rotacionales): En un escenario donde componentes gaussianas rotan sobre un círculo, los métodos basados en gradiente fallan en capturar el movimiento rotacional simple, produciendo en su lugar potenciales altamente oscilatorios. NGIF con un regularizador de rotacional recupera con precisión el campo de velocidad rotacional suave con solo un aumento marginal en la energía cinética en comparación con la solución óptima de gradiente.
• Partículas Trazadoras en Flujo Turbulento: En un flujo turbulento forzado por Kolmogorov (donde la verdad fundamental es libre de divergencia), NGIF con un regularizador de divergencia ($G_{DIV}$) produce trayectorias de partículas y reconstrucciones de densidad que coinciden estrechamente con la verdad fundamental. Las líneas base basadas en gradiente (Action Matching, DICE, HOAM) tienen un rendimiento significativamente peor.
• Inestabilidades de Vlasov-Poisson: En problemas de dos corrientes, joroba en la cola y amortiguamiento de Landau fuerte, las variantes de NGIF superan a los métodos basados en gradiente más avanzados (AM, HOAM, DICE) en la predicción de la evolución de la energía eléctrica. El regularizador de divergencia permite un movimiento de partículas más rico, mejorando la precisión de la distribución final.
• Turbulencia de Alta Dimensión: En una simulación de campo de vorticidad de 128x128, NGIF aplicado a un espacio latente captura con precisión la descomposición de la enstrofía (una medida estadística del flujo), demostrando que el flujo aprendido reproduce estadísticas salientes a nivel de población incluso cuando las trayectorias individuales son caóticas e irrecuperables.

5. Significado y Afirmaciones

El artículo afirma que la inferencia de dinámicas de población es inherentemente subdeterminada, y la elección del gauge es una decisión de modelado en lugar de una búsqueda de una solución universalmente correcta. El significado de NGIF radica en su capacidad para:

• Evitar Objetivos Patológicos: Al no codificar en el código el gauge de energía cinética mínima/gradiente, el método evita la necesidad de aprender potenciales altamente oscilatorios que surgen al forzar estructuras de gradiente sobre dinámicas rotacionales.
• Incorporar Priors Físicos: Proporciona un marco modular para incorporar conocimiento físico específico (por ejemplo, incompresibilidad) directamente a través de la elección del regularizador de gauge, sin requerir flujos de referencia o datos de velocidad a nivel de muestra.
• Mejorar la Precisión Distribucional: Los experimentos sugieren que para sistemas con transporte no potencial, los enfoques no gradientes producen una precisión distribucional superior en comparación con las líneas base restringidas a gradientes.

Los autores concluyen que, aunque la ecuación de continuidad débil restringe solo la divergencia del flujo, la selección explícita de un regularizador de gauge permite el aprendizaje de modelos dinámicos que están mejor alineados con la física subyacente del sistema, particularmente en regímenes donde los datos de trayectorias no están disponibles o no son informativos.