Résumé technique : Localisation stochastique discrète pour la génération non-autorégressive

1. Énoncé du problème

Les modèles de diffusion continus sont devenus une référence pour la génération de données de haute dimension dans les domaines de l'image, de la vidéo et de l'audio. Cependant, leur application à la génération de langage via une diffusion continue sur les plongements de tokens a systématiquement sous-performé par rapport aux modèles de diffusion discrète masquée (MDM). Les modèles de langage à diffusion continue existants accusent généralement un retard par rapport aux MDM en termes de fidélité distributionnelle (mesurée par des métriques telles que MAUVE).

Les auteurs soutiennent que le goulot d'étranglement ne réside pas dans la continuité elle-même, mais dans la représentation utilisée. Les approches standard (préservation de la variance ou explosion de la variance) représentent la limite d'information nulle comme un bruit gaussien isotrope plutôt que comme un état de token de type masque. De plus, dans ces cadres, le débruiteur doit recevoir une étiquette de pas de temps ou de niveau de bruit pour calibrer la quantité de preuves contenue dans le plongement bruité. Cela rend les débruiteurs hautement sensibles à la condition temporelle, nécessitant un traitement distinct pour les états masqués par rapport aux états incertains. Par conséquent, les modèles de diffusion continue n'ont pas réussi à exploiter la supervision plus riche fournie par le continuum d'états finis de rapport signal-sur-bruit (SNR).

2. Méthodologie : Localisation stochastique discrète (DSL)

L'article introduit la localisation stochastique discrète (DSL), un cadre à état continu conçu pour rendre le processus de débruitage invariant au SNR nominal.

Choix de conception fondamentaux

1. Plongements de tokens sur la sphère unité : DSL contraint tous les plongements de tokens propres à se situer sur la sphère unité ($\|x_i\|_2 = 1$). Il s'agit d'une contrainte structurelle sur la géométrie, et non d'un régularisateur de perte.
2. Canal de localisation stochastique : Le processus de corruption est défini par $z_t = t x + \sqrt{t} \epsilon$, où $t$ représente le SNR. Sous la contrainte de la sphère unité, le terme impliquant la norme de $x$ dans le posterior bayésien devient constant et s'annule lors de la normalisation.
3. Débruiteur invariant au SNR : Un résultat théorique clé est que, sous cette géométrie, le débruiteur optimal au sens de Bayes (la moyenne du posterior) dépend uniquement de l'état du token bruité $z$, et non du SNR nominal $t$.
   $$\hat{x}(z, t) = \hat{x}(z) = \frac{\mathbb{E}_P[x e^{x \cdot z}]}{\mathbb{E}_P[e^{x \cdot z}]}$$
   Cela permet à un réseau unique, agnostique au temps, de servir de débruiteur pour tous les régimes de SNR (du entièrement masqué au presque propre).

Entraînement et architecture

• Entraînement à support mixte : Pour exploiter à la fois la vraisemblance de l'intégrale de chemin continue et la vraisemblance de l'extrémité de la diffusion masquée, DSL emploie un objectif d'entraînement à support mixte. Il échantillonne des SNR à partir de deux branches :
  • Chemin continu : SNR tirés d'une distribution log-normale (états à SNR fini).
  • Extrémité ROAR (Autorégressive à ordre aléatoire) : SNR tirés de valeurs extrêmes (proches de 0 pour masqué, proches de $\infty$ pour révélé).
    Les deux branches sont entraînées en utilisant une seule perte d'entropie croisée sur le posterior des tokens, qui agit comme un substitut à l'erreur quadratique moyenne (MSE) de l'intégrale de chemin, tout en évitant les pathologies d'optimisation dans l'espace des plongements (par exemple, l'effondrement des plongements).
• Convertisseur de vue posterior : Puisque l'état bruité $z$ réside dans un espace de plongement continu mais que le modèle de base (Transformer/DiT) s'attend à des entrées de type token, un convertisseur mappe $z$ vers une représentation de « mélange de tokens ». Cela mappe l'état bruité vers une combinaison bornée de plongements de norme unité, préservant la géométrie du posterior (la direction contrôle la préférence de token, la norme contrôle la confiance) tout en assurant une attention self-stable.

Flexibilité de l'inférence

Comme le débruiteur est agnostique au temps, DSL prend en charge des chemins de SNR arbitraires par token au sein d'une seule exécution de génération :

• Raffinement masqué : Remasquage itératif standard (ReMDM).
• Autorégression à ordre aléatoire (ROAR) : Révélation des tokens un par un dans un ordre aléatoire sans remasquage.
• Échantillonnage hybride : Une étape de débruitage continue suivie d'une étape de raffinement masqué discrète.

3. Contributions clés

1. Un cadre à état continu pour la génération discrète : DSL introduit un canal de localisation stochastique sur des plongements de sphère unité où le débruiteur optimal au sens de Bayes est une fonction de l'état bruité uniquement. Cela permet à un réseau unique agnostique au temps de supporter toute une famille de chemins de SNR par token.
2. Une recette d'entraînement pratique : Les auteurs proposent une stratégie d'échantillonnage de SNR à support mixte qui entraîne le même débruiteur sur les états d'extrémité (masqué/révélé) et les états à SNR fini. Ils utilisent un convertisseur de mélange de tokens pour aligner la vue du posterior continu avec les architectures Transformer standard.
3. Preuves à travers les chemins : Un seul point de contrôle DSL démontre une génération en peu d'états améliorée sur OpenWebText sous raffinement masqué, prend en charge l'échantillonnage AR à ordre aléatoire sans réentraînement, et atteint une vraisemblance compétitive sur Text8.

4. Résultats expérimentaux

Les auteurs ont évalué DSL en affinant un point de contrôle MDLM pré-entraîné sur OpenWebText (OWT) et Text8.

• Génération OpenWebText (MAUVE) :
  • Sans raffinement : Même sans remasquage itératif, le point de contrôle affiné avec DSL surpasse significativement le MDLM de base. À $T=1024$, le MAUVE passe de 0,042 (MDLM) à 0,495 (DSL-FT + MDLM) et 0,551 (DSL-FT + ROAR).
  • Raffinement itératif (ReMDM) : Combiné avec des décodeurs de la famille ReMDM, DSL a atteint des résultats de pointe parmi les bases de référence sans référence. La meilleure configuration (boucle ReMDM avec $\eta_{cap}$ principiel) a atteint un MAUVE de 0,722 à $T=1024$, surpassant les méthodes précédentes telles que PRISM-loop (0,527) et ADLM (0,699).
  • Efficacité en peu d'états : DSL a montré de fortes performances avec des budgets d'états faibles. À $T=128$, DSL-FT + boucle ReMDM a atteint un MAUVE de 0,639, contre 0,015 pour le MDLM de base.
• Échantillonnage hybride : Un échantillonneur hybride (débruitage continu suivi d'un raffinement discret) a atteint un MAUVE de 0,662 en seulement 48 étapes au total, dépassant le résultat du MDLM sans raffinement à $T=1024$.
• Vraisemblance Text8 : DSL a atteint un nombre de bits par caractère (BPC) de $\le$ 1,45, s'approchant de la plage des modèles de diffusion discrète masquée (MDLM $\le$ 1,40, MD4 $\le$ 1,37) et s'améliorant par rapport aux bases de référence de diffusion continue antérieures (Plaid $\le$ 1,48).

5. Importance et revendications

L'article revendique que l'avantage principal de DSL est l'unification des chemins d'inférence. En paramétrant la diffusion continue comme un débruitage d'état posterior sur des tokens bruités (plutôt que comme un débruitage indexé par le pas de temps), DSL permet :

• La révélation d'extrémité, le raffinement masqué et le débruitage continu/hybride d'être servis par un seul estimateur de posterior.
• La compatibilité des points de contrôle : Un seul modèle peut être interrogé le long de multiples chemins de SNR (par exemple, AR à ordre aléatoire, raffinement masqué, hybride) sans réentraînement ni distillation.
• Une fidélité distributionnelle améliorée : Les états finis de SNR continus fournissent une supervision plus riche, conduisant à des scores MAUVE nettement meilleurs sur tous les budgets d'états par rapport aux bases de référence standard de diffusion continue ou de diffusion masquée.

Les auteurs concluent que l'ingrédient manquant dans la diffusion continue pour le langage n'était pas la continuité, mais le choix de la représentation. En adoptant des plongements de sphère unité et une localisation stochastique, ils comblent le fossé entre diffusion continue et diffusion discrète, permettant à un seul modèle de supporter diverses stratégies de génération avec une haute efficacité et qualité.

Limites : L'évaluation est actuellement limitée à OpenWebText et Text8. Les échantillonneurs rapportés (famille ReMDM, hybride) ne sont pas des politiques de décodage exhaustivement optimisées mais servent à démontrer la flexibilité du posterior DSL sous-jacent. Les directions de travail futures incluent l'optimisation des calendriers de remasquage, des règles de transfert hybride, et l'exploration de la génération conditionnelle.