Résumé technique : Échantillonnage par flux

Énoncé du problème

L'article aborde le défi de l'échantillonnage à partir de distributions de probabilité de haute dimension définies uniquement à une constante de normalisation près (densités non normalisées). De telles distributions sont omniprésentes en sciences computationnelles, définies par une fonction d'énergie $r(x)$ où la densité cible est $q(x) \propto \exp(r(x))$. Bien que l'évaluation de $r(x)$ et de son gradient $\nabla r(x)$ soit possible, elle est souvent coûteuse en calcul, et l'accès direct à des échantillons de $q(x)$ n'est pas disponible.

Les méthodes traditionnelles de Monte Carlo par chaîne de Markov (MCMC), telles que la dynamique de Langevin, convergent asymptotiquement mais souffrent d'un mélange lent et d'une génération séquentielle, ce qui les rend inefficaces pour des applications nécessitant des échantillonneurs rapides et réutilisables (par exemple, la conception de matériaux). Les approches basées sur l'apprentissage existantes, y compris les modèles de diffusion et l'appariement de flux, nécessitent généralement des échantillons de données provenant de la distribution cible ou reposent sur des corrections stochastiques (par exemple, pondération par importance, rééchantillonnage ou MCMC auxiliaire) qui entraînent des coûts computationnels significatifs dus aux évaluations répétées de la fonction d'énergie.

Méthodologie : Échantillonnage par flux

Les auteurs introduisent l'Échantillonnage par flux, un cadre pour apprendre des échantillonneurs amortis directement à partir de fonctions de densité non normalisées sans nécessiter d'échantillons de données cibles. La méthode est fondée sur les modèles de diffusion et l'appariement de flux, mais adapte l'objectif d'entraînement au contexte sans données.

Concept central

Contrairement à l'appariement de flux standard, qui conditionne sur un échantillon de données $X_1 \sim q$ pour apprendre une dérive de bruitage, l'Échantillonnage par flux conditionne sur un échantillon de bruit $X_0 \sim p_0$ (provenant d'une distribution source facile à échantillonner) pour apprendre une dérive de débruitage. La distribution marginale de ce processus appris est conçue pour correspondre à la cible $q$ par construction.

Objectif d'entraînement

Le cadre utilise un chemin de probabilité conditionnelle défini comme une image directe de la distribution cible conditionnée à la source :
$$p_{t|0}(x|x_0) = \frac{1}{\alpha_t^d} q\left(\frac{x - \sigma_t x_0}{\alpha_t}\right)$$
L'objectif d'entraînement consiste à régresser une dérive apprise $u_\theta$ vers une dérive conditionnelle $u_{t|0}$ construite à partir de la fonction d'énergie. Une insight théorique clé (Proposition 3.2) fournit une expression fermée pour cette dérive conditionnelle le long d'un processus d'interpolation $X_t = \sigma_t X_0 + \alpha_t X_1$ :
$$u_{t|0}(X_t|x_0) = \dot{\alpha}_t X_1 + \dot{\sigma}_t x_0 + \frac{g_t^2}{2\alpha_t} \nabla r(X_1)$$
Cette formulation permet au modèle de réutiliser le gradient d'énergie $\nabla r(X_1)$ évalué en un seul point $X_1$ pour tous les pas de temps $t \in [0, 1]$, réduisant considérablement le nombre d'évaluations de l'énergie requises pendant l'entraînement.

Algorithme : Itération de point fixe

Puisque les échantillons cibles $X_1 \sim q$ ne sont pas disponibles, les auteurs proposent un algorithme d'entraînement par itération de point fixe (Algorithme 1) composé de deux phases :

1. Phase d'exploration : En utilisant l'état actuel détaché du modèle $\bar{\theta}$, l'algorithme génère des échantillons $X_1^{\bar{\theta}}$ par simulation (Euler–Maruyama), évalue le gradient d'énergie $\nabla r(X_1^{\bar{\theta}})$, et stocke les paires $(X_1^{\bar{\theta}}, \nabla r(X_1^{\bar{\theta}}))$ dans un tampon de replay.
2. Phase d'optimisation : Le modèle $u_\theta$ est entraîné en minimisant l'erreur quadratique moyenne entre sa dérive prédite et la cible à forme fermée dérivée des échantillons du tampon de replay. Ce processus se répète jusqu'à convergence.

Extension aux variétés riemanniennes

Le cadre s'étend naturellement aux géométries non euclidiennes. En remplaçant l'interpolant affine par un interpolant géodésique et en dérivant une formule à forme fermée pour la dérive conditionnelle sur des variétés à courbure constante (incluant les hypersphères et les espaces hyperboliques), la méthode permet un échantillonnage basé sur la diffusion sur des variétés riemanniennes sans reparamétrisation.

Contributions clés

1. Cadre d'Échantillonnage par flux : Une méthode fondée sur des principes pour apprendre des échantillonneurs de diffusion à partir de densités non normalisées en appariant des dérives de diffusion de débruitage conditionnelles, éliminant ainsi le besoin d'échantillons de données cibles.
2. Objectif d'entraînement efficace : Dérivation d'une cible de régression simple à forme fermée qui réutilise les gradients d'énergie aux extrémités le long du processus d'interpolation, permettant un algorithme d'entraînement par tampon de replay efficace qui réduit drastiquement les évaluations de la fonction d'énergie.
3. Extension riemannienne : Extension des échantillonneurs basés sur la diffusion à des géométries au-delà de l'espace euclidien, incluant une formule à forme fermée pour la dérive conditionnelle sur des variétés à courbure constante.
4. Performance empirique : Démonstration de performances de pointe sur des benchmarks d'énergie synthétiques et la génération de conformères moléculaires, avec une réduction rapportée des coûts d'entraînement par simulation de 4 à 8 fois par rapport aux méthodes précédentes.

Résultats expérimentaux

Les auteurs évaluent l'Échantillonnage par flux sur quatre benchmarks distincts :

• Benchmarks d'énergie synthétiques : Sur les potentiels Double-Well (DW-4) et Lennard-Jones (LJ-13, LJ-55), l'Échantillonnage par flux atteint des distances de Wasserstein-2 ($W_2$) géométriques et énergétiques inférieures par rapport aux références de pointe telles que iDEM, PIS, DDS et l'échantillonnage adjoint (AS).
• Échantillonnage de peptides : Sur le dipeptide d'alanine (Ala2) et le tétrapeptide (Ala4), la méthode récupère avec précision les densités de vérité terrain (mesurées par la divergence de Kullback-Leibler et la divergence de Jensen-Shannon) et démontre une robustesse lors de la réduction du nombre d'évaluations de fonctions (NFE) pendant l'entraînement, surpassant ASBS qui se dégrade considérablement à des NFE plus faibles.
• Génération amortie de conformères : Sur des ensembles de données moléculaires à grande échelle (SPICE et GEOM-DRUGS), l'Échantillonnage par flux atteint une couverture et une précision élevées dans la génération de conformères moléculaires, surpassant les références RDKit et ASBS tout en nécessitant moins d'étapes d'entraînement.
• Échantillonnage riemannien : Sur la sphère unité 2D ($S^2$) avec un mélange de distributions de von Mises–Fisher, la méthode apprend avec succès à échantillonner à partir de la cible multimodale, validant l'extension théorique aux géométries courbes.

Importance et affirmations

L'article affirme que l'Échantillonnage par flux offre une alternative flexible et facile à mettre en œuvre aux formulations complexes de contrôle stochastique optimal ou de pont de Schrödinger. En évitant les couplages optimaux et les réseaux correcteurs auxiliaires, il simplifie le schéma d'entraînement tout en maintenant une haute efficacité.

Les auteurs soulignent que la capacité de la méthode à réduire le nombre d'évaluations de la fonction d'énergie (via le tampon de replay et la dérive à forme fermée) la rend particulièrement évolutive pour les grands modèles et ensembles de données où l'évaluation de l'énergie constitue le principal goulot d'étranglement. De plus, l'extension aux variétés riemanniennes fournit une approche unifiée pour l'échantillonnage dans des contextes géométriques divers, tels que la robotique et la dynamique moléculaire, sans nécessiter de reparamétrisations spécialisées. Ce travail positionne l'Échantillonnage par flux comme une étape significative vers un échantillonnage amorti et efficace pour les densités non normalisées dans les espaces euclidiens et non euclidiens.