Résumé technique : Dérive de l'attention dans le décodage spéculatif autorégressif

Énoncé du problème

Le décodage spéculatif accélère l'inférence des modèles de langage de grande taille (LLM) en utilisant un modèle « rédacteur » léger pour prédire les futurs jetons, lesquels sont ensuite vérifiés par un modèle « cible » plus grand. Bien qu'efficace, les modèles rédacteurs actuels (spécifiquement EAGLE-3 et les têtes de prédiction multi-jetons) présentent une fragilité significative face aux variations de déploiement, telles que les perturbations de modèles et les entrées à contexte long. Les praticiens doivent souvent réentraîner les rédacteurs pour des moteurs d'inférence ou des invites spécifiques afin de maintenir les performances, ce qui suggère une instabilité architecturale sous-jacente.

L'article identifie un phénomène précédemment non signalé appelé dérive de l'attention. Alors qu'un rédacteur génère une séquence de jetons spéculatifs, son mécanisme d'attention se déplace progressivement de l'« puits d'attention » (une ancre stable, généralement située près du début de l'invite) pour se concentrer de plus en plus sur ses propres jetons récemment générés. Cette dérive est corrélée à une dégradation des taux d'acceptation, en particulier dans les scénarios à contexte long et lorsque les modèles d'entrée s'écartent des distributions d'entraînement.

Méthodologie et analyse

Les auteurs enquêtent sur les causes profondes de la dérive de l'attention grâce à une combinaison d'observations empiriques et d'analyses architecturales sur plusieurs familles de modèles (Llama 3.1, Qwen 3/3.5 et GPT-oss).

1. Observation phénoménologique

En utilisant des cartes thermiques d'attention et des métriques quantitatives, les auteurs démontrent que :

• Migration du puits : Dans les modèles dotés de puits d'attention forts (par exemple, Llama, Qwen), l'attention du rédacteur envers le jeton puits diminue à mesure que la chaîne de spéculation s'approfondit, tandis que l'attention envers le jeton rédigé le plus récent augmente.
• Universalité : Cette dérive est observée aussi bien dans les rédacteurs EAGLE-3 (entraînés a posteriori sur des cibles figées) que dans les têtes MTP (entraînées conjointement avec la cible), suggérant qu'il s'agit d'une propriété des conceptions autorégressives qui consomment leurs propres sorties.
• Absence de puits : Même dans les modèles conçus pour supprimer les puits (par exemple, GPT-oss avec attention à porte ou biais apprenables), une forme de dérive se produit où l'attention se concentre sur les jetons récurrents du modèle plutôt que sur le début de l'invite.

2. Analyse des causes profondes : Magnitude de l'état caché

L'article rattache la dérive de l'attention au chemin résiduel non normalisé entre les étapes de spéculation dans les architectures pré-norm standard :

• Croissance monotone : Dans les transformateurs pré-norm, les connexions résiduelles accumulent de la magnitude sans normalisation. Les auteurs montrent que la magnitude de l'état caché du rédacteur ($\|h^*_k\|$) croît de manière monotone avec la profondeur de spéculation ($k$).
• Empilement de couches vs autorégression : Cette croissance amène le rédacteur à se comporter moins comme un prédicteur autorégressif autonome et davantage comme une séquence de couches de transformateurs supplémentaires empilées au-dessus de la cible. À l'étape $k$, le rédacteur approxime efficacement le modèle cible à la profondeur $N+k$, plutôt que de prédire le jeton suivant basé sur une distribution stable.
• Entrées déséquilibrées : Les états cachés fusionnés provenant du modèle cible ($h_{FC}$) sont souvent dominés par des couches plus profondes ayant des magnitudes plus grandes, exacerbant davantage le déséquilibre d'échelle lorsqu'ils sont injectés dans le rédacteur.

3. Interventions proposées

Pour contrer l'accumulation de magnitude et stabiliser le processus de rédaction, les auteurs proposent deux modifications architecturales :

1. Post-norm sur les états cachés du rédacteur : L'application de RMSNorm aux états cachés du rédacteur après la connexion résiduelle (entre les étapes de spéculation) empêche la croissance monotone de la magnitude. Cela force le rédacteur à apprendre une règle autorégressive invariante en profondeur plutôt que des raffinements spécifiques à la profondeur.
2. RMSNorm par état caché : L'ajout de couches RMSNorm avant la fusion des états cachés de la cible ($h_{low}, h_{mid}, h_{high}$) dans l'entrée du rédacteur pour équilibrer la magnitude du signal entrant.

Les auteurs évaluent également l'attention à porte (suppression de la dépendance à un puits fixe) comme une intervention distincte. Ils constatent que si la porte élimine le puits, elle n'empêche pas la croissance de la magnitude, et que la combinaison de la porte avec la post-norm peut entraîner une sur-régularisation (effondrement de l'entropie).

Résultats clés

L'architecture post-norm proposée a été évaluée par rapport à l'EAGLE-3 pré-norm actuel, état de l'art, sur diverses références et conditions :

• Perturbation de modèle : Sous des perturbations de modèle (par exemple, suppression des modèles de chat ou des jetons de début de séquence), le rédacteur post-norm a amélioré la longueur d'acceptation jusqu'à 2× par rapport au pré-norm. Les modèles pré-norm ont subi des chutes catastrophiques (jusqu'à 52 % de perte), tandis que les modèles post-norm sont restés résilients (perte maximale de 5 %).
• Tâches à contexte long : Sur des tâches à contexte long (LongBench, chat multi-tours avec attention à fenêtre glissante), la post-norm a amélioré la longueur d'acceptation de 1,18×. Dans des paramètres de contexte long avec attention complète, les modèles pré-norm se sont effondrés vers une acceptation proche de zéro, tandis que la post-norm a maintenu des performances utilisables.
• Références standard : Sur sept références standard (chat multi-tours, mathématiques, codage), la post-norm a montré une amélioration constante de 1,10× de la longueur d'acceptation.
• Efficacité de l'entraînement : Les rédacteurs post-norm généralisent aux profondeurs d'inférence au-delà de leur horizon d'entraînement. Cela permet un entraînement avec des profondeurs « entraînement-test » (TTT) plus courtes (par exemple, réduction de 8 à 4 étapes), réduisant le temps d'entraînement d'environ un tiers sans sacrifier les performances.
• Robustesse au bruit : Les architectures post-norm tolèrent significativement plus de bruit dans les états cachés (58 % de rétention d'acceptation à un bruit $\alpha=0,5$ contre 5 % pour le pré-norm), suggérant une meilleure robustesse face à la quantification et aux décalages de distribution.

Importance et affirmations

L'article affirme que la dérive de l'attention est un mode de défaillance critique, précédemment non signalé, dans le décodage spéculatif, qui découle de l'interaction entre les chemins résiduels non normalisés et la rétroaction autorégressive.

• Diagnostic : Les auteurs soutiennent que la dérive de l'attention est un symptôme d'un problème plus profond : l'accumulation de la magnitude de l'état caché, qui amène le rédacteur à apprendre des transformations dépendantes de la profondeur plutôt qu'une prédiction de jeton stable.
• Solution : Ils affirment que simplement corriger le puits d'attention (par exemple, via une porte) est insuffisant ; la dynamique sous-jacente de la magnitude doit être traitée. L'architecture post-norm proposée est présentée comme une correction simple et efficace qui stabilise le processus de rédaction.
• Généralisation : Les résultats suggèrent que la post-norm devrait être considérée comme un choix architectural par défaut pour les rédacteurs de décodage spéculatif, en particulier pour les systèmes de production nécessitant une robustesse face aux invites variables et aux contextes longs.

L'article conclut que nommer et caractériser la « dérive de l'attention » fournit un diagnostic utile pour la robustesse des rédacteurs, et que les interventions architecturales ciblant l'accumulation de magnitude sont essentielles pour un décodage spéculatif fiable.