Resumen Técnico: Residuos de Wasserstein: Aprendizaje de Flujos de Gradiente a partir de la Dinámica de Poblaciones

Planteamiento del Problema

El artículo aborda el problema inverso de reconstruir la dinámica de poblaciones a partir de un conjunto disperso de instantáneas observacionales. En muchos campos científicos, como la biología computacional y la dinámica de multitudes, la evolución de la población se modela como un Flujo de Gradiente de Wasserstein (WGF): una curva de distribuciones de probabilidad $\rho_t$ impulsada por el descenso más pronunciado de un funcional de energía $F$. El objetivo es recuperar el funcional de energía subyacente $F$ de tal manera que su flujo de gradiente se ajuste a las distribuciones marginales observadas $\{q_t\}$ en tiempos de observación discretos $t \in \mathcal{T}_{obs}$.

Los enfoques dominantes existentes dependen del esquema Jordan–Kinderlehrer–Otto (JKO), que define los WGF como una secuencia de pasos proximales que minimizan una energía más una penalización de transporte óptimo (OT). Sin embargo, los métodos basados en JKO sufren dos limitaciones principales:

1. Inflexibilidad a la Discretización Temporal: Requieren resolver costosos problemas de OT entre instantáneas consecutivas.
2. Sensibilidad a los Intervalos: Tienen dificultades con grandes brechas temporales entre observaciones, fallando a menudo al capturar trayectorias curvas (por ejemplo, partículas moviéndose a lo largo de un valle sinusoidal) porque dependen de una interpolación de cuerda lineal entre las instantáneas.

Metodología: El Marco de Residuos de Wasserstein

Los autores proponen un enfoque basado en residuos que evita el esquema JKO y los acoplamientos de transporte óptimo. En lugar de imponer la condición JKO, imponen las ecuaciones diferenciales fundamentales que caracterizan los flujos de gradiente directamente a través de una función de pérdida no negativa.

1. La Formulación del Residuo

El artículo utiliza las formulaciones de Tangente y de Igualdad de Disipación de Energía (EDE) de los flujos de gradiente. Específicamente, se centra en la restricción de velocidad:
$$v_t(x) = -\nabla_x \frac{\delta F}{\delta \rho_t}(x)$$
donde $v_t$ es el campo de velocidad y $\frac{\delta F}{\delta \rho_t}$ es la primera variación del funcional de energía.

Los autores definen un residuo de velocidad $R_{vel}$ que es cero si y solo si $\rho$ es un flujo de gradiente de $F$:
$$R_{vel}[F, (\rho, v)] = \int_0^T \int_{\mathbb{R}^D} \left\| v_t(x) + \nabla_x \frac{\delta F}{\delta \rho_t}(x) \right\|^2 \rho_t(x) \, dx \, dt$$

El objetivo global combina este residuo con una divergencia de ajuste de datos $D$ en los tiempos observados:
$$\mathcal{L}(F, \rho, v) = \lambda R_{vel}[F, \rho, v] + \sum_{t \in \mathcal{T}_{obs}} D(\rho_t, q_t)$$
Este marco unifica métodos existentes como Path-Finding y Action Matching como instancias de un único principio de residuo.

2. El Algoritmo "Stitching"

Para minimizar este objetivo, los autores introducen Stitching, un método basado en partículas y libre de simulación con las siguientes características:

• Trayectorias Aprendibles: A diferencia de los métodos JKO donde la curva $\rho$ es implícita, Stitching parametriza explícitamente la curva $\rho_\theta$ como una Estimación de Densidad de Kernel (KDE) sobre un conjunto de $N$ trayectorias de partículas aprendibles $x_{t,k}^\theta$.
• Libre de Simulación: El método no requiere resolver Neural ODEs o integrar SDEs. Las posiciones y velocidades de las partículas son parámetros aprendibles directos.
• Aproximación Empírica: Para hacer la computación eficiente, la KDE se aproxima mediante una medida empírica en los centros de las partículas. Esto simplifica el residuo de velocidad a una suma sobre las partículas:
  \mathcal{L}_{stitch} \approx \int_0^T \sum_{k=1}^N w_k \left\| \dot{x}_{t,k} + \nabla_x \frac{\delta F_\theta}{\delta \rho_\theta_t}(x_{t,k}) \right\|^2 dt + \sum_{t \in \mathcal{T}_{obs}} D(\rho_\theta_t, q_t)
• Parametrización Funcional: El funcional de energía $F_\theta$ se parametriza como una suma de un término de potencial $V_\theta$, un término de entropía (log-densidad) y un término de interacción $W_\theta$, todos representados por redes neuronales.

Contribuciones Clave

1. Marco de Residuos: Los autores reformulan el aprendizaje de WGF como la minimización de un término de residuo no negativo, proporcionando una perspectiva unificada que subsume a Path-Finding y Action Matching.
2. Algoritmo de Stitching: Introducen un nuevo método de partículas, libre de simulación, que trata la curva de trayectoria $\rho$ como una variable aprendible de primer orden. Esto permite al método tolerar grandes brechas entre instantáneas y recuperar trayectorias curvas que los predictores de cuerda basados en JKO omiten.
3. Rendimiento de Vanguardia: El método logra resultados superiores en las pruebas de referencia de inferencia de trayectorias, particularmente en escenarios con datos dispersos o no emparejados.

Resultados Experimentales

El artículo evalúa Stitching en tres bancos de pruebas principales:

1. Ejemplo Ilustrativo (Flujo Continuo vs. JKO)

En un potencial sintético de "valle ondulado", las partículas siguen una trayectoria sinusoidal.

• Resultado: Los métodos basados en JKO (específicamente JKOnet*) fallan al capturar la curvatura entre instantáneas dispersas ($t=10$ a $t=20$), produciendo interpolaciones de línea recta. Stitching recupera correctamente la trayectoria curva, lo que conduce a una recuperación significativamente mejor del potencial subyacente ($R^2 = 0.62$ frente a $0.51$).

2. Recuperación de Potencial Sintético

El método se prueba en 15 potenciales bidimensionales bajo regímenes emparejados (trayectorias correlacionadas) y no emparejados (instantáneas descorrelacionadas).

• Resultado: Stitching mantiene un rendimiento consistente en ambos regímenes. En contraste, los métodos basados en JKO se degradan significativamente o colapsan en el entorno no emparejado porque dependen de acoplamientos de OT entre instantáneas consecutivas. Los errores de Stitching se mantienen dentro de un factor de $\sim 2\times$ entre regímenes, mientras que los métodos JKO a menudo fallan en recuperar campos de gradiente informativos en el caso no emparejado.

3. Inferencia de Trayectoria de Célula Única (Conjunto de Datos de Cuerpo Embrioide)

Aplicado a datos de diferenciación de células madre embrionarias humanas (5 instantáneas a lo largo de 27 días).

• Resultado: Stitching logra el mejor rendimiento (menores distancias $W_1$ y $W_2$) en comparación con las líneas base, incluyendo Neural SDEs, TrajectoryNet y varios métodos basados en JKO. Supera a sus competidores tanto en el entrenamiento con datos completos como en un protocolo de interpolación temporal de "dejar dos fuera".

4. Recuperación de la Dinámica de Interacción

El método se utiliza para desenredar los potenciales ambientales ($V$) y los núcleos de interacción ($W$) en un sistema de McKean–Vlasov.

• Resultado: Stitching recupera con éxito tanto el potencial de confinamiento como el núcleo de interacción atractiva, logrando puntuaciones de patrón $R^2$ elevadas ($0.84$ para $V$, $0.81$ para $W$). También es significativamente más rápido por época que los métodos basados en JKO, que requieren resolver problemas de OT en cada paso.

5. Flujos No Gradientes

Los autores demuestran que el marco de residuos puede extenderse más allá de los flujos de gradiente reemplazando el término de gradiente con un campo vectorial general, rastreando con éxito la dinámica quiral (no conservativa).

Significado y Reivindicaciones

El artículo afirma que la perspectiva del Residuo de Wasserstein ofrece una alternativa más flexible y robusta al esquema JKO para aprender la dinámica de poblaciones. Al desacoplar el aprendizaje del funcional de energía de las restricciones de los acoplamientos de transporte óptimo, Stitching permite:

• Robustez a la Dispersión: La capacidad de aprender la dinámica a partir de observaciones ampliamente espaciadas sin la "pérdida de curvatura" inherente a la interpolación lineal.
• Compatibilidad con Datos No Emparejados: La capacidad de aprender de conjuntos de datos donde las identidades de las partículas no se rastrean a través del tiempo (instantáneas no emparejadas), un escenario común en la imagenología biológica.
• Eficiencia Libre de Simulación: Evitar la sobrecarga computacional de resolver ODEs o problemas de OT durante el entrenamiento.

Los autores posicionan este trabajo como una unificación de paradigmas existentes y un avance práctico para el descubrimiento científico en campos que van desde la biología de célula única hasta el comportamiento colectivo, reconociendo limitaciones como la complejidad computacional $O(N^2)$ para los términos de interacción y la falta de garantías teóricas de convergencia cuando $N \to \infty$.