Résumé technique : Décodage guidé primal-dual pour la diffusion discrète contrainte

Énoncé du problème

Les modèles de diffusion discrète sont devenus une alternative puissante aux modèles autorégressifs (AR) pour générer des séquences structurées (par exemple, texte, molécules, listes de lecture musicales) en dévoilant progressivement des jetons. Cependant, un défi majeur subsiste : imposer des contraintes de propriétés globales (par exemple, des comptes spécifiques de vocabulaire, des masses moléculaires cibles ou des compositions de genres) lors de la génération sans compromettre la distribution de données apprise par le modèle.

Les approches existantes font face à un compromis entre la satisfaction des contraintes et la fidélité distributionnelle :

• Guidage par prédicteur/classifieur : Nécessite l'entraînement d'un modèle auxiliaire pour chaque propriété et opère souvent dans des espaces continus non directement applicables aux jetons discrets.
• Diffusion discrète contrainte (CDD) : Évite l'entraînement auxiliaire mais repose sur des boucles d'optimisation internes (par exemple, Gumbel-softmax) à chaque étape de débruitage, entraînant une surcharge computationnelle de 5 à 20 fois.
• Biais statique : Des biais additifs simples échouent souvent à s'adapter au contexte évolutif du processus de génération, conduisant à un dépassement ou à une satisfaction insuffisante des contraintes.

Le problème central consiste à formuler la génération contrainte comme un problème d'optimisation pouvant être résolu en ligne durant le processus de débruitage, tout en maintenant l'efficacité computationnelle de l'échantillonnage standard tout en fournissant des garanties formelles sur la violation des contraintes.

Méthodologie : Décodage guidé primal-dual

Les auteurs proposent le Décodage guidé primal-dual (SPDD), une méthode d'inférence qui formule la génération contrainte comme un problème d'optimisation régularisé par la divergence KL. La méthode résout ce problème en ligne en utilisant des multiplicateurs de Lagrange adaptatifs mis à jour via la descente de miroir.

1. Formulation de l'optimisation

L'objectif est de trouver une distribution d'échantillonnage $r(x_0)$ qui satisfait une contrainte linéaire (par exemple, une contribution totale attendue $\ge R$) tout en minimisant la divergence KL par rapport à la distribution du modèle original $p_\theta(x_0)$.
$$\min_{r} D_{KL}(r(x_0) \| p_\theta(x_0)) \quad \text{sous contrainte} \quad \sum b_{\ell j} r(x_0^\ell = j) \ge R$$
où $b_{\ell j}$ représente la contribution du jeton $j$ à la position $\ell$ envers la contrainte.

2. Approximation primal-dual en ligne

Au lieu de résoudre la trajectoire complète hors ligne, la méthode traite le problème comme une optimisation en ligne sur les décisions d'échantillonnage au niveau des jetons à chaque étape de débruitage $t$.

• Mise à jour Primal (Échantillonnage) : À chaque étape, les probabilités des jetons sont modifiées par un biais additif dépendant de la contrainte, dérivé du multiplicateur de Lagrange $\lambda$. La distribution d'échantillonnage optimale est :
  $$a_{\ell j}^t = \frac{p_\theta(x_\ell^t = j | x_{t+1}) e^{\lambda_{t+1} b_{\ell j}}}{\sum_i p_\theta(x_\ell^t = i | x_{t+1}) e^{\lambda_{t+1} b_{\ell i}}}$$
  Cela re-pèse efficacement les logits du modèle, augmentant le poids des jetons qui contribuent à la contrainte.
• Mise à jour Duale (Multiplicateur) : Le multiplicateur de Lagrange $\lambda_t$ est mis à jour en utilisant la descente de miroir basée sur la violation cumulative de la contrainte (marge) observée durant la génération :
  $$\lambda_t = \lambda_T \exp\left\{ -\eta \sum_{\tau=t}^{T-1} \sum_{\ell, j} \left(b_{\ell j} - \frac{R}{L}\right) a_{\ell j}^\tau \right\}$$
  Ici, $\eta$ est un paramètre de taille de pas. Si la contrainte est sous-satisfaite, $\lambda$ augmente, renforçant le biais ; si elle est satisfaite, $\lambda$ diminue, réduisant l'interférence avec la distribution originale.

3. Variantes

L'article explore différentes définitions de « marge » pour la mise à jour de $\lambda$ :

• Marge accumulée : Agrège les contributions au fil du temps (utilisée dans les expériences principales).
• Marge instantanée : Compare l'accumulation actuelle directement à l'objectif global.
• Marge optimiste : Intègre une prédiction de la satisfaction future de la contrainte (Descente de miroir optimiste).

Contributions clés

1. Efficacité au moment de l'inférence : La méthode ne nécessite aucun réentraînement et aucun modèle auxiliaire. Elle n'ajoute que des opérations en temps constant à la boucle d'échantillonnage standard, préservant le profil computationnel du modèle de base (coût 1×).
2. Garanties théoriques : Les auteurs dérivent des bornes formelles sur la violation attendue des contraintes, montrant que la violation est majorée par $O(1/\eta \log(\dots))$. Ils fournissent également des bornes de regret caractérisant le compromis entre la satisfaction des contraintes et l'erreur d'approximation.
3. Prise en charge de multiples contraintes : La formulation primal-dual s'étend naturellement à plusieurs contraintes simultanées en sommant les biais de logits additifs provenant de multiples multiplicateurs de Lagrange.
4. Compromis contrôlable : Le paramètre $\eta$ agit comme un seul « bouton » permettant aux praticiens de tracer une frontière lisse entre la satisfaction des contraintes et la fidélité distributionnelle (mesurée par la divergence KL).

Résultats expérimentaux

La méthode a été évaluée sur trois domaines distincts en utilisant un algorithme unique instancié avec des fonctions de scoring spécifiques au domaine :

1. Génération de texte thématique (TinyStories) :

   • Tâche : Générer des histoires avec au moins 10 jetons liés à l'océan.
   • Résultat : SPDD a atteint 92,1 % de taux de réussite (satisfaction) avec une divergence KL de 0,420. Cela surpasse nettement CDD (38,0 % de taux de réussite) et le Biais statique (81,9 % de taux de réussite) tout en maintenant une KL inférieure à celle du Biais statique. D-CBG n'a pas réussi à satisfaire la contrainte sparse.

2. Génération moléculaire (ZINC20) :

   • Tâche : Générer des molécules avec une masse moléculaire (MW) $\ge 350$.
   • Résultat : SPDD a atteint 48,5 % de taux de réussite (au-dessus de 350 MW) avec 49,9 % de validité, surpassant toutes les lignes de base. Il a déplacé la distribution plus loin dans la queue supérieure de la distribution MW que CDD ou D-CBG tout en maintenant la validité chimique.

3. Génération de listes de lecture musicales :

   • Tâche : Générer des listes de lecture avec une composition de sous-genres cible.
   • Résultat : SPDD a égalé la satisfaction au niveau des pistes de CDD (9,28 % contre 9,18 %) mais avec moins de la moitié de la divergence KL (0,355 contre 0,779) et une diversité de jetons significativement plus élevée.

Dans tous les domaines, SPDD a atteint la satisfaction de contraintes la plus élevée tout en préservant la proximité la plus forte avec la distribution du modèle non contraint par rapport aux lignes de base.

Importance et affirmations

L'article affirme que le Décodage guidé primal-dual offre une solution légère et théoriquement fondée à la diffusion discrète contrainte.

• Efficacité : Contrairement à CDD, il évite les boucles d'optimisation internes coûteuses. Contrairement à D-CBG, il ne nécessite pas l'entraînement de classificateurs séparés pour chaque contrainte.
• Flexibilité : Il prend en charge plusieurs contraintes et s'adapte dynamiquement au contexte de génération via le multiplicateur de Lagrange.
• Robustesse : L'analyse théorique fournit des bornes sur la violation des contraintes, et les résultats empiriques montrent que la méthode fonctionne sur divers types de données (texte, chimie, plongements audio) avec un seul cadre algorithmique.
• Limites : Les auteurs reconnaissent que dans la diffusion à état absorbant, les engagements précoces de jetons sont irréversibles, rendant la conception de la mise à jour de la marge critique dans les premières étapes. Ils notent également que pour les domaines hautement structurés (comme les molécules), l'hypothèse théorique de dérives de logits de jetons indépendantes est une approximation, bien que la méthode reste empiriquement efficace.

Ce travail se positionne comme un outil pratique pour les praticiens qui doivent orienter les modèles génératifs vers des propriétés spécifiques sans la surcharge du réentraînement ou de l'optimisation complexe au moment de l'inférence.