Neural Pushforward Samplers for the Fokker-Planck Equation on Embedded Riemannian Manifolds

1. Problématique

L'article s'attaque à la résolution numérique de l'équation de Fokker-Planck (FPE) définie sur une variété riemannienne compacte, lisse et sans bord $(M, g)$, plongée isométriquement dans un espace euclidien $\mathbb{R}^n$.

L'équation gouverne l'évolution temporelle de la densité de probabilité $\rho$ d'un processus stochastique :
$$\frac{\partial \rho}{\partial t} = \frac{\sigma^2}{2} \Delta_g \rho - \text{div}_g(\rho b)$$
où $\Delta_g$ est l'opérateur de Laplace-Beltrami, $b$ est un champ de vecteurs de dérive tangentielle, et $\sigma$ est le coefficient de diffusion.

Défis majeurs :

• Complexité géométrique : Les méthodes classiques (différences finies, éléments finis) nécessitent une triangulation explicite de la variété, ce qui devient prohibitif en dimensions modérées à élevées.
• Malédiction de la dimension : Même pour des variétés simples comme la sphère $S^{n-1}$ avec $n \ge 4$, les méthodes basées sur des grilles sont impraticables.
• Limites des méthodes existantes : Les approches par "flots normalisants" (normalizing flows) exigent souvent que la dimension de l'espace latent soit égale à celle de la variété et que la carte soit inversible, ce qui limite la capacité de représentation.

2. Méthodologie : WANPF sur Variété

Les auteurs étendent la méthode Weak Adversarial Neural Pushforward (WANPF), initialement conçue pour les espaces euclidiens plats, au cadre riemannien. La méthode repose sur trois piliers structurels :

A. Formulation Faible et Intégration par Parties

Au lieu de résoudre l'équation forte pour la densité $\rho$, la méthode utilise une formulation faible. En multipliant l'équation par une fonction test $f$ et en intégrant par parties (les termes de bord disparaissant grâce à la compacité de $M$), les opérateurs différentiels agissent sur la fonction test $f$ plutôt que sur $\rho$.
L'objectif est de satisfaire la condition d'attente :
$$\mathbb{E}_{x \sim \rho} [\mathcal{L}^{ss} f(x)] = 0$$
où $\mathcal{L}^{ss}$ est l'opérateur de Kolmogorov arrière.

B. Architecture de Poussée Contrainte (Pushforward Architecture)

La solution est représentée non pas par une densité discrète, mais par une application de poussée (pushforward map) $F_\vartheta$ qui transforme une distribution de base simple (ex: Gaussienne) en une distribution sur la variété $M$.

• Contrainte de variété : Pour garantir que tous les échantillons restent sur $M$, la sortie d'un réseau de neurones non contraint $\tilde{F}_\vartheta$ est projetée sur la variété via une rétraction $\Pi_M$ (ex: normalisation pour la sphère, enveloppement pour le tore).
  $$F_\vartheta(r) = \Pi_M(\tilde{F}_\vartheta(r))$$
• Avantage : Cette approche n'exige pas que la dimension de l'espace latent soit égale à $\dim(M)$, permettant une capacité de représentation accrue.

C. Fonctions Test : Ondes Planes Ambiantes

Au lieu d'utiliser des fonctions définies intrinsèquement sur la variété (difficiles à paramétrer), les auteurs utilisent des ondes planes ambiantes définies sur l'espace $\mathbb{R}^n$ :
$$f^{(k)}(x) = \sin(w^{(k)} \cdot x + b^{(k)})$$
Le défi majeur résolu dans cet article est le calcul de l'action de l'opérateur de Laplace-Beltrami $\Delta_g$ sur ces ondes planes. Les auteurs dérivent une formule fermée ne nécessitant pas de différenciation automatique (autodiff) à travers la géométrie intrinsèque :
$$\Delta_g f^{(k)}|_M = -\sin(\phi) \, w^\top P(x) w - \cos(\phi) \, H(x) \cdot w$$
où :

• $P(x)$ est la matrice de projection orthogonale sur l'espace tangent.
• $H(x)$ est le vecteur de courbure moyenne de l'immersion.
• $\phi = w \cdot x + b$.

Cette formule permet d'évaluer les termes de l'équation faible uniquement à l'aide de quantités ambiantes calculables analytiquement pour des variétés spécifiques (sphère, tore).

D. Entraînement Adversarial

Le problème est formulé comme un jeu min-max :

• Le Générateur (Minimiseur) : Ajuste les paramètres $\vartheta$ de $F_\vartheta$ pour minimiser l'erreur de la formulation faible.
• L'Adversaire (Maximiseur) : Ajuste les paramètres des fonctions test $\{w^{(k)}, b^{(k)}\}$ pour maximiser l'erreur, agissant comme un critique qui détecte les écarts par rapport à la solution stationnaire ou temporelle.

3. Contributions Clés

1. Formulation WANPF Stationnaire et Temporelle : Extension complète de la méthode aux équations de Fokker-Planck sur variétés, avec des formulations explicites pour les états stationnaires et dépendants du temps.
2. Formules de Laplace-Beltrami Fermées : Dérivation de formules analytiques pour l'action de $\Delta_g$ sur les ondes planes pour la sphère $S^{n-1}$ et le tore plat $T^n$, éliminant le besoin de calculs de gradients géométriques coûteux via l'autodiff.
3. Architecture de Rétraction : Conception d'un générateur qui garantit par construction que les échantillons appartiennent à la variété, évitant ainsi les problèmes de singularités de cartes locales.
4. Preuve Numérique : Implémentation réussie sur un problème d'équation de Fokker-Planck à double puits sur la sphère $S^2$.

4. Résultats Numériques

Les auteurs ont testé la méthode sur une sphère unité $S^2 \subset \mathbb{R}^3$ avec un potentiel à double puits :
$$V(x, y, z) = \alpha(x^2 - 1)^2 + \beta z^2$$

• Configuration : Le potentiel crée deux minima stables en $(\pm 1, 0, 0)$ et confine la masse près de l'équateur grâce au terme $\beta z^2$.
• Résultats :
  • Le générateur entraîné concentre correctement la masse de probabilité autour des deux minima du potentiel, en accord avec la distribution de Gibbs théorique $\rho_\infty \propto e^{-2V/\sigma^2}$.
  • La distribution est fortement supprimée aux pôles, respectant la contrainte du potentiel.
  • La symétrie du problème est préservée.
  • La perte d'entraînement converge vers un niveau stationnaire après environ 2 000 étapes, démontrant la stabilité du jeu min-max.
• Efficacité : La méthode est sans maillage (mesh-free), sans atlas de cartes et sans calcul de déterminant jacobien, ce qui la rend applicable à des dimensions plus élevées.

5. Signification et Perspectives

Cet article représente une avancée significative pour la simulation de processus stochastiques sur des espaces contraints complexes :

• Indépendance de la dimension : La méthode contourne la malédiction de la dimensionnalité en évitant les grilles, permettant potentiellement de traiter des variétés de haute dimension (ex: $S^{n-1}$ pour $n \ge 10$).
• Généralité Géométrique : L'approche est applicable à toute variété plongée où les projections tangentes $P(x)$ et les courbures moyennes $H(x)$ sont connues (sphères, tores, variétés de Stiefel, groupes de Lie).
• Efficacité Computationnelle : En évitant l'autodiff à travers la géométrie intrinsèque, le coût par étape d'entraînement est réduit et devient indépendant de la complexité des opérateurs différentiels intrinsèques.
• Applications Futures : Les auteurs suggèrent des extensions vers l'équation de la chaleur sur des groupes de Lie (comme $SO(3)$), des équations de Fokker-Planck fractionnaires, et l'utilisation de fonctions harmoniques sphériques pour compléter les ondes planes.

En résumé, cette méthode offre un cadre robuste, flexible et efficace pour apprendre des distributions de probabilité complexes sur des variétés riemanniennes, combinant la puissance des réseaux de neurones avec la rigueur de la théorie géométrique des mesures.