Résumé Technique : FlashBlock

1. Énoncé du Problème

Bien que la diffusion par blocs ait amélioré l'efficacité de l'inférence dans les modèles génératifs à contexte long (par exemple, les modèles de diffusion de langage et de génération de vidéo) en combinant le cache KV avec l'inférence causale par blocs, elle souffre toujours d'un surcoût important dans les configurations à contexte long. Dans la diffusion par blocs standard, chaque étape de diffusion nécessite de recomput l'attention sur l'ensemble du cache KV en pleine croissance. À mesure que la longueur du contexte augmente, le coût combiné du calcul de l'attention et de l'accès au cache KV devient le facteur dominant de la latence d'inférence.

L'article identifie une propriété sous-explorée de la diffusion par blocs : la redondance inter-étapes. Plus précisément, durant le processus de débruitage, les jetons situés en dehors du bloc actuel en cours de mise à jour restent largement statiques, ce qui entraîne des contributions d'attention très stables de la part de ces jetons « externes au bloc » à travers les étapes de diffusion adjacentes. À l'inverse, l'attention au sein du bloc actuel (« interne au bloc ») varie considérablement à mesure que les jetons sont activement affinés ou démasqués. Les méthodes actuelles ne parviennent pas à exploiter cette stabilité, recalculant de manière répétée l'attention redondante externe au bloc.

2. Méthodologie : FlashBlock

Les auteurs proposent FlashBlock, un mécanisme qui met en cache et réutilise les sorties d'attention externes au bloc pour réduire le calcul et les entrées/sorties (I/O) mémoire sans modifier le processus de diffusion sous-jacent.

Mécanisme Central

FlashBlock décompose la sortie d'attention complète en deux composantes :

1. Attention Interne au Bloc ($A_{in}$) : Calculée sur les jetons à l'intérieur du bloc actuel ($J_{in}$), qui sont en cours de changement actif.
2. Attention Externe au Bloc ($A_{out}$) : Calculée sur les jetons en dehors du bloc actuel ($J_{out}$), qui restent stables à travers les étapes.

La méthode met en cache la sortie d'attention externe au bloc ($A_{out}$) et son normalisateur logarithmique ($L_{out} = \log Z_{out}$) de l'étape de diffusion précédente. Dans les étapes suivantes, au lieu de réaccéder à l'intégralité du cache KV pour $J_{out}$, le modèle réutilise les valeurs mises en cache. La sortie d'attention complète est reconstruite en combinant l' $A_{out}$ mise en cache avec la nouvelle $A_{in}$ calculée en utilisant une composition dans l'espace logarithmique (similaire à FlashAttention) pour garantir la stabilité numérique :
$$A_{full} = \frac{e^{L_{out}-m} A_{out} + e^{L_{in}-m} A_{in}}{e^{L_{out}-m} + e^{L_{in}-m}}$$
où $m = \max(L_{out}, L_{in})$.

Stratégies de Réutilisation Adaptatives

• Réutilisation basée sur un seuil : Dans les modèles de diffusion discrets, les mises à jour de jetons peuvent être abruptes. La méthode introduit un seuil $\tau$ sur le nombre de jetons mis à jour ($M^{s+1}$). Si $M^{s+1} < \tau$, l'attention mise en cache est réutilisée ; sinon, elle est recalculée pour garantir l'exactitude.
• Sélection par tête (Vidéo) : Pour la diffusion vidéo, où la variabilité inter-étapes est plus élevée, la méthode estime la similitude par tête d'attention. Seules les têtes dont la similitude dépasse un seuil $\gamma$ réutilisent le cache.
• Compatibilité avec l'attention parcimonieuse (Sparse Attention) : FlashBlock est orthogonal à l'attention parcimonieuse. Il traite les clés non sélectionnées dans l'attention parcimonieuse comme l'ensemble « externe au bloc », mettant en cache leurs contributions résiduelles. Cela permet à FlashBlock de compenser la perte d'information causée par une parcimonisation agressive.

Adaptation de l'Entraînement (Optionnel)

Pour traiter les potentiels décalages de distribution dans les modèles de langage à diffusion discrète (dLLM) causés par le démasquage abrupt de jetons, les auteurs proposent une stratégie de distillation sensible à la réutilisation. Un modèle étudiant employant FlashBlock est entraîné pour correspondre à la distribution de sortie d'un modèle enseignant gelé utilisant une attention dense, en utilisant des adaptateurs LoRA pour garder l'adaptation légère.

3. Contributions Clés

1. Insight Empirique : L'article fournit une étude empirique révélant une forte stabilité inter-étapes dans l'attention externe au bloc, contrastant avec la haute variabilité de l'attention interne au bloc.
2. Mécanisme FlashBlock : Un nouveau mécanisme de mise en cache qui exploite cette redondance pour réduire le calcul de l'attention et l'accès au cache KV, implémenté via une agrégation dans l'espace logarithmique.
3. Synergie avec la Parcimonie : Démonstration que FlashBlock sert de stratégie complémentaire de réutilisation résiduelle pour l'attention parcimonieuse, atténuant la dégradation de la qualité due à une parcimonisation agressive.
4. Gains d'Efficacité : Accélération significative du débit d'inférence et réduction de la latence pour les modèles de langage de diffusion et les modèles de génération de vidéo, avec un impact négligeable sur la qualité de la génération.

4. Résultats Expérimentaux

Les expériences ont été menées sur des modèles de langage de diffusion (Trado-8B/4B) et des modèles de diffusion vidéo (LongLive-1.3B).

• Débit (Throughput) : FlashBlock a atteint un débit de jetons jusqu'à 1,44× plus élevé sur les modèles de langage de diffusion (par exemple, passant de 312 à 451 jetons/s pour une taille de bloc de 4).
• Latence : Le temps d'attention a été réduit jusqu'à 1,6×.
• Qualité : La qualité de génération est restée comparable aux bases de référence denses. Sur les tests de raisonnement mathématique (GSM8K, MATH500) et de codage (HumanEval, MBPP), les différences de précision étaient minimales (entre 0,2 % et 1,6 % selon la taille du bloc).
• Mise à l'échelle du Contexte Long : À mesure que la longueur du contexte passait de 100k à 800k, l'écart de latence entre FlashBlock et la base de référence s'est creusé, FlashBlock montrant une croissance de latence d'environ la moitié de celle de la base de référence sous certains seuils.
• Combinaison avec l'Attention Parcimonieuse : Lorsqu'il est combiné avec SparseD (à 20 %, 30 % et 40 % de densité), FlashBlock a systématiquement amélioré la précision par rapport à SparseD seul (améliorant par exemple la précision GSM8K de 34,72 % à 42,68 % à 20 % de densité), récupérant efficacement l'information perdue par la parcimonisation.

5. Signification

L'article affirme que FlashBlock répond à un goulot d'étranglement critique de la diffusion par blocs à contexte long : le recalcul redondant des contributions d'attention stables. En découplant le contexte « externe au bloc » stable des mises à jour « internes au bloc » dynamiques, la méthode réduit considérablement les coûts de calcul et d'I/O mémoire associés aux caches KV en croissance.

Les auteurs soulignent que cette approche est orthogonale aux techniques existantes d'attention parcimonieuse, offrant une stratégie complémentaire qui améliore à la fois l'efficacité et la précision. Les résultats suggèrent que FlashBlock fournit une solution évolutive pour une génération efficace à contexte long dans les domaines du texte et de la vidéo, permettant une inférence plus rapide sans compromettre la qualité du contenu généré.