Résumé Technique : Résidus de Wasserstein : Apprentissage des flux de gradient à partir de la dynamique des populations

Énoncé du Problème

Le papier traite du problème inverse de la reconstruction de la dynamique des populations à partir d'un ensemble épars de clichés observationnels. Dans de nombreux domaines scientifiques, tels que la biologie computationnelle et la dynamique des foules, l'évolution des populations est modélisée comme un Flux de Gradient de Wasserstein (WGF) : une courbe de distributions de probabilité $\rho_t$ pilotée par la descente la plus raide d'une fonctionnelle d'énergie $F$. L'objectif est de récupérer la fonctionnelle d'énergie sous-jacente $F$ telle que son flux de gradient s'ajuste aux distributions marginales observées $\{q_t\}$ à des instants de mesure discrets $t \in \mathcal{T}_{obs}$.

Les approches dominantes existantes reposent sur le schéma de Jordan–Kinderlehrer–Otto (JKO), qui définit les WGF comme une séquence d'étapes proximales minimisant une énergie plus une pénalité de transport optimal (OT). Cependant, les méthodes basées sur JKO souffrent de deux limitations principales :

1. Inflexibilité à la discrétisation temporelle : Elles nécessitent de résoudre des problèmes de transport optimal (OT) coûteux entre des instantanés consécutifs.
2. Sensibilité aux écarts : Elles peinent face aux grands intervalles temporels entre les observations, échouant souvent à capturer les trajectoires courbes (par exemple, des particules se déplaçant le long d'une vallée sinusoïdale) car elles reposent sur une interpolation par cordes linéaires entre les instantanés.

Méthodologie : Le Cadre des Résidus de Wasserstein

Les auteurs proposent une approche basée sur les résidus qui contourne le schéma JKO et les couplages de transport optimal. Au lieu d'imposer la condition JKO, ils imposent les équations différentielles fondamentales caractérisant les flux de gradient directement via une fonction de perte non négative.

1. La formulation du résidu

Le papier utilise les formulations de Tangente et d'Égalité de Dissipation d'Énergie (EDE) des flux de gradient. Plus précisément, il se concentre sur la contrainte de vitesse :
$$v_t(x) = -\nabla_x \frac{\delta F}{\delta \rho_t}(x)$$
où $v_t$ est le champ de vitesse et $\frac{\delta F}{\delta \rho_t}$ est la variation première de la fonctionnelle d'énergie.

Les auteurs définissent un résidu de vitesse $R_{vel}$ qui est nul si et seulement si $\rho$ est un flux de gradient 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$$

L'objectif global combine ce résidu avec une divergence d'ajustement aux données $D$ aux instants observés :
$$\mathcal{L}(F, \rho, v) = \lambda R_{vel}[F, \rho, v] + \sum_{t \in \mathcal{T}_{obs}} D(\rho_t, q_t)$$
Ce cadre unifie les méthodes existantes comme le Path-Finding et l'Action Matching en tant qu'instanciations d'un principe de résidu unique.

2. L'algorithme "Stitching"

Pour minimiser cet objectif, les auteurs introduisent Stitching, une méthode basée sur les particules et sans simulation, présentant les caractéristiques suivantes :

• Trajectoires apprenables : Contrairement aux méthodes JKO où la courbe $\rho$ est implicite, Stitching paramètre explicitement la courbe $\rho_\theta$ comme une Estimation de Densité par Noyau (KDE) sur un ensemble de $N$ trajectoires de particules apprenables $x_{t,k}^\theta$.
• Sans simulation : La méthode ne nécessite pas de résoudre des ODE neuronales ou d'intégrer des SDE. Les positions et les vitesses des particules sont des paramètres directs apprenables.
• Approximation empirique : Pour rendre le calcul efficace, la KDE est approximée par une mesure empirique aux centres des particules. Cela simplifie le résidu de vitesse en une somme sur les particules :
  \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)
• Paramétrage fonctionnel : La fonctionnelle d'énergie $F_\theta$ est paramétrée comme une somme d'un terme de potentiel $V_\theta$, d'un terme d'entropie (densité logarithmique) et d'un terme d'interaction $W_\theta$, tous représentés par des réseaux de neurones.

Contributions Clés

1. Cadre de Résidus : Les auteurs recadrent l'apprentissage des WGF comme la minimisation d'un terme de résidu non négatif, offrant une perspective unifiée qui englobe le Path-Finding et l'Action Matching.
2. Algorithme Stitching : Ils introduisent une nouvelle méthode de particules sans simulation qui traite la courbe de trajectoire $\rho$ comme une variable apprenable de premier rang. Cela permet à la méthode de tolérer de grands écarts entre les instantanés et de récupérer des dynamiques courbes que les prédicteurs de cordes basés sur JKO manquent.
3. Performance de l'État de l'Art : La méthode obtient des résultats supérieurs sur les benchmarks d'inférence de trajectoire, particulièrement dans les scénarios de données éparses ou non appariées.

Résultats Expérimentaux

Le papier évalue Stitching sur trois benchmarks principaux :

1. Exemple illustratif (Flux continu vs JKO)

Dans un potentiel synthétique de type "vallée ondulée", les particules suivent une trajectoire sinusoïdale.

• Résultat : Les méthodes basées sur JKO (spécifiquement JKOnet*) échouent à capturer la courbure entre les instantanés éparses ($t=10$ à $t=20$), produisant des interpolations linéaires. Stitching récupère correctement la trajectoire courbe, menant à une meilleure récupération du potentiel sous-jacent ($R^2 = 0,62$ contre $0,51$).

2. Récupération de Potentiel Synthétique

La méthode est testée sur 15 potentiels en deux dimensions sous régimes appariés (trajectoires corrélées) et non appariés (instantanés décorrelés).

• Résultat : Stitching maintient des performances cohérentes dans les deux régimes. En revanche, les méthodes basées sur JKO se dégradent considérablement ou s'effondrent dans le cas non apparié car elles dépendent des couplages OT entre les instantanés consécutifs. Les erreurs de Stitching restent dans un facteur de $\sim 2\times$ entre les régimes, tandis que les méthodes JKO échouent souvent à récupérer des champs de gradient informatifs dans le cas non apparié.

3. Inférence de Trajectoire de Cellule Unique (Jeu de données Embryoid Body)

Appliquée aux données de différenciation de cellules souches embryonnaires humaines (5 instantanés sur 27 jours).

• Résultat : Stitching obtient les meilleures performances (distances $W_1$ et $W_2$ les plus faibles) par rapport aux bases de comparaison incluant les Neural SDEs, TrajectoryNet et diverses méthodes basées sur JKO. Il surpasse ses concurrents tant en entraînement sur données complètes qu'en protocole d'interpolation temporelle "leave-two-out".

4. Récupération de la Dynamique d'Interaction

La méthode est utilisée pour dissocier les potentiels environnementaux ($V$) et les noyaux d'interaction ($W$) dans un système de McKean–Vlasov.

• Résultat : Stitching récupère avec succès à la fois le potentiel de confinement et le noyau d'attraction, atteignant des scores $R^2$ élevés ($0,84$ pour $V$, $0,81$ pour $W$). Il est également nettement plus rapide par époque que les méthodes basées sur JKO, qui nécessitent de résoudre des problèmes de OT à chaque étape.

5. Flux Non-Gradient

Les auteurs démontrent que le cadre de résidus peut être étendu au-delà des flux de gradient en remplaçant le terme de gradient par un champ de vecteurs général, suivant ainsi avec succès des dynamiques chirales (non conservatives).

Signification et Revendications

Le papier affirme que la perspective du Résidu de Wasserstein offre une alternative plus flexible et robuste au schéma JKO pour l'apprentissage des dynamiques de population. En découplant l'apprentissage de la fonctionnelle d'énergie des contraintes de couplage du transport optimal, Stitching permet :

• Robustesse à la parcimonie : La capacité d'apprendre des dynamiques à partir d'observations largement espacées sans la "perte de courbure" inhérente à l'interpolation linéaire.
• Compatibilité avec les données non appariées : La capacité d'apprendre à partir de jeux de données où l'identité des particules n'est pas suivie à travers le temps (instantanés non appariés), un scénario courant en imagerie biologique.
• Efficacité sans simulation : Éviter la surcharge de calcul liée à la résolution d'ODE ou de problèmes de transport optimal pendant l'entraînement.

Les auteurs positionnent ce travail comme une unification des paradigmes existants et une avancée pratique pour la découverte scientifique dans des domaines allant de la biologie cellulaire unique à la dynamique collective, tout en reconnaissant des limitations telles que la complexité computationnelle en $O(N^2)$ pour les termes d'interaction et l'absence de garanties théoriques de convergence lorsque $N \to \infty$.