Résumé Technique : Stein Diffusion Guidance (SDG)

Énoncé du Problème

Les modèles de diffusion ont démontré un succès remarquable dans la modélisation de distributions complexes et de haute dimension. Cependant, lorsqu'ils sont entraînés sur des données non étiquetées, ils échantillonnent principalement depuis les régions de haute densité de la variété des données, négligeant souvent les régions de faible densité où peuvent résider des échantillons rares à haute valeur (par exemple, de nouveaux candidats médicaments).

Les méthodes existantes pour remédier à cette limitation se divisent généralement en deux catégories, toutes deux présentant des inconvénients majeurs :

1. Guidage basé sur un classifieur : Ces méthodes exploitent des classifieurs auxiliaires entraînés sur des données bruitées à plusieurs niveaux de bruit pour guider l'échantillonnage. Cependant, elles introduisent une complexité dans l'entraînement et peuvent égarer les échantillons si les gradients du classifieur sont bruités, provoquant un dérive des échantillons hors de la variété générative, en particulier dans les régimes de faible densité.
2. Contrôle Optimal Stochastique (SOC) : Le SOC offre un cadre principiel pour l'échantillonnage a posteriori en orientant le processus de diffusion via un contrôleur optimal. Bien que théoriquement solide et capable d'exploiter des classifieurs prêts à l'emploi entraînés sur des données propres, le calcul du contrôle optimal nécessite de rétropropager les signaux de récompense à travers l'ensemble des trajectoires de réseaux de neurones-SDE. Cela rend le SOC prohibitif sur le plan computationnel pour une inférence efficace.

Les approches récentes « sans entraînement » tentent d'approximer l'a posteriori de diffusion en utilisant la formule de Tweedie pour éviter le coût computationnel du SOC. Cependant, l'article soutient que cette approximation est intrinsèquement biaisée et sous-optimale, conduisant à un guidage peu fiable, en particulier dans les régions de faible densité où les modèles basés sur le score sont les moins précis.

Méthodologie : Stein Diffusion Guidance (SDG)

Les auteurs proposent Stein Diffusion Guidance (SDG), un cadre sans entraînement qui concilie l'efficacité des méthodes sans entraînement avec la nature principielle du SOC. SDG est fondé sur un objectif SOC de substitution et introduit un mécanisme de correction de Stein pour affiner les a posteriori approximatifs.

1. Formulation SOC de Substitution

Les auteurs formulent l'échantillonnage de diffusion en faible densité comme un problème de contrôle optimal stochastique avec une nouvelle fonctionnelle de coût $\tilde{J}(u, x, t)$. Cette fonctionnelle inclut :

• Un coût d'état qui pénalise les régions de haute densité (encourageant l'exploration des zones de faible densité via un programme de recuit $\alpha(t)$).
• Un coût terminal basé sur un classifieur prêt à l'emploi ou un modèle de récompense $r(x_T)$ (orientant les échantillons vers les propriétés souhaitées via une intensité $\beta(t)$).

Ils dérivent une borne supérieure variationnelle sur la fonction de valeur SOC $V(x,t)$. Cette borne révèle que l'optimisation du guidage sans entraînement standard (qui repose sur la formule de Tweedie) correspond à la minimisation des deux premiers termes de la borne, négligeant la divergence de Kullback-Leibler (KL) entre l'a posteriori approximatif $q$ et l'a posteriori vrai $p$. Cette négligence conduit à un guidage sous-optimal.

2. Correction de Stein Aller-Retour

Pour traiter le terme de divergence KL sans encourir le coût computationnel du SOC complet, SDG introduit un mécanisme de correction de Stein aller-retour :

• Mapping Arrière : Les particules bruitées $x_t$ sur la variété de diffusion $M_t$ sont mappées en arrière vers la variété des données $M_T$ en utilisant la formule de Tweedie pour obtenir un ensemble initial d'échantillons d'a posteriori approximatifs $x_T$.
• Correction de Stein : Ces particules sont ensuite évoluées sur $M_T$ en utilisant la Descente de Gradient Variationnel de Stein (SVGD). Les particules sont mises à jour dans la direction de la descente la plus raide pour minimiser la divergence KL entre l'a posteriori approximatif et l'a posteriori vrai. Cela utilise le score de l'a posteriori conditionnel, approximé via les scores du modèle : $\nabla_{x_T} \log p(x_T|x_t) \approx s_\theta(x_T) - \eta(t)s_\theta(x_t)$.
• Mapping Avant : Les particules corrigées sont mappées vers l'avant vers la variété courante $M_t$ pour être utilisées lors de l'étape de diffusion suivante.

Ce mécanisme régularise efficacement les échantillons a posteriori, garantissant qu'ils restent proches de la dynamique de diffusion réelle et de la variété générative, même lors de l'utilisation de classifieurs prêts à l'emploi.

3. Algorithme

L'algorithme SDG intègre cette correction dans le processus de diffusion inverse standard. À chaque étape, il :

1. Génère des estimations initiales d'a posteriori via la formule de Tweedie.
2. Applique l'étape de correction de Stein (mises à jour des particules basées sur les interactions de noyau et les différences de scores).
3. Utilise les échantillons corrigés pour calculer le signal de guidage (gradient de récompense) pour la mise à jour de l'EDS inverse.

Contributions Clés

1. Cadre de Faible Densité : Une formulation novatrice basée sur le SOC pour le guidage de diffusion ciblant explicitement les régions de faible densité via une nouvelle fonction de coût à venir.
2. Insight Théorique : Une preuve théorique démontrant que l'approximation de l'a posteriori de diffusion via la formule de Tweedie est sous-optimale et nécessite une étape de correction pour tenir compte de la divergence KL entre les a posteriori approximatifs et vrais.
3. Mécanisme de Correction de Stein : L'introduction d'une correction de Stein sans entraînement et efficace sur le plan computationnel qui minimise itérativement la divergence KL entre les a posteriori approximatifs et vrais, permettant un guidage robuste dans les régimes de faible densité.
4. Généralisabilité : Un cadre plug-and-play qui peut être intégré dans divers problèmes d'échantillonnage a posteriori de diffusion sans nécessiter de finetuning du modèle ni de distillation.

Résultats Expérimentaux

Les auteurs évaluent SDG sur diverses tâches, notamment le guidage d'images et l'échantillonnage moléculaire pour l'amarrage de protéines.

• Tâches de Guidage d'Images : Sur des tâches telles que le guidage par étiquette d'image, le défloutage gaussien, la super-résolution et le transfert de style texte-à-image, SDG (spécifiquement la variante sans échantillonnage de faible densité, SDG$\heartsuit$) surpasse constamment les bases sans entraînement standard (par exemple, DPS, LGD) et même des méthodes avancées avec des stratégies récurrentes (UGD, MPGD) sur la plupart des métriques. Cela met en évidence l'importance de corriger les a posteriori approximatifs, même dans des contextes non de faible densité.
• Échantillonnage Moléculaire (Faible Densité) : Dans les tâches d'amarrage ligand-protéine (Fa7, 5ht1b, Jak2, Parp1), SDG surpasse nettement les bases (GDSS, MOOD, FREED, etc.) dans la génération de nouvelles molécules « hits » satisfaisant des critères stricts (nouveauté, accessibilité synthétique, ressemblance à un médicament et score d'amarrage).
  • Études d'Ablation : Le retrait de la correction de Stein (SDG$\clubsuit$) conduit à de mauvaises performances, les échantillons échouant souvent à satisfaire les critères de hit ou tombant hors de la variété générative. Le retrait de l'échantillonnage de faible densité (SDG$\heartsuit$) donne également des résultats sous-optimaux, confirmant la nécessité de cibler les régions de faible densité pour la découverte de nouveautés.
  • Surestimation de la Récompense : L'article démontre que sans correction de Stein, les modèles de récompense ont tendance à surestimer les récompenses dans les régions de faible densité, conduisant à des molécules irréalistes. SDG agit comme un mécanisme de régularisation, atténuant cette surestimation et préservant la validité des échantillons.
  • Efficacité : Bien que SDG entraîne une surcharge de temps d'exécution (environ 36-85 % selon le solveur et la tâche), il reste computationnellement faisable et s'étend à de grandes tailles de particules (jusqu'à 3000) avec une surcharge mémoire minimale par rapport au SOC complet.

Importance et Revendications

L'article revendique que SDG fournit une alternative sans entraînement et efficace sur le plan computationnel aux méthodes de Contrôle Optimal Stochastique prohibitivement coûteuses pour l'échantillonnage a posteriori. En établissant théoriquement la nécessité d'une correction a posteriori et en l'implémentant via l'inférence variationnelle de Stein, SDG permet un échantillonnage efficace dans les régions de faible densité où des découvertes à haute valeur (telles que de nouveaux candidats médicaments) sont susceptibles d'exister.

Les auteurs soulignent que leur méthode :

• Surmonte le biais des approximations basées sur Tweedie.
• Empêche les échantillons de dériver hors de la variété générative.
• Atténue la surestimation des récompenses dans les régimes de faible densité.
• Offre un cadre général applicable à divers problèmes d'échantillonnage a posteriori au-delà des régions de haute densité.

Ce travail suggère que la correction des a posteriori approximatifs est un composant critique, souvent négligé, dans le guidage de diffusion sans entraînement, et que SDG débloque le potentiel des modèles de diffusion pour découvrir des échantillons rares et à haute valeur dans les domaines scientifiques.