Résumé Technique : Adaptive Block Diffusion (ABD)

1. Énoncé du Problème

Les modèles de langage par diffusion (DLM) offrent une alternative prometteuse aux modèles autorégressifs (AR) en permettant un décodage parallèle, ce qui peut réduire la latence d'inférence. Cependant, les approches récentes de diffusion structurée, telles que les modèles de langage à diffusion discrète masquée (MDLM) et la diffusion par blocs (ex: BD3LM), reposent sur des structures de contexte fixes. Dans ces modèles, l'ensemble des jetons à débruiter suit un schéma prédéterminé (par exemple, des tailles de blocs fixes et des préfixes alignés) tout au long de l'entraînement.

Cette conception crée un décalage fondamental entre l'entraînement et l'inférence :

• Entraînement : Le modèle optimise le risque de débruitage sur un sous-ensemble restreint et fixe de configurations (la « grille d'entraînement »).
• Inférence : Les stratégies de décodage pratiques nécessitent souvent des longueurs de préfixe et des tailles de fenêtres arbitraires (par exemple, des fenêtres glissantes, des tailles de blocs adaptatives ou du streaming).

Lorsque les modèles à blocs fixes sont évalués sur des configurations hors de leur support d'entraînement (« off-grid »), ils présentent une dégradation brutale de la perplexité et un comportement de génération instable. Ils ne parviennent pas à apprendre une procédure de débruitage générale, mais se spécialisent uniquement dans les configurations spécifiques vues durant l'entraînement. L'article postule qu'un modèle de diffusion unique devrait être capable d'éliminer ce décalage pour rester invariant aux configurations de décodage.

2. Méthodologie : Adaptive Block Diffusion (ABD)

Les auteurs proposent l'Adaptive Block Diffusion (ABD), un cadre qui résout le décalage entraînement-inférence en traitant la configuration de débruitage comme une variable stochastique plutôt que comme un hyperparamètre fixe.

Mécanisme Central

Au lieu de restreindre l'entraînement à un sous-ensemble aligné sur des blocs fixes, l'ABD définit un espace de configuration $\Omega := \{(k, \ell) \mid 0 \le k \le L-1, 1 \le \ell \le L-k\}$, où $k$ est la longueur du préfixe et $\ell$ est la longueur de la fenêtre de débruitage.

• Entraînement Stochastique : Le modèle est entraîné sur une distribution $\pi$ définie sur $\Omega$. À chaque étape d'entraînement, une configuration $(k, \ell)$ est échantillonnée de $\pi$.
• Processus de Débruitage : Les jetons indexés par la fenêtre $S(k, \ell) = \{k+1, \dots, k+\ell\}$ sont corrompus par le processus de diffusion directe, tandis que le préfixe $x_{1:k}$ reste propre et sert de contexte de conditionnement.
• Objectif : L'objectif d'entraînement minimise le risque de débruitage attendu sur la distribution $\pi$ :
  $$\mathcal{L}_{ABD}(\theta) = \mathbb{E}_{(k,\ell)\sim\pi, x_0\sim D, t} \left[ \sum_{i \in S(k,\ell)} \text{KL}\left(q(x^i_0 | x^i_t, t) \parallel p_\theta(x^i_0 | x^i_t, x^i_{1:k}, t)\right) \right]$$

Détails d'Implémentation

• Architecture : L'ABD ne nécessite aucun changement architectural au modèle de diffusion sous-jacent (par exemple, le transformer). Il modifie uniquement la distribution des configurations rencontrées lors de l'entraînement.
• Masquage d'Attention : Un masque d'attention dynamique est construit pour chaque étape en fonction du $(k, \ell)$ échantillonné. Ce masque garantit que les jetons bruités n'attendent que les autres jetons bruités au sein du même bloc et les jetons propres dans les blocs précédents, tandis que les jetons propres attendent causalement les jetons propres précédents.
• Choix de la Distribution : L'article explore diverses distributions pour $\pi$ (exponentielle catégorielle, uniforme, log-normale). Dans leurs expériences principales, ils utilisent une distribution catégorielle sur les puissances de deux $\{1, 2, 4, \dots, 128\}$ avec une décroissance exponentielle, ce qui privilégie les petites tailles de blocs.

3. Analyse Théorique

L'article fournit des garanties formelles liant le support de la distribution d'entraînement à la performance de généralisation.

• Minimisation du Risque : Les auteurs formalisent l'entraînement de la diffusion comme une minimisation du risque sur l'espace de configuration $\Omega$.
• Théorème 1 (Optimalité) : Tout minimiseur global de l'objectif ABD atteint une performance de débruitage optimale pour presque tout configuration $(k, \ell)$ selon $\pi$. Les configurations en dehors du support de $\pi$ restent non contraintes.
• Théorème 2 (Alignement Entraînement–Inférence) : Si la politique d'inférence induit une distribution $\Pi_{inf}$ telle que $\Pi_{inf} \ll \pi$ (c'est-à-dire que le support d'inférence est contenu dans le support d'entraînement), alors le modèle entraîné pour minimiser le risque sous $\pi$ est également un minimiseur global pour le risque d'inférence.
• Implication : La généralisation n'est pas une propriété de l'architecture du modèle, mais du couverture du support de la distribution d'entraînement. Les modèles à blocs fixes échouent car $\Pi_{inf} \not\ll \pi_{fixed}$ pour des stratégies de décodage arbitraires.

4. Résultats Expérimentaux

Les auteurs évaluent l'ABD sur les benchmarks LM1B (One Billion Word) et OpenWebText (OWT), en comparant avec l'Autorégressif (AR), MDLM, SEDD et les spécialistes à blocs fixes (BD3LM).

Principales Conclusions

1. Invariance Structurelle :
   • Les spécialistes à blocs fixes (ex: BD3LM entraîné sur une taille de bloc de 4) présentent des pics de perplexité brutaux lorsqu'ils sont évalués sur des tailles de blocs supérieures à leur configuration d'entraînement.
   • L'ABD rétablit une relation monotone entre la taille du bloc et la perplexité sur toute la plage. À mesure que la taille du bloc diminue (se rapprochant de la précision autorégressive), la perplexité s'améliore de manière monotone, sans effondrement hors grille.
2. Parité de Performance :
   • Un seul modèle ABD égale ou surpasse les spécialistes à blocs fixes à leurs configurations « domicile » respectives. Par exemple, sur LM1B, l'ABD entraîné sur une distribution large atteint une perplexité compétitive à des tailles de blocs 4 et 16 par rapport aux spécialistes BD3LM entraînés spécifiquement pour ces tailles.
3. Généralisation Zero-Shot :
   • Sur des benchmarks non vus (PTB, Wikitext, Lambda, etc.), l'ABD démontre une généralisation robuste. Plus précisément, l'ABD à une taille de bloc de 4 surpasse les spécialistes BD3LM sur les domaines scientifiques (Pubmed, Arxiv), suggérant qu'un entraînement multi-échelle améliore la robustesse au changement de domaine.
4. Qualité Générative :
   • En termes de perplexité générative (mesurée via GPT2-Large), l'ABD égale la qualité des plus puissants spécialistes à blocs fixes à coûts de calcul équivalents (NFE), confirmant qu'un entraînement large n'altère pas la qualité des échantillons.

Études d'Ablation

• Sensibilité de la Distribution : Le choix de $\pi$ permet aux praticiens de négocier le compromis entre les échelles. Le pondération exponentielle favorise les petits blocs ; la pondération uniforme favorise les grands blocs ; la log-normale offre un équilibre.
• Dynamique d'Entraînement : L'entraînement multi-échelle agit comme un régulariseur stable, améliorant les performances pour toutes les tailles de blocs simultanément plutôt que de se spécialiser dans une seule.
• Discret vs Continu : Les distributions discrètes sur les puissances de deux se sont révélées plus maniables et performantes que les distributions continues, probablement en raison des points de fonctionnement distincts qu'elles représentent.

5. Signification et Revendications

L'article affirme que le support de configuration est le facteur central régissant la généralisation dans les modèles de langage par diffusion. En déplaçant l'attention des contraintes architecturales fixes vers la distribution des configurations d'entraînement, l'ABD établit une alternative principled à des objectifs à structure fixe.

• Résolution du Décalage : L'ABD élimine avec succès le décalage entraînement-inférence, permettant à un seul modèle de gérer des stratégies de décodage arbitraires (fenêtres glissantes, blocs adaptatifs) sans changements architecturaux ni heuristiques au moment de l'inférence.
• Modèle Unifié : Il démontre qu'un seul modèle peut atteindre l'invariance structurelle, maintenant une haute qualité à travers toutes les échelles de décodage tout en récupérant la relation monotone attendue entre la taille du bloc et la perplexité.
• Limites : Les auteurs notent modestement que la performance dépend toujours du choix de $\pi$ (qui dicte le compromis entre les échelles) et que l'efficacité de l'inférence (sélection de la politique optimale au moment du test) reste un problème ouvert. Les garanties théoriques reposent sur la couverture du support et ne fournissent pas de bornes à échantillon fini.

En conclusion, l'Adaptive Block Diffusion offre un cadre robuste pour entraîner des modèles de langage par diffusion qui sont invariants aux configurations de décodage, répondant à une limitation critique des approches actuelles à structure fixe.