Résumé Technique : Dynamique Stochastique de Données à Énergie

Énoncé du Problème

Le problème du pont de Schrödinger (Schrödinger Bridge - SB) cherche à trouver un système dynamique stochastique qui relie deux distributions marginales, $p_0$ et $p_1$, tout en minimisant un coût de transport par rapport à un processus de référence (généralement une équation différentielle stochastique d'Itô). Ce problème généralise le transport optimal au cadre stochastique et sous-tend les techniques modernes de modélisation générative comme les modèles de diffusion et le flow matching.

Les algorithmes existants pour résoudre le problème SB, particulièrement ceux basés sur la procédure d'ajustement proportionnel itératif (Iterative Proportional Fitting - IPF), nécessitent généralement l'accès à des échantillons des deux distributions marginales, $p_0$ et $p_1$. Cette limitation empêche l'application des méthodes SB dans des scénarios où l'une ou les deux distributions sont définies uniquement par des densités non normalisées (fonctions d'énergie) $p(x) = e^{-E(x)}/Z$, où la fonction de partition $Z$ est inconnue et aucun échantillon direct n'est disponible. Cette lacune est critique pour des applications telles que l'inférence bayésienne des postulats dans les espaces latents de modèles génératifs, où le postérieur est souvent connu uniquement à une constante de normalisation près via une fonction de récompense ou de vraisemblance.

Méthodologie

Le papier propose le Data-to-Energy (et Energy-to-Energy) IPF, une méthode générale pour modéliser les ponts de Schrödinger lorsque les distributions sont données par des densités non normalisées. L'approche étend l'algorithme IPF standard en intégrant des techniques d'apprentissage par renforcement hors-politique (off-policy) et de l'échantillonnage par diffusion.

1. Cadre IPF Généralisé

L'algorithme IPF standard met à jour de manière itérative les processus de temps direct et inverse ($\overrightarrow{P}_t$ et $\overleftarrow{P}_t$) en résolvant des problèmes de demi-pont.

• Étape Arrière (Data-to-Data) : Lorsque des échantillons de $p_0$ sont disponibles, le processus inverse est entraîné via le maximum de vraisemblance (étape IPF standard).
• Étape Avant (Data-to-Energy) : Lorsque $p_1$ est donnée par une densité non normalisée $e^{-E_1(x)}$, le processus direct ne peut pas être entraîné via la vraisemblance standard. À la place, les auteurs introduisent une perte de log-variance (VarGrad conditionnelle). Cette perte minimise la variance du log-rapport entre la probabilité de la trajectoire directe et la densité cible, pondérée par la probabilité de la trajectoire inverse :
  $$\text{LLV}(x_0, \theta) = \text{Var} \left( \log \frac{\overrightarrow{p}_\theta(\tau | x_0)}{p_1(x_1) \overleftarrow{p}_\phi(\tau | x_1)} \right)$$
  Cet objectif impose la proportionnalité requise par la solution SB sans nécessiter d'échantillons de $p_1$.

2. Stratégies d'Apprentissage par Renforcement Hors-Politique

Pour répondre au défi de l'échantillonnage de densités non normalisées complexes et de haute dimension où l'entraînement sur-politique (on-policy) échoue à explorer tous les modes, les auteurs adaptent des techniques hors-politique :

• Buffer de Rejeu (Replay Buffer) : Un buffer stocke les échantillons finaux $x_1$ du processus direct. Pendant l'entraînement, $x_0$ est généré en échantillonnant $x_1$ depuis le buffer et en exécutant une trajectoire inverse. Cela concentre l'exploration sur les régions de haute densité de la cible.
• Mises à jour de Langevin : Le buffer est périodiquement mis à jour à l'aide de la dynamique de Langevin non ajustée sur la densité cible $p_1$ pour corriger les ajustements imparfaits.
• Réutilisation de Trajectoires : L'algorithme réutilise les trajectoires inverses pour former des lots de trajectoires partageant un même point initial $x_0$, améliorant ainsi l'efficacité de l'estimation de la perte basée sur la variance.
• Mixage : La politique d'entraînement mélange des échantillons sur-politique avec des échantillons hors-politique du buffer pour équilibrer l'exploration et l'exploitation.

3. Coefficients de Diffusion Apprenables

Une contribution méthodologique secondaire implique la généralisation de l'algorithme IPF pour apprendre non seulement les fonctions de dérive (drift) mais aussi les coefficients de diffusion ($\sigma_t$) de la dynamique. En remplaçant les variances fixes dans les étapes de discrétisation d'Euler-Maruyama par des fonctions apprenables, la méthode vise à corriger les erreurs de discrétisation temporelle, un concept inspiré par des travaux récents sur les échantillonneurs de diffusion.

Contributions Clés

1. Première Méthode SB Data-to-Energy Générale : Le papier introduit le premier algorithme capable d'inférer des ponts de Schrödinger lorsque l'une ou les deux distributions sont fournies sous forme de densités non normalisées sans accès aux échantillons. Il apprend avec succès des ponts stochastiques à faible coût de transport entre des jeux de données synthétiques et des densités, avec des performances comparables à l'IPF Data-to-Data utilisant des échantillons réels.
2. Coefficients de Diffusion Apprenables : Les auteurs démontrent que les algorithmes IPF Data-to-Data existants peuvent être considérablement améliorés en apprenant le coefficient de diffusion en plus de la dérive. Cela généralise les découvertes récentes concernant les échantillonneurs de diffusion au cadre plus large du SB, montrant qu'une variance apprenable permet une modélisation plus précise, particulièrement avec un nombre faible d'étapes de discrétisation.
3. Échantillonnage de Postérieur Externalisé : La méthode est appliquée au problème d'échantillonnage des distributions postérieures dans les espaces latents de modèles génératifs (ex: GAN, VAE). En modélisant un pont entre la distribution a priori latente et un postérieur pondéré par une récompense, les auteurs créent une méthode de traduction d'image à image sans données (data-free). Cette approche préserve le contenu sémantique (style, arrière-plan) non contraint par la cible, contrairement aux échantillonneurs de diffusion standards.

Résultats Expérimentaux

L'efficacité de la méthode proposée est démontrée à travers plusieurs tests de référence :

• Benchmarks Synthétiques 2D : Sur des tâches comme passer d'une Gaussienne à un mélange de Gaussiennes (GMM) et des distributions "Two Moons", l'IPF Data-to-Energy atteint des performances comparables à l'IPF Data-to-Data. L'inclusion d'une variance apprenable réduit significativement la distance Wasserstein-2 et le Path KL, surtout lorsque le nombre d'étapes de discrétisation est faible.
• Pontage Énergie-à-Énergie : Des résultats préliminaires montrent que l'algorithme peut réussir à relier deux distributions définies uniquement par des fonctions d'énergie, sans accès à des échantillons de l'une ou l'autre.
• Traduction d'Images (MNIST & CIFAR-10) : Dans les expériences d'espace latent utilisant SN-GAN et StyleGAN, la méthode traduit des échantillons a priori vers les postérieures de classes cibles (ex: générer des "chiens" ou des "cinq").
  • Métriques : La méthode obtient des scores FID plus bas que l'échantillonnage par rejet et la dynamique de Langevin pour plusieurs classes.
  • Qualitatif : L'inspection visuelle confirme que le pont préserve les caractéristiques de style (ex: épaisseur du chiffre, orientation de l'arrière-plan) tout en transformant le contenu vers la classe cible, suggérant une préservation efficace de l'information sémantique non contrainte.
• Études d'Ablation : Les expériences confirment que les buffers de rejeu, les mises à jour de Langevin et la réutilisation de trajectoires améliorent significativement les métriques de Path KL. La réutilisation des trajectoires inverses améliore le Path KL mais peut impacter négativement la moyenne du log-récompense, suggérant un compromis entre la couverture des modes et le coût de transport.

Signification et Revendications

Le papier revendique de fournir un cadre général pour les ponts de Schrödinger de type données-vers-énergie, comblant une lacune dans la littérature où les méthodes SB étaient auparavant restreintes aux contextes disposant d'un accès aux échantillons pour les deux marginales.

Les auteurs positionnent leur travail comme une étape significative vers la modélisation générative et l'inférence de postérieur scalables et sans données. En exploitant les techniques d'apprentissage par renforcement hors-politique, la méthode permet d'apprendre des dynamiques de transport optimal dans des espaces latents de haute dimension où seules les fonctions d'énergie sont disponibles. La capacité d'apprendre les coefficients de diffusion renforce davantage la robustesse de ces dynamiques face aux erreurs de discrétisation.

Les auteurs notent modestement que bien que la méthode soit scalable aux tâches de génération d'images de haute dimension, des travaux futurs sont nécessaires pour traiter l'effondrement de mode (mode collapse) et améliorer davantage la couverture des modes. Ils identifient également des directions futures potentielles, incluant l'ajustement fin (fine-tuning) de récompenses textuelles-conditionnelles et les problèmes de pont de Schrödinger discrets, sans toutefois revendiquer de succès immédiat dans ces domaines inexplorés.