Consistent Diffusion Language Models

1. Énoncé du problème

Les modèles de diffusion linguistique (DLM) offrent une alternative prometteuse aux modèles autoregressifs (AR) en permettant une génération parallèle de tokens et une inférence en temps sous-linéaire. Cependant, un goulot d'étranglement critique subsiste : la génération de texte de haute qualité nécessite généralement des centaines d'étapes de raffinement, annulant ainsi les gains d'efficacité promis par le parallélisme.

Dans les domaines continus (par exemple, les images), les modèles de cohérence ont réussi à accélérer la diffusion en entraînant des modèles à être invariants le long d'une trajectoire unique et déterministe définie par l'équation différentielle ordinaire du flux de probabilité (PF-ODE).

• Le fossé : Dans les domaines discrets (comme le texte), aucune PF-ODE analogue dans l'espace des échantillons n'existe. Les processus de corruption discrets (par exemple, le masquage ou le bruit uniforme) ne suivent pas un chemin déterministe unique ; ils impliquent plutôt des transitions stochastiques.
• La conséquence : Adapter directement les méthodes de cohérence continues aux données discrètes est mal défini. Les solutions existantes reposent souvent sur des pipelines de distillation multi-étapes (utilisant un modèle enseignant) ou sur des substituts de relaxation continue, qui sont complexes, coûteux en calcul et souffrent souvent d'effondrement de mode ou de perte de diversité.

2. Méthodologie : Cohérence discrète multi-chemins (MPDC)

Les auteurs proposent un changement fondamental : au lieu de chercher un chemin déterministe inexistant, ils exploitent les ponts postérieurs exacts inhérents aux processus de diffusion discrets.

Concept central : Le pont postérieur

Pour deux niveaux de bruit quelconques $s < t$ dans un processus de diffusion discret, il existe un pont postérieur exact sous forme close $q(x_s | x_t, x_0)$. Ce pont définit un chemin stochastique valide reliant un état bruité $x_t$ à un état plus propre $x_s$, conditionné aux données réelles $x_0$.

• Contrairement à la PF-ODE continue, ce pont est stochastique.
• Crucialement, ces ponts forment un semi-groupe (propriété de Chapman-Kolmogorov), ce qui signifie que les chemins peuvent être composés de manière transitive.

Le principe MPDC

Les auteurs introduisent la cohérence discrète multi-chemins (MPDC), qui entraîne un débruiteur à être invariant de chemin en espérance.

• Définition : Un modèle $f_\theta(x_t, t)$ est cohérent si sa prédiction à un état bruité $x_t$ s'accorde (en espérance) avec la prédiction faite après avoir effectué un "saut" intermédiaire vers un état plus propre $x_s$ via le pont postérieur exact.
• Objectif : Le modèle minimise la divergence entre sa prédiction directe depuis $x_t$ et la prédiction espérée depuis $x_s$ (où $x_s$ est échantillonné à partir du pont).
• Formulation mathématique :
  $$\mathcal{L}_{CDLM} = \mathbb{E}_{t, \delta, x_t} \left[ \sum_{i \in \mathcal{M}(x_t)} w(t, \delta) \cdot D\left( f_\theta(x_t, t)_i \parallel \tilde{f}_\theta(x_s, s)_i \right) \right]$$
  Où $x_s \sim q(x_s | x_t, x_0)$, $\tilde{f}_\theta$ est un réseau cible (EMA), et $D$ est une métrique de divergence (généralement la divergence de Jensen-Shannon).

Choix de conception clés pour la stabilité

Pour éviter l'effondrement de mode (où le modèle génère du texte répétitif pour satisfaire trivialement la cohérence) et la dérive uniforme, les auteurs introduisent trois mécanismes stabilisateurs :

1. Taille de pas comme curriculum : Échantillonnage des tailles de pas $\delta$ à partir d'une plage spécifique (par exemple, de $1/8$ à $3/8$) et pondération par $1/\delta$ pour équilibrer les signaux d'apprentissage entre les chemins courts et longs.
2. Ancre de diffusion à pas maximal : Une petite fraction de la perte d'entraînement est ancrée aux données réelles ($s=0$) en utilisant un objectif de diffusion standard. Cela empêche le modèle de converger vers des solutions constantes dégénérées.
3. Asymétrie d'optimisation et divergence : Utilisation d'un stop-gradient sur le réseau cible et emploi de la divergence de Jensen-Shannon (JSD) plutôt que des KL avant/arrière. La JSD offre un compromis stable entre la recherche de modes et la couverture de modes, évitant l'instabilité de la divergence KL lors de l'entraînement à partir de zéro.

3. Contributions clés

1. Un nouveau principe pour la modélisation générative discrète : Le papier établit que les ponts postérieurs exacts servent de substitut rigoureux et sous forme close à la PF-ODE manquante, transformant l'obstacle de l'"accélération discrète" en un atout analytique.
2. Un cadre d'entraînement unifié : CDLM fournit un objectif en une seule étape, sans enseignant. Les auteurs prouvent que la diffusion masquée standard, les modèles de cohérence continus, la distillation progressive et les récentes méthodes de distillation en deux étapes (comme SDTT et DUO-DCD) sont tous des cas particuliers, des limites ou des couplages approximatifs du cadre CDLM.
3. Performances de pointe : CDLM obtient des résultats supérieurs dans la génération de texte conditionnelle et inconditionnelle, surpassant les modèles de base puissants et égalant ou dépassant les modèles distillés complexes multi-étapes sans nécessiter d'enseignant.

4. Résultats expérimentaux

Les auteurs ont évalué CDLM (spécifiquement MCDLM pour la diffusion masquée et UCDLM pour la diffusion uniforme) par rapport à des modèles de base tels que MDLM, DUO, SDTT et DUO-DCD sur le jeu de données OpenWebText.

• Génération inconditionnelle :

  • Efficacité en peu d'étapes : CDLM établit de nouveaux résultats de pointe (SOTA) pour les modèles en une seule étape. Il surpasse nettement les modèles de base (MDLM, DUO) sur tous les budgets d'échantillonnage.
  • Comparaison avec la distillation : De manière remarquable, le CDLM en une seule étape surpasse souvent les modèles distillés multi-étapes (SDTT, DUO-DCD) en termes de perplexité (PPL) tout en maintenant une entropie (diversité) beaucoup plus élevée.
  • Effondrement de mode : Les modèles distillés utilisant des échantillonneurs gourmands (par exemple, DUO-DCD) ont montré une faible entropie (3,9), indiquant un effondrement de mode sévère (caractères répétitifs). CDLM a maintenu une entropie saine (5,0+), prouvant sa capacité à préserver la diversité.
  • Vitesse : CDLM réalise des accélérations de 32x à 64x par rapport aux modèles autoregressifs de base et de 64x à 128x par rapport au MDLM standard, tout en maintenant une qualité comparable.

• Génération conditionnelle :

  • CDLM a démontré des performances robustes dans la récupération de phrases originales à partir d'entrées corrompues à 50 % sur plusieurs jeux de données (OpenWebText, Lambada, Wikitext-103, PTB).
  • Il a obtenu de meilleurs scores BLEU et une perplexité plus faible que SDTT et DUO, en particulier dans la préservation des conditions d'entrée où d'autres modèles ont échoué.

5. Signification et impact

• Unification théorique : Le papier offre une "lentille unificatrice" pour la modélisation générative discrète, montrant que diverses techniques d'accélération sont fondamentalement connectées par le concept de cohérence de chemin.
• Évolutivité : En éliminant le besoin de distillation multi-étapes et de modèles enseignants, CDLM offre une base principée, évolutive et sans enseignant pour les modèles génératifs discrets rapides de nouvelle génération.
• Généralisabilité : Bien que démontré sur le texte, le cadre s'applique à tout domaine discret avec un pont postérieur traitable (par exemple, les séquences biologiques, les graphes, la synthèse de programmes).
• Changement de paradigme : Il redéfinit la diffusion discrète rapide non pas comme une approximation des EDO continues, mais comme l'apprentissage d'un débruiteur indépendant du chemin sur des ponts stochastiques, résolvant avec succès le problème de la génération en peu d'étapes dans les espaces discrets.

En conclusion, CDLM comble le fossé entre l'efficacité des modèles de cohérence continus et les réalités des données discrètes, offrant une solution robuste en une seule étape qui réalise une génération de texte rapide et de haute fidélité sans la complexité des pipelines de distillation.