Résumé Technique : Appariement de Distribution Contrastive pour le Monte Carlo Séquentiel Amorti en Diffusion Discrète

Énoncé du Problème

Les modèles de diffusion discrète sont devenus des cadres puissants pour générer des données catégorielles structurées, telles que du texte, des séquences d'ADN et des protéines. Une application centrale de ces modèles est l'alignement par récompense, où l'objectif est d'échantillonner à partir d'une distribution cible qui biaise l'a priori préentraîné vers des valeurs plus élevées d'une fonction de récompense en aval (par exemple, la préférence humaine ou l'aptitude au design de protéines).

Bien que le Monte Carlo Séquentiel (SMC) Tordu offre une méthode asymptotiquement exacte pour échantillonner à partir de telles distributions biaisées par la récompense, son application dans les espaces d'états discrets se heurte à un goulot d'étranglement fondamental : l'estimation de la fonction de torsion optimale.

• Dans les domaines continus, la fonction de torsion peut être approximée efficacement en utilisant la formule de Tweedie ou un guidage basé sur le gradient.
• Dans les domaines discrets, l'absence d'une relation analogue force la dépendance à des approximations Monte Carlo (MC) coûteuses pour estimer la fonction de torsion à chaque étape d'inférence.
• Cette dépendance crée une surcharge computationnelle sévère, en particulier lorsque la fonction de récompense elle-même est coûteuse à évaluer (par exemple, des simulations de repliement de protéines ou un score de préférence de LLM basé sur une API). Augmenter le nombre d'échantillons MC ($M$) améliore la précision mais augmente linéairement le temps d'inférence, rendant l'approche peu pratique pour de nombreuses applications réelles.

Méthodologie

Les auteurs proposent l'Appariement de Distribution Contrastive (CDM), un cadre conçu pour amortir le coût de l'inférence SMC en apprenant une fonction de torsion paramétrée via un objectif d'apprentissage contrastif.

1. Objectif d'Appariement de Distribution Contrastive

Au lieu d'apprendre la fonction de torsion par régression directe (en minimisant l'Erreur Quadratique Moyenne par rapport aux estimations MC), le CDM minimise la divergence de Kullback-Leibler (KL) directe entre la distribution cible $p^*_t$ et une distribution paramétrée $p^\phi_t$ modulée par la fonction de torsion apprise $\psi^\phi_t$.

L'objectif est formulé comme suit :
$$\mathcal{L}_{\text{CDM}}(\phi) = \mathbb{E}_{t}\left[ \mathcal{D}_{\text{KL}}\left( p^*_t(\mathbf{x}_t) \big\| p^\phi_t(\mathbf{x}_t) \right) \right]$$

Le gradient de cette perte présente une structure contrastive :
$$-\nabla_\phi \mathcal{L} = \mathbb{E}_{t}\left[ \underbrace{\mathbb{E}_{p^*_t}[\nabla_\phi \log \psi^\phi_t]}_{\text{Terme Positif}} - \underbrace{\mathbb{E}_{p^\phi_t}[\nabla_\phi \log \psi^\phi_t]}_{\text{Terme Négatif}} \right]$$

• Terme Positif : Augmente le poids du log-torsion sur les échantillons tirés de la distribution cible optimale (régions à haute récompense).
• Terme Négatif : Réduit le poids du log-torsion sur les échantillons tirés de l'approximation actuelle (régions sous-optimales).

2. Schéma d'Entraînement Efficace

Pour remédier à l'inefficacité de l'échantillonnage à partir de la distribution cible $p^*_t$ pour le terme positif, les auteurs exploitent une propriété unique des modèles de diffusion : le noyau direct sous forme fermée.

• Estimation de Gradient Basée sur le Direct : Au lieu d'échantillonner indépendamment à partir de chaque cible intermédiaire $p^*_t$, la méthode échantillonne d'abord des échantillons positifs propres ($\mathbf{x}_0$) à partir de la distribution cible en utilisant le SMC. Elle utilise ensuite le noyau direct sous forme fermée $p^{\text{base}}(\mathbf{x}_t | \mathbf{x}_0)$ pour générer plusieurs échantillons positifs à n'importe quel pas de temps intermédiaire $t$ à un coût négligeable.
• Entraînement Basé sur un Tampon : Un tampon d'échantillons positifs propres est maintenu et réutilisé sur plusieurs pas de temps et mises à jour de gradient, réduisant considérablement le nombre d'évaluations de récompense coûteuses requises pendant l'entraînement.
• Échantillonnage Négatif : Les échantillons négatifs sont tirés par Échantillonnage par Importance (IS) en utilisant la torsion paramétrée actuelle, ce qui est computationnellement efficace.

3. Paramétrisation

Pour garantir que la fonction de torsion apprise n'introduise pas de surcharge d'inférence, la fonction de torsion $\psi^\phi_t$ est paramétrée comme une tête scalaire légère attachée à la dernière couche de caractéristiques du noyau de diffusion préentraîné. Cela permet au modèle de produire à la fois les logits et l'estimation de la torsion en un seul passage avant, ajoutant moins de 5 % de surcharge computationnelle par rapport au modèle de base.

Contributions Clés

1. Cadre CDM : Introduction d'un nouveau cadre d'amortissement SMC pour la diffusion discrète qui réduit le coût de l'application de la fonction de torsion au moment de l'inférence à une opération de temps constant.
2. Entraînement Natif à la Diffusion : Conception d'un schéma d'entraînement qui exploite le processus direct sous forme fermée des modèles de diffusion pour permettre un entraînement efficace avec des fonctions de récompense coûteuses, évitant l'inefficacité d'échantillonnage des approches basées sur la régression standard.
3. Polyvalence et Synergie : Démonstration que le CDM est agnostique à la distribution de proposition, permettant de le coupler avec des propositions affinées (par exemple, d1, DRAKES) pour des gains de performance synergiques.
4. Objectif Supérieur : Validation que l'objectif d'apprentissage contrastif produit des performances supérieures et une convergence plus rapide par rapport aux objectifs de torsion basés sur la régression standard (Soft Value).

Résultats Expérimentaux

Les auteurs ont évalué le CDM sur quatre applications diverses : génération de texte toxique, conception de séquences d'ADN réglementaires, aptitude au design de protéines et alignement de grands modèles de langage par diffusion (dLLM).

• Performance de Mise à l'Échelle : Dans toutes les tâches, le CDM a systématiquement surpassé les lignes de base (Best-of-N, SMC standard avec $M$ variable, et Soft Value) lorsque les temps d'horloge étaient équivalents.
• Récompenses Coûteuses : Dans la tâche d'aptitude au design de protéines, où l'évaluation de la récompense est le goulot d'étranglement dominant, le CDM a atteint une mise à l'échelle supérieure par rapport à la fois à Best-of-N et à Soft Value, tandis que le SMC standard n'a pas réussi à se mettre à l'échelle efficacement en raison de la surcharge d'inférence.
• Compatibilité avec l'Affinement (Fine-Tuning) : Lorsqu'il est intégré avec des propositions affinées (d1 et DRAKES), le CDM améliore davantage les performances et atténue les problèmes d'effondrement de mode souvent observés dans les modèles affinés, maintenant la diversité tout en atteignant des récompenses élevées.
• Efficacité de l'Entraînement : Le CDM a démontré une convergence plus rapide en temps d'horloge par rapport à Soft Value, attribué à l'efficacité de l'objectif d'apprentissage contrastif avec des échantillons négatifs.

Signification et Revendications

L'article revendique que le CDM adresse le principal goulot d'étranglement dans l'inférence SMC tordue pour la diffusion discrète : la dépendance à une estimation Monte Carlo coûteuse. En apprenant la fonction de torsion via un objectif contrastif qui exploite le processus direct de diffusion, la méthode permet un alignement par récompense efficace et évolutif.

Les auteurs soulignent que leur approche est orthogonale à l'affinement de la proposition ; elle peut être combinée avec des méthodes d'affinement existantes pour obtenir des gains de performance supplémentaires. De plus, la méthode reste applicable même lorsque la fonction de récompense est non différentiable ou computationnellement coûteuse, un scénario où le guidage basé sur le gradient et le SMC standard peinent. Les résultats suggèrent que le CDM établit une nouvelle frontière de Pareto pour le compromis entre la qualité de la récompense et le temps d'inférence dans la modélisation générative discrète.