💬 NLP

Prefilling-dLLM: Predictive Prefilling for Long-Context Inference in Diffusion Language Models

Le document présente Prefilling-dLLM, un cadre sans entraînement qui accélère l'inférence à contexte long dans les modèles de langage de diffusion en mettant en cache et en sélectionnant de manière éparse des blocs de préfixes pertinents afin de réduire la complexité computationnelle d'une croissance quadratique par rapport à la longueur totale de la séquence à une croissance quadratique par rapport à la longueur de décodage, atteignant ainsi des accélérations de pointe et atténuant le phénomène de perte au milieu (lost-in-the-middle).

Auteurs originaux : Jing Xiong, Qi Han, Shansan Gong, Yunta Hsieh, Chengyue Wu, Chaofan Tao, Chenyang Zhao, Ngai Wong

Publié 2026-06-10

📖 5 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Jing Xiong, Qi Han, Shansan Gong, Yunta Hsieh, Chengyue Wu, Chaofan Tao, Chenyang Zhao, Ngai Wong

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous essayez de résoudre un immense puzzle, mais qu'au lieu de regarder l'image entière d'un coup, vous devez réexaminer l'intégralité de la boîte de pièces à chaque fois que vous posez une seule nouvelle pièce sur la table.

C'est essentiellement ainsi que fonctionnent actuellement les Modèles de Langage de Diffusion (dLLM) lorsqu'ils traitent des histoires ou des documents longs. Chaque fois que le modèle essaie de générer un nouveau mot, il relit et traite à nouveau tout l'historique de la conversation depuis le tout début. À mesure que l'histoire s'allonge, ce « re-lecture » devient si lente et coûteuse que c'est comme essayer de courir un marathon en portant un sac à dos qui devient plus lourd à chaque pas.

L'article introduit une nouvelle méthode appelée Prefilling-dLLM pour corrir cela. Voici comment elle fonctionne, en utilisant quelques analogies de la vie quotidienne :

1. Le Problème : Le piège de la « Re-lecture »

Dans les modèles d'IA traditionnels, une fois que vous avez lu une page d'un livre, vous vous en souvenez et vous ne vous concentrez que sur la nouvelle page que vous lisez. Mais dans ces modèles de diffusion, l'ordinateur oublie le livre et relit l'intégralité du livre, de la page 1 à la page 1 000, chaque fois qu'il écrit un seul nouveau mot.

Le Résultat : Si le livre est court, tout va bien. Si le livre fait 32 000 pages, l'ordinateur passe 99 % de son temps à simplement relire les anciennes pages et seulement 1 % à écrire de nouvelles pages.

2. La Solution : Le système du « Bibliothécaire Intelligent »

Les auteurs proposent un système appelé Prefilling-dLLM qui agit comme un bibliothécaire super efficace. Au lieu de relire toute la bibliothèque à chaque fois, le bibliothécaire fait deux choses :

Étape A : Le « Prefill » (Organiser la bibliothèque)

Avant que l'écriture ne commence, le bibliothérapeute prend le texte d'entrée massif (le « préfixe ») et le découpe en morceaux (comme des chapitres dans un livre).

L'astuce : Le bibliothécaire lit chaque chapitre une seule fois, crée une « fiche de résumé » (appelée cache KV) pour celui-ci, et la pose sur une étagère.
L'innovation : Le bibliothécaire ne lit pas chaque mot de chaque chapitre. Ils utilisent une astuce spéciale pour ne conserver que les phrases les plus importantes de chaque chapitre, en jetant le superflu. Cela rend les fiches de résumé beaucoup plus petites et rapides à manipuler.

Étape B : Le « Decode » (Écrire l'histoire)

Maintenant, quand le modèle doit écrire le mot suivant, il ne relit pas toute la bibliothèque.

La Sélection : Le modèle demande : « Quels chapitres sont réellement pertinents par rapport à ce que j'écris en ce moment ? » Il utilise un système de notation intelligent pour choisir seulement les quelques chapitres supérieurs (les morceaux les plus pertinents) sur l'étagère.
L'Efficacité : Il ignore les 90 % de la bibliothèque qui ne sont pas nécessaires pour cette phrase spécifique. Il ne regarde que les « fiches de résumé » pertinentes qu'il a créées précédemment.

3. Pourquoi c'est un changement de donne (Game-Changer)

L'article affirme que cette approche est une accélération massive car :

Plus de Re-lecture : Le gros travail de lecture du texte long se fait une seule fois au début.
Saut Intelligent : Pendant la phase d'écriture réelle, le modèle ne regarde qu'une infime fraction du texte (les morceaux les plus pertinents), plutôt que l'ensemble.
La Vitesse : Sur des contextes longs (8 000 à 32 000 mots), cette méthode est 9 à 28 fois plus rapide que les méthodes précédentes, tout en donnant les bonnes réponses.

4. Résoudre le mystère du « Perdu au Milieu » (Lost in the Middle)

Il existe un problème connu avec les modèles d'IA longs appelé « Lost in the Middle ». Imaginez une personne lisant une longue liste de faits ; elle se souvient parfaitement des premiers et des derniers, mais elle oublie complètement ce qui se trouve au milieu.

La Correction de l'Article : Les auteurs ont découvert qu'en découpant le texte en morceaux et en ajoutant un « jeton d'ancrage » spécial (comme un titre de chapitre) au début de chaque morceau, le modèle peut trouver des informations n'importe où dans le texte, même au milieu. C'est comme avoir une table des matières qui fonctionne parfaitement, de sorte que le modèle ne se perd jamais, peu importe la longueur de l'histoire.

5. L'équilibre de la « Taille des Morceaux »

L'article a également testé la taille de ces « chapitres ».

Trop Grands : Si les morceaux sont énormes, le modèle pourrait manquer des détails importants au milieu du texte.
Trop Petits : Si les morceaux sont minuscules, l'ordinateur passe trop de temps à gérer la liste des morceaux.
Le Juste Milieu : Ils ont constaté que pour des textes très longs, l'utilisation de nombreux morceaux plus petits fonctionne mieux pour maintenir une précision élevée.

Résumé

Prefilling-dLLM est comme passer d'un étudiant qui relit un manuel de 500 pages à chaque fois qu'il répond à une question de quiz, à un étudiant qui crée un guide d'étude condensé et intelligent une seule fois, puis qui ne feuillette que les pages spécifiques dont il a besoin pour répondre.

Le résultat ? L'IA peut gérer des conversations et des documents beaucoup plus longs sans ralentir jusqu'à l'arrêt total, et elle n'oublie pas les détails importants cachés au milieu du texte.

Résumé Technique : Prefilling-dLLM

Énoncé du Problème

Les modèles de langage à diffusion (dLLM) offrent une alternative prometteuse aux modèles autorégressifs (AR) en permettant une génération de jetons en parallèle via un débruitage itératif. Cependant, les dLLM sont confrontés à une inefficacité critique dans les scénarios à contexte long : à chaque étape de débruitage, le modèle doit ré-encoder l'intégralité du préfixe d'entrée pour calculer l'auto-attention. Cela entraîne une complexité de calcul qui croît de manière quadratique avec la longueur totale de la séquence ( $O((L_p + L_d)^2 \cdot T)$ , où $L_p$ est la longueur du préfixe, $L_d$ la longueur du décodage et $T$ les étapes de débruitage). Contrairement aux modèles AR, où la phase de prefill est limitée par le calcul (compute-bound) et la phase de décodage par la mémoire (memory-bound), l'inférence des dLLM reste limitée par le calcul durant toutes les itérations car la séquence complète (préfixe + réponse) est traitée conjointement à chaque étape. Cette recomputation répétée du préfixe devient prohibitive pour les applications à contexte long (ex. : 8K–32K+ jetons).

Méthodologie : Prefilling-dLLM

Les auteurs proposent PREFILLING-DLLM, un cadre de désagrégation prefill-decode sans entraînement, conçu pour éliminer la recomputation du préfixe et introduire de la parcimonie. La méthode opère en trois étapes :

1. Prefill par blocs et mise en cache KV

Au lieu de traiter l'intégralité du préfixe comme un bloc unique, le préfixe d'entrée est partitionné en $N$ blocs non chevauchants de taille fixe $C$ .

Traitement indépendant : Chaque bloc est traité indépendamment avec une attention bidirectionnelle intra-bloc. Un jeton spécial de début de séquence (BOS) est prépendé à chaque bloc pour servir de délimiteur.
Scoring prédictif : Pour déterminer quels blocs sont pertinents sans cibles de vérité terrain, le cadre génère des pseudo-étiquettes ( $m$ $m$ ) via un passage de diffusion préliminaire. Il score ensuite chaque bloc selon deux signaux :
- Score d'auto-information ( $s_I$ ) : La log-vraisemblance négative de la fenêtre de requête conditionnée par le bloc.
- Score de pseudo-étiquette ( $s_P$ ) : La vraisemblance des pseudo-étiquettes conditionnée par le bloc et la requête.
Sélection : Les $K$ blocs les plus pertinents sont sélectionnés sur la base du score combiné ( $s = s_I + s_P$ ). Les représentations KV de ces blocs sélectionnés sont mises en cache.

2. Parcimonie à deux niveaux

Le cadre introduit de la parcimonie tant au niveau des blocs qu'au niveau des jetons pour réduire la complexité :

Parcimonie inter-blocs : Seuls les $K$ meilleurs blocs participent au mécanisme d'attention pendant le décodage, plutôt que l'ensemble des $N$ blocs.
Parcimonie intra-bloc : Au sein de chaque bloc sélectionné, les jetons sont classés selon leur poids d'attention bidirectionnelle cumulée avec les jetons de requête. Seuls les $B$ jetons les plus pertinents par bloc sont conservés dans le cache KV, évincant les jetons non pertinents.

3. Décodage avec réutilisation statique du KV

Pendant le processus de débruitage itératif :

Les représentations KV mises en cache des $K$ blocs sélectionnés restent statiques et sont réutilisées à travers les $T$ étapes de débruitage.
Le modèle ne recompute que les représentations KV pour la requête et les jetons de réponse en évolution.
Cela réduit la complexité par étape de $O((L_p + L_d)^2 \cdot T)$ à $O(N \cdot C^2 + (L_d^2 + K \cdot C) \cdot T)$ , isolant ainsi efficacement le coût de décodage de la longueur totale du préfixe.

Optimisation du Noyau d'Attention

Pour traiter les goulots d'étranglement de la mémoire causés par le chargement de caches KV de blocs non contigus, les auteurs implémentent un noyau Split-S FlexAttention. Ce noyau parallélise le décodage sur les KVs de blocs stockés de manière non contiguë en partitionnant la dimension KV et en traitant les blocs indépendamment avant de fusionner les résultats via une réduction log-sum-exp, évitant ainsi la nécessité de regrouper les données dans une mémoire contiguë.

Principaux Résultats et Constats

Performance et Efficacité

Accélération : Sur LongBench et InfiniteBench, PREFILLING-DLLM atteint une accélération de 9,1× à 28,0× pour des longueurs de contexte de 8K à 32K par rapport aux modèles de référence Transformers standards.
Qualité : La méthode atteint une qualité de pointe parmi les méthodes d'accélération de dLLM. Sur LongBench, elle obtient un score moyen de 34,59 (Dream-7B) et 37,02 (UltraLLaMA), surpassant les bases de référence comme Fast-dLLM, Sparse-dLLM et dKV-Cache.
Mise à l'échelle du contexte long : Sur InfiniteBench (contextes >128K), la méthode atteint une précision moyenne de 43,62 % sur Dream-7B, dépassant nettement la base de référence la plus forte (Fast-dLLM v2 à 30,32 %). Notamment, elle atteint 95,42 % de précision sur la tâche de récupération de Passkey et 70,00 % sur la récupération de nombres.

Traitement du phénomène "Lost-in-the-Middle"

Phénomène : L'article observe que les dLLM présentent un phénomène de "perdu au milieu" (lost-in-the-middle) lors de l'extrapolation de contexte, où la précision diminue pour les informations situées au milieu du contexte, montrant un biais de décroissance monotone vers la réponse plutôt que la courbe en U observée dans les modèles AR.
Atténuation : En sélectionnant les blocs sur la base de la pertinence prédictive plutôt que de la proximité positionnelle, PREFILLING-DLLM élimine ce biais. La méthode maintient une haute précision de correspondance exacte à toutes les positions du contexte, même à des longueurs de 32K et 128K.

Insights sur les mécanismes

Ancres d'attention : Les jetons BOS prépendés à chaque bloc agissent comme des ancres d'attention périodiques. L'analyse montre que ces jetons distribuent la masse d'attention uniformément à travers le contexte (formant des crêtes stables de faible magnitude) plutôt que de dégénérer en "puits d'attention" (attention sinks) qui absorbent une masse d'attention disproportionnée. Ce mécanisme permet une récupération invariante à la position.
Compromis de taille de bloc : Il existe un compromis entre la taille du bloc et la précision de la récupération. Les petits blocs (ex. : 256 jetons) maintiennent une haute précision même à 128K contextes, tandis que les blocs plus larges (ex. : 4096 jetons) se dégradent brutalement au-delà de 32K. La sélection multi-blocs est essentielle pour l'inférence à contexte long.

Signification et Revendications

L'article affirme que PREFILLING-dLLM comble l'écart entre la nature limitée par le calcul de l'inférence des dLLM et l'efficacité des architectures de désagrégation prefill-decode utilisées dans les modèles AR. Ses principales contributions sont :

Désagrégation sans entraînement : Il introduit un cadre qui met en cache les représentations KV du préfixe une seule fois et les réutilise, éliminant le coût de recomputation quadratique du préfixe à chaque étape de débruitage sans nécessifier de réentraînement du modèle.
Prefill parcimonieux : Il démontre qu'un prefill parcimonieux (sélection de blocs et de jetons pertinents) peut surpasser l'attention dense dans les dLLM, réduisant la complexité tout en améliorant ou en maintenant la qualité de génération.
Optimisation au niveau système : Il fournit un noyau d'attention spécialisé (Split-S FlexAttention) qui exploite la structure par blocs du préfixe pour obtenir des réductions de latence significatives (jusqu'à 10,2× par rapport au FlexAttention classique) durant la phase de décodage.
Robustesse aux contextes longs : Il montre que cette approche atténue le phénomène de "lost-in-the-middle" et permet aux dLLM de passer à l'échelle vers 128K+ jetons avec une haute précision de récupération, une capacité auparavant inexplorée dans l'accélération des dLLM.

Les auteurs notent des limites, notamment la nature statique de la sélection des blocs (pas de resélection dynamique pendant le décodage), la nécessité d'un réglage spécifique à la tâche pour la taille des blocs et $K$ , et le fait que l'évaluation actuelle soit limitée aux benchmarks en anglais sur des architectures de modèles spécifiques (Dream-7B et UltraLLaMA).