Résumé Technique : Un Cadre de Champ Moyen pour le Contrôle Distributionnel à l'Inférence des Modèles de Diffusion

Énoncé du Problème

Les modèles de diffusion sont de plus en plus utilisés comme échantillonneurs contrôlables, où les processus de génération sont guidés lors de l'inférence pour satisfaire des objectifs spécifiques. Alors que les méthodes existantes gèrent efficacement les récompenses ponctuelles (définies sur des échantillons individuels, $r: \mathbb{R}^d \to \mathbb{R}$), de nombreuses applications critiques nécessitent des récompenses distributionnelles (définies sur la mesure générée, $R: \mathcal{P}(\mathbb{R}^d) \to \mathbb{R}$). Ces objectifs distributionnels incluent le calibrage de populations par rapport à des observations expérimentales, l'encouragement de la diversité ou l'équilibrage des modes.

Les approches actuelles pour le pilotage distributionnel incorporent souvent des termes de gradient dans la dynamique d'inférence (par exemple, des potentiels répulsifs ou des vraisemblances d'ensemble). Cependant, à l'instar des premières méthodes de pilotage ponctuel, ces heuristiques manquent de garanties théoriques concernant la distribution échantillonnée. En effet, l'ajout simple d'un terme de gradient ne garantit pas que les particules suivent la mesure inclinée $\mu^*$ qui maximise l'objectif régularisé par la divergence KL :
$$\arg \max_{\mu} \left\{ R(\mu) - \text{KL}(\mu \| p_1) \right\}$$
Contrairement au cas ponctuel, où la cible inclinée est une inclinaison exponentielle explicite ($\mu^* \propto p_1 e^r$), le cas distributionnel aboutit à une cible implicite où le potentiel d'inclinaison dépend de la mesure elle-même. Il n'existe actuellement aucun cadre théoriquement fondé pour corriger la dynamique et les poids afin de cibler cette distribution implicite.

Méthodologie

Les auteurs proposent un cadre de champ moyen (Mean-Field Framework) pour formaliser et résoudre le contrôle distributionnel à l'inférence.

1. Formulation Théorique

Le problème est formulé comme le ciblage d'une mesure inclinée $\mu^*_t$ à chaque étape de temps $t \in [0,1]$, définie implicitement par :
$$\mu^*_t(dx) = \frac{1}{Z_t} e^{\Psi_t(x, \mu^*_t)} p_t(x) dx$$
où $\Psi_t(x, \mu) = \frac{\delta R_t}{\delta \mu}(x, \mu)$ est la première variation (dérivée fonctionnelle) de la récompense dépendante du temps $R_t$. Cela crée une relation d'auto-cohérence où la mesure cible dépend du potentiel, lequel dépend lui-même de la mesure.

2. Dynamique de McKean-Vlasov Pondérée

Pour suivre cette trajectoire, les auteurs dérivent un système de McKean-Vlasov pondéré. Contrairement au pilotage par gradient standard, le terme de pilotage $\nabla_x \Psi_t(X_t; \mu_t)$ induit une interaction entre les particules car il dépend de la mesure empirique évolutive $\mu_t$.

La dynamique idéalisée pour une particule $X_t$ et son log-poids $A_t$ est :
$$dX_t = \left[ b_t(X_t) + \epsilon_t s_t(X_t) + \epsilon_t \nabla_x \Psi_t(X_t; \mu_t) \right] dt + \sqrt{2\epsilon_t} dB_t$$
$$dA_t = \left[ b_t(X_t) \cdot \nabla_x \Psi_t(X_t; \mu_t) + \dot{\Psi}_t(X_t) \right] dt$$
Ici, $\dot{\Psi}_t(x) = \frac{d}{dt} \Psi_t(x, \mu_t)$ est un terme de correction crucial. Contrairement au cas ponctuel où $\dot{\Psi}_t$ est simplement $\partial_t r$, il implique ici l'évolution temporelle de la mesure $\mu_t$. Les auteurs démontrent que $\dot{\Psi}_t$ satisfait une équation intégrale implicite impliquant la seconde variation du fonctionnel de récompense.

3. Algorithme à Particules Finies

Pour implémenter cela, les auteurs proposent un schéma de particules interagissantes pondérées (Algorithme 1) :

• Mises à jour de position : Les particules sont mises à jour en utilisant le drift de base, le score et le gradient de la récompense par rapport à la mesure empirique actuelle $\hat{\mu}_t$.
• Mises à jour de log-poids : Les poids sont mis à jour pour corriger l'écart entre la dynamique pilotée et la trajectoire inclinée cible. Cela nécessite de calculer $\dot{\Psi}_t$.
• Résolution de $\dot{\Psi}_t$ : Le papier fournit deux méthodes pour calculer le terme de dérivée temporelle implicite :
  1. Système Linéaire Implicite : Résolution d'un système linéaire dérivé de l'équation intégrale (Proposition 3.3).
  2. Itération de Point Fixe : Un schéma d'itération de Picard léger approximant les mises à jour de poids.
• Rééchantillonnage : Un rééchantillonnage standard par Monte Carlo séquentiel (SMC) est employé pour éviter la dégénérescence des poids.

Le cadre récupère le pilotage standard de récompense ponctuelle comme un cas particulier où $R(\mu) = \int r(x) d\mu(x)$.

Contributions Clés

1. Formalisation du Pilotage Distributionnel : Les auteurs formalisent le pilotage au niveau de la distribution comme le ciblage d'une mesure inclinée où l'inclinaison dépend de la mesure elle-même, établissant un analogue distributionnel aux cibles de récompense inclinées ponctuelles.
2. Dérivation de la Dynamique de McKean-Vlasov Pondérée : Ils dérivent un système dynamique fondé avec un repondérage Feynman-Kac correctif qui garantit que les particules interagissantes pondérées suivent la trajectoire inclinée désirée $\mu^*_t$.
3. Récupération Théorique et Fondations : Le cadre récupère le pilotage ponctuel standard comme un cas particulier et fournit une base théorique pour les méthodes heuristiques de pilotage par lots existantes (ex: échantillonnage de diversité, guidage expérimental).
4. Garanties de Convergence : Le papier prouve que le système à particules finies converge vers la dynamique idéalisée lorsque le nombre de particules $N \to \infty$.

Résultats Expérimentaux

Les auteurs valident le cadre dans des contextes de faible et haute dimension :

• Faible Dimension (Mélange 1D) :

  • Tâche : Incliner un mélange de Gaussiennes bimodales vers une distribution cible ayant des poids de mode différents en utilisant une récompense de type MMD (Maximum Mean Discrepancy) au carré.
  • Résultats : Le pilotage par gradient standard pousse les particules vers la cible mais échoue à correspondre précisément à la forme de la distribution inclinée $\mu^*$. En revanche, la méthode de champ moyen cible précisément $\mu^*$, minimisant l'écart d'optimalité $J(\mu) - J(\mu^*)$ lorsque la force de pilotage $\lambda$ correspond au paramètre de l'objectif. Cela est vrai pour divers schémas de bruit.

• Haute Dimension (Conformation de Protéines) :

  • Tâche 1 (HIV-1 Protease) : Orienter les ensembles générés par Boltz-2 pour correspondre aux distributions de distances de résidus observées expérimentalement (données DEER). Le modèle de base était biaisé vers l'état "fermé" ; l'objectif était de récupérer les proportions de l'état "semi-ouvert".
    • Résultat : Le pilotage par champ moyen a correspondu avec précision à la distribution de distance de résidus empirique, tandis que le pilotage par gradient seul a produit une distribution avec une forme légèrement différente et des pics plus marqués.
  • Tâche 2 (Adénylate Kinase) : Orienter vers une distribution 2D d'angles de liaison dérivée de simulations de dynamique moléculaire.
    • Résultat : Le pilotage par champ moyen a produit un décalage plus marqué vers la distribution cible par rapport aux méthodes basées uniquement sur le gradient.
  • Tâche 3 (Cristallographie aux Rayons X) : Orienter les générations de Protenix pour s'ajuster à une carte de densité électronique (structure 4OLE) présentant des conformations hétérogènes.
    • Résultat : Pour une même force de pilotage, le pilotage par champ moyen a obtenu un alignement cosinus plus élevé avec la densité cible et a placé plus d'échantillons dans la conformation alternative (Conformation B) par rapport au pilotage par gradient seul, suggérant que les méthodes par gradient seul sous-pilotent par rapport à la véritable cible du problème inverse.

Signification et Revendications

Le papier affirme combler le fossé entre le contrôle ponctuel et le contrôle distributionnel en fournissant un fondement théorique rigoureux pour ce dernier.

• Correction : La revendication principale est que la procédure proposée cible exactement la distribution inclinée $\mu^*$ prescrite dans la limite de particules infinies, corrigeant les écarts inhérents au pilotage par gradient heuristique.
• Interprétabilité : Le cadre clarifie que les méthodes de pilotage par lots existantes effectuent effectivement des étapes de gradient mais manquent du repondérage correctif nécessaire pour échantillonner la véritable mesure postérieure/inclinée.
• Praticité : Bien que la méthode introduise un coût computationnel supplémentaire (résolution de $\dot{\Psi}_t$), les auteurs notent que cela est souvent négligeable par rapport aux évaluations de réseaux dans les contextes de haute dimension (ex: protéines) et que le solveur de point fixe est efficace.
• Limites : Les auteurs reconnaissent que la dépendance au rééchantillonnage peut être une limite pour des modalités comme les images (pouvant mener à des échantillons dupliqués) et que la méthode nécessite des fonctions de récompense différentiables. Ils suggèrent que la combinaison avec des techniques de réduction de variance (ex: solveurs d'équation de Poisson) pourrait atténuer les problèmes de rééchantillonnage.

Ce travail se positionne non pas comme un remplacement de toutes les méthodes existantes, mais comme un ancrage rigoureux expliquant pourquoi certaines heuristiques fonctionnent et fournissant un mécanisme pour garantir qu'elles ciblent la bonne distribution.