Résumé technique : SMOOTHIE (Lissage de la diffusion sur les plongements de tokens)

Énoncé du problème

Les modèles de diffusion ont obtenu des résultats de pointe dans des domaines continus comme les images et l'audio, mais rencontrent des défis majeurs lorsqu'ils sont appliqués à la génération de texte en raison du caractère discret du langage. Les approches existantes se divisent généralement en deux catégories, chacune présentant des limitations distinctes :

1. Diffusion gaussienne dans les espaces latents : Ces méthodes appliquent un bruit gaussien aux plongements de tokens continus. Bien qu'elles préservent la structure sémantique (selon l'hypothèse de l'espace sémantique euclidien), elles peinent avec l'étape de décodage discret, nécessitant souvent une correspondance explicite des vecteurs continus vers les tokens, ce qui peut être instable ou entraînant une perte d'information.
2. Diffusion catégorielle/simplexe : Ces méthodes opèrent directement sur le simplexe de probabilités du vocabulaire ou sur des distributions catégorielles. Bien qu'elles respectent le caractère discret du texte et éliminent le besoin d'un décodage explicite, elles négligent généralement les relations sémantiques entre les tokens durant le processus d'ajout de bruit. Cela conduit à une dégradation de l'information plus erratique et moins significative.

Le problème central est l'absence d'un cadre de diffusion qui respecte simultanément le caractère discret du texte et supprime progressivement l'information sémantique d'une manière qui maintient la cohérence sémantique durant le processus direct.

Méthodologie : SMOOTHIE

Les auteurs proposent SMOOTHIE (Lissage de la diffusion sur les plongements de tokens), un cadre novateur qui comble le fossé entre la diffusion continue et la diffusion discrète en lissant les plongements de tokens basés sur la similarité sémantique.

1. Processus de diffusion directe (Ajout de bruit)

Au lieu d'ajouter du bruit directement aux plongements ou aux probabilités catégorielles, SMOOTHIE représente chaque token $w_i$ dans une séquence cible comme un vecteur de distances euclidiennes au carré négatives entre son plongement et les plongements de tous les tokens du vocabulaire :
$$D_0 = \left( -\|E_{w_i} - E_j\|^2 \right)_{i,j=1}^{m,V}$$
où $E$ est une matrice de plongement fixe (par exemple, issue de BERT), $m$ est la longueur de la séquence et $V$ est la taille du vocabulaire.

Le processus direct perturbe progressivement ces vecteurs de distances en ajoutant un bruit gaussien :
$$q(D_t | D_0) = \mathcal{N}\left(D_t \mid \frac{1}{\sigma_t^2}D_0, \delta^2 I\right)$$
Ici, $\sigma_t$ est un planificateur de bruit contrôlant la bande passante, et $\delta$ contrôle la stochasticité. Le vecteur bruité $D_t$ est converti en une distribution de probabilité sur le vocabulaire à l'aide d'une fonction softmax : $p_t = \text{softmax}(D_t)$.

Insight clé : À mesure que le bruit augmente (et que $\sigma_t$ croît), la masse de probabilité initialement centrée sur le token correct se distribue progressivement vers d'autres tokens. Crucialement, parce que l'entrée est basée sur des distances euclidiennes, la masse de probabilité se diffuse d'abord vers les tokens sémantiquement similaires (ceux avec des distances euclidiennes plus petites) avant d'atteindre les tokens distants. Cela imite un principe de "dissipation de chaleur", assurant que l'information sémantique est supprimée progressivement et non de manière chaotique.

2. Processus de diffusion inverse (Débruitage)

Le processus inverse est modélisé comme une chaîne de Markov où un réseau de neurones $f_\theta$ prédit le vecteur de distance original $D_0$ (ou une version mise à l'échelle) à partir de l'entrée bruitée $p_t$. Le modèle est entraîné pour minimiser l'erreur quadratique moyenne (MSE) entre les vecteurs de distance prédits et réels :
$$L(\theta) = \mathbb{E}[\|D_0 - g_\theta(p_t, t)\|^2]$$
Les auteurs démontrent que prédire le vecteur de distance équivaut à prédire le plongement attendu, permettant au modèle de produire une estimation de plongement propre qui est ensuite convertie en un token via une recherche du plus proche voisin ou un argmax.

3. Échantillonnage et solveur

Les auteurs introduisent un schéma d'échantillonnage spécifique (Éq. 9) qui met à jour l'état latent en appliquant le processus direct directement à la prédiction du modèle, plutôt que de s'appuyer sur l'état latent précédent comme le font les solveurs DDPM standards.
$$D_{t-1} = \frac{1}{\sigma_{t-1}^2}D_0(f_\theta(p_t, t)) + \tilde{\delta}\epsilon$$
Cette approche ignore l'état caché précédent durant l'étape de mise à jour. Les auteurs soutiennent que cela est bénéfique pour la distribution de données sparse des plongements de texte, réduisant le risque que la représentation latente dérive hors de la variété des données.

4. Architecture et conditionnement

• Architecture : Basée sur des couches de décodeur Transformer avec des connexions de saut de style UNet.
• Conditionnement : Pour les tâches séquence-à-séquence, un encodeur Transformer supplémentaire traite la séquence source, intégré via une attention croisée.
• Auto-conditionnement : Le modèle utilise un auto-conditionnement, alimentant sa propre prédiction précédente comme entrée pour les étapes suivantes afin d'améliorer la stabilité et la qualité.

Contributions clés

1. Cadre de diffusion novateur : SMOOTHIE est la première méthode de diffusion textuelle qui respecte simultanément le caractère discret du texte et supprime progressivement l'information sémantique en lissant les plongements de tokens basés sur des distances euclidiennes.
2. Généralisation de la diffusion simplexe : Les auteurs démontrent théoriquement que leur espace latent basé sur les distances généralise l'espace latent basé sur le simplexe (utilisé dans SSD-LM et TESS). La diffusion simplexe est montrée comme un cas particulier de SMOOTHIE où la métrique de distance est triviale (0 pour le même token, 1 sinon).
3. Solveur d'échantillonnage amélioré : L'article propose un solveur qui ignore l'état caché précédent durant les mises à jour, ce qui surpasse empiriquement les solveurs DDPM et DDIM standards sur les tâches séquence-à-séquence.
4. Supériorité empirique : Des expériences extensives montrent que SMOOTHIE surpasse les modèles de diffusion existants (y compris les variantes basées sur les plongements et sur le simplexe) et atteint une qualité comparable aux modèles autorégressifs (comme BART et GPT-2) sur plusieurs tâches.

Résultats expérimentaux

Les auteurs ont évalué SMOOTHIE sur une tâche de génération non conditionnée (ROCStories) et quatre tâches séquence-à-séquence (Quora Question Pairs, Quasar-T, XSum, ParaDetox).

• Performance : SMOOTHIE a constamment surpassé les autres bases de diffusion (DiffuSeq, TESS, SSD-LM, TEncDM, etc.) et les implémentations alternatives d'espaces latents (diffusion basée sur les plongements et diffusion simplexe entraînées dans le même cadre).
• Métriques :
  • XSum : A obtenu un BERTScore de 68,8 et un ROUGE-L de 26,0, surpassant les autres méthodes de diffusion et s'approchant des bases autorégressives.
  • ROCStories : A obtenu le score MAUVE le plus élevé (80,7) parmi les modèles de diffusion, égalant presque GPT-2 (78,9), tout en maintenant une diversité élevée.
  • ParaDetox : A obtenu le score J le plus élevé (51,7) pour la détoxification de texte.
• Études d'ablation :
  • $\tilde{\delta}$ (Stochasticité) : Des valeurs plus faibles de $\tilde{\delta}$ (par exemple 0,1) ont amélioré la perplexité et la cohérence pour les tâches séquence-à-séquence, tandis que des valeurs légèrement plus élevées étaient optimales pour la génération non conditionnée afin de maintenir la diversité.
  • Solveur : Le solveur proposé (Éq. 9) a constamment surpassé DDPM et DDIM sur toutes les tâches.
  • Plongements : Le modèle repose sur des plongements BERT pré-entraînés. L'ablation a montré que bien que les plongements GPT-2 aient donné des performances similaires, les plongements GloVe ont résulté en une qualité inférieure, soulignant l'importance de la qualité des plongements.

Importance et affirmations

L'article affirme que SMOOTHIE représente une avancée significative dans la diffusion textuelle en résolvant le compromis entre la préservation sémantique et la structure discrète. En intégrant directement les relations sémantiques dans le processus d'ajout de bruit via des distances euclidiennes, la méthode permet un "processus de décodage naturel" où l'information est supprimée progressivement et de manière significative.

Les auteurs notent modestement que bien que la méthode soit applicable à d'autres domaines catégoriels avec des similarités inhérentes (par exemple, les graphes, les séquences de protéines), elle repose actuellement sur des plongements pré-entraînés plutôt que sur un entraînement de bout en bout, ce qui reste un défi en raison de problèmes de stabilité avec la perte MSE. Ils reconnaissent également le goulot d'étranglement computationnel du calcul des distances par paires à chaque étape, mais soutiennent que les gains de performance justifient le coût, et que la méthode offre une alternative robuste aux modèles autorégressifs pour les tâches nécessitant une grande diversité et une contrôlabilité.

Le travail suggère que le choix de l'espace latent et du solveur de diffusion spécifique sont aussi critiques que l'architecture du modèle elle-même, SMOOTHIE fournissant un cadre unifié qui exploite les forces des approches de diffusion continue et discrète.