Proximal Diffusion Neural Sampler

Résumé Technique : Échantillonneur Neuronal de Diffusion Proximal (PDNS)

Énoncé du Problème
Apprendre un échantillonneur neuronal basé sur la diffusion pour prélever des échantillons à partir d'une distribution cible non normalisée (par exemple, une distribution de Boltzmann $\pi(x) \propto e^{-\beta E(x)}$) est une tâche fondamentale en statistiques computationnelles, en inférence bayésienne et en physique statistique. Ce problème est souvent formulé comme un problème de Contrôle Optimal Stochastique (SOC) sur les mesures de trajectoire, où un réseau de neurones paramétrise un processus de contrôle pour transporter une distribution de référence simple vers la cible.

Cependant, l'entraînement de ces échantillonneurs rencontre des défis majeurs lorsque la distribution cible est multimodale avec des barrières d'énergie élevées séparant les modes. Dans de tels scénarios, les méthodes standard basées sur le SOC (y compris les variantes de l'Entropie Croisée et de l'Entropie Croisée de Débruitage Pondérée) souffrent souvent d'un effondrement des modes. Au début de l'entraînement, le modèle ne capture que quelques modes ; la fonction de perte globale renforce ensuite ces modes capturés sur la base des trajectoires limitées générées, provoquant une concentration de la masse de probabilité sur un sous-ensemble de modes tout en ignorant les autres. Cela conduit à une optimisation instable et à une couverture médiocre de la distribution cible.

Méthodologie
Les auteurs proposent le PDNS (Proximal Diffusion Neural Sampler), un cadre unifié conçu pour résoudre l'effondrement des modes en reformulant le processus d'entraînement comme une séquence de sous-problèmes locaux plutôt qu'une minimisation globale unique.

1. Cadre Unifié : L'article unifie d'abord les échantillonneurs de diffusion continus et discrets sous un seul cadre SOC. Il définit la mesure de trajectoire optimale $P^*$ comme la solution minimisant un objectif d'entropie relative soumis à une mesure de trajectoire de référence $P_{ref}$ satisfaisant une condition de « sans mémoire » (indépendance entre les distributions initiale et terminale).
2. Méthode du Point Proximal : Pour atténuer l'effondrement des modes, le PDNS applique la méthode du point proximal à l'espace des mesures de trajectoire. Au lieu d'optimiser directement l'objectif global, l'algorithme résout itérativement un sous-problème local à l'étape $k$ :
   $$P_{\theta^*_k} = \arg\min_{P_\theta} \left( -E_{P_\theta}[r(X_T)] + KL(P_\theta \| P_{ref}) + \frac{1}{\eta_k} KL(P_\theta \| P_{\theta_{k-1}}) \right)$$
   Ici, $\eta_k$ est une taille de pas proximal, et le terme $KL(P_\theta \| P_{\theta_{k-1}})$ agit comme un régularisateur maintenant l'itération actuelle proche de la précédente.
3. Propriétés Théoriques : Les auteurs démontrent que la solution optimale de l'étape proximale $k$ est une interpolation géométrique entre l'itération précédente et la mesure de trajectoire optimale. Plus précisément, la distribution terminale de l'étape $k$, $P_k$, est une interpolation géométrique entre la distribution terminale de référence $\nu$ et la cible $\pi$. Lorsque $k \to \infty$, la séquence de mesures de trajectoire converge vers la cible $P^*$. Cette procédure par étapes garantit que le modèle explore l'espace d'état progressivement, empêchant un engagement prématuré envers un mode unique.
4. Instantiation Pratique (Proximal WDCE) : Pour l'efficacité, les auteurs instancient le PDNS en utilisant un objectif Proximal Weighted Denoising Cross-Entropy (Proximal WDCE). Cette approche évite de stocker les trajectoires complètes en utilisant des poids d'importance accumulés en ligne.
   • Cas Continu : L'objectif implique l'appariement d'un score de débruitage (ou d'un pont) avec un terme de poids dérivé de la taille de pas proximal $\eta_k$ et de la fonction d'énergie.
   • Cas Discret : L'objectif adapte la diffusion discrète masquée, utilisant un schéma de pondération similaire basé sur la dérivée de Radon-Nikodym des mesures de trajectoire.
5. Conception de l'Ordonnanceur : L'article discute des stratégies pour sélectionner la taille de pas $\eta_k$ (ou son équivalent $\lambda_k$). Un $\eta_k$ plus petit (régularisation plus forte) préserve la couverture des modes mais ralentit la convergence, tandis qu'un $\eta_k$ plus grand risque l'effondrement des modes. Les auteurs proposent à la fois des ordonnancements prédéfinis (diminution linéaire de $\lambda_k$) et une stratégie adaptative qui ajuste $\eta_k$ en fonction d'une contrainte de région de confiance sur la divergence KL entre le modèle actuel et la prochaine cible.

Contributions Clés

1. Cadre Unifié : L'article intègre les échantillonneurs basés sur le SOC continus et discrets dans une seule formulation de mesure de trajectoire.
2. PDNS (Proximal Diffusion Neural Sampler) : Il introduit un cadre novateur qui exploite les itérations de points proximaux dans l'espace des trajectoires pour stabiliser l'apprentissage et atténuer l'effondrement des modes dans des contextes multimodaux.
3. Proximal WDCE : Il dérive des instantiations concrètes et computationnellement efficaces du PDNS pour les domaines continus et discrets en utilisant des pertes d'entropie croisée de débruitage pondérée proximale.
4. Stratégies d'Ordonnanceur : Les auteurs exposent des stratégies fondées sur des principes pour sélectionner les tailles de pas proximaux, y compris un ordonnanceur adaptatif qui équilibre la vitesse de convergence et la stabilité.

Résultats Expérimentaux
Les auteurs évaluent le PDNS sur un large éventail de références continues et discrètes :

• Tâches Continues :
  • Fonctions d'Énergie Synthétiques : Le PDNS surpasse des bases solides (incluant SMC, DDS, PIS et AS) sur des références multimodales telles que le potentiel à 32 puits (MW54), l'entonnoir (Funnel) et les mélanges gaussiens de haute dimension (GMM40). Il obtient des résultats de l'état de l'art sur le mélange de Student (MoS) et les potentiels basés sur des particules (systèmes de Lennard-Jones LJ-13 et LJ-55).
  • Dynamique Moléculaire : Sur le système du Dipeptide d'Alanine (AD), le PDNS atteint des divergences KL comparables à la méthode de l'état de l'art ASBS et capture avec précision les structures conformationnelles dans les diagrammes de Ramachandran.
• Tâches Discrètes :
  • Physique Statistique : Le PDNS surpasse significativement LEAPS et Metropolis-Hastings (MH) dans l'échantillonnage de modèles d'Ising et de Potts à basse température (près ou en dessous des points critiques), où la séparation des modes est sévère. Il obtient des erreurs plus faibles en aimantation et en corrélations à 2 points.
  • Optimisation Combinatoire : Appliqué au problème du Maximum Cut sur les graphes Barabási–Albert et Erdős–Rényi, le PDNS démontre sa capacité à atteindre des qualités de solution comparables à celles des solveurs de vérité terrain (Gurobi).

Signification et Revendications
L'article revendique que le PDNS fournit une solution robuste au problème d'effondrement des modes inhérent à l'entraînement d'échantillonneurs neuronaux basés sur la diffusion pour des distributions complexes et multimodales. En décomposant l'optimisation globale en une séquence d'étapes proximales, la méthode favorise une exploration approfondie à travers les modes et stabilise l'optimisation dans les espaces de haute dimension.

Les auteurs notent deux limitations principales :

1. L'enquête se concentre sur une seule instantiation du PDNS (proximal WDCE), laissant la place à l'exploration d'autres objectifs proximaux ou priors.
2. Les expériences sont menées sur des références contrôlées plutôt que sur des applications réelles à grande échelle.

L'ouvrage suggère que les directions futures incluent l'exploration de problèmes SOC avec des dynamiques de référence non sans mémoire, l'extension du cadre à des espaces d'états plus généraux, et le développement d'autres instantiations de l'objectif proximal.