Attention Sinks and Outliers in Attention Residuals

Résumé technique : OASIS – Routage conscient des valeurs aberrantes et des puits pour les architectures AttnResidual

Énoncé du problème
L'article aborde des problèmes critiques de stabilité et de robustesse survenant dans les architectures AttnResidual, une conception récente qui remplace la normalisation Softmax standard à un seul niveau par un schéma de double normalisation régissant à la fois le routage de l'attention au niveau des tokens et au niveau de la profondeur. Bien qu'AttnResidual améliore l'expressivité représentationnelle, les auteurs identifient un mode de défaillance structurel : la conception à double normalisation exacerbe les puits d'attention (attention sinks) et les valeurs aberrantes d'activation.

Dans les Transformers standards, les puits d'attention (par exemple, le token <|begin_of_text|>) accumulent une masse de probabilité pour stabiliser les représentations. Cependant, dans AttnResidual, la contrainte supplémentaire de simplexe sur le routage de la profondeur force les mises à jour « sans opération » (identité) à être exprimées via de vrais tokens et de vraies branches de profondeur plutôt que par un canal nul explicite. Cette contrainte nécessite des logits pré-Softmax extrêmes pour approximer des mises à jour d'identité, conduisant à :

1. Intensification des puits d'attention : Concentration accrue sur les tokens précoces.
2. Amplification des valeurs aberrantes d'activation : Les états cachés présentent des distributions à queues lourdes (kurtose élevée) et de grandes normes infinies.
3. Fragilité de la quantification : Ces valeurs aberrantes dégradent la stabilité de l'inférence et rendent la quantification à faible nombre de bits (par exemple, W4A4, W8A8) significativement moins robuste.
4. Effondrement de la profondeur : Les poids de routage s'effondrent vers une seule branche dominante, ramenant effectivement l'architecture à un comportement résiduel de type vanilla.

Méthodologie : OASIS
Les auteurs proposent OASIS (Routage conscient des valeurs aberrantes et des puits par signalisation nulle inter-couche), une intervention légère qui modifie la couche de normalisation sans altérer l'objectif d'entraînement ni le squelette architectural central. OASIS opère sur deux fronts :

1. Routage conscient du nul (Softmax1) :
   La méthode remplace les opérateurs Softmax standards aux niveaux des tokens et de la profondeur par Softmax1. Cette fonction introduit un état nul explicite (noté $\emptyset$) dans le dénominateur de la normalisation.

   • Niveau des tokens : Permet à la masse d'attention d'être routée vers un état « nul » plutôt que de la forcer sur de vrais tokens.
   • Niveau de la profondeur : Permet aux poids de routage de la profondeur d'absorber la pression « sans opération » via une couche nulle, empêchant la concentration de masse sur des branches de profondeur spécifiques.
   • Mécanisme : En découplant la masse non informative de la compétition entre les vrais tokens/branches, Softmax1 réduit le besoin d'amplification extrême des logits.

2. Couplage token-à-profondeur nul :
   Reconnaissant que le routage au niveau des tokens et au niveau de la profondeur reste découplé dans Softmax1 standard, OASIS introduit un mécanisme de couplage. Il agrège les preuves nulles au niveau des tokens (la masse de probabilité routée vers l'état nul à travers les têtes d'attention) en une statistique scalaire pour chaque branche source. Cette statistique est ensuite injectée dans les logits de routage de la profondeur.

   • Effet : Les branches présentant un comportement « sans opération » élevé (masse nulle élevée) sont pondérées à la baisse dans le routage de la profondeur, tandis que les branches informatives sont pondérées à la hausse. Cela empêche le routeur de profondeur d'attribuer un poids significatif à des couches non informatives.

Contributions clés

• Identification d'une défaillance structurelle : L'article identifie formellement que la conception à double normalisation d'AttnResidual amplifie structurellement les valeurs aberrantes et les puits car elle manque d'un chemin nul explicite pour les mises à jour proches de l'identité.
• Analyse théorique : Les auteurs fournissent des preuves théoriques (Lemmes 5.1, 5.2 ; Théorèmes 5.1, 5.2) démontrant que :
  • Sans canal nul, les contraintes proches du « sans opération » forcent la masse d'attention sur des tokens à faible contribution, créant des valeurs aberrantes.
  • La double normalisation conduit à un effondrement de la profondeur où le routage se concentre sur des branches à mise à jour minimale.
  • Softmax1 introduit une route nulle réalisable qui découple le comportement « sans opération » de l'allocation de vrais tokens, réduisant théoriquement la pression structurelle.
• Cadre OASIS : Une modification pratique et minimale (remplacement de Softmax par Softmax1 et ajout d'un terme de couplage) qui atténue ces pathologies.

Résultats expérimentaux
Les auteurs ont évalué OASIS sur les architectures de base Llama-3.2-1B et Qwen3-0.6B à travers trois jeux de données réels (BookCorpus, Wiki40B, WikiText-2, GSM8K).

• Suppression des valeurs aberrantes : Par rapport à la référence Vanilla AttnResidual, OASIS a réalisé une réduction moyenne de 84,14 % de la norme infinie maximale et de 96,77 % de la kurtose moyenne.
• Robustesse de la quantification :
  • Sous quantification W8A8, OASIS a réduit la perplexité de 73,55 % par rapport à Vanilla.
  • Sous quantification W4A4, OASIS a amélioré la précision Pass@1 sur GSM8K de 23,47 %.
• Atténuation des puits d'attention : L'analyse qualitative des cartes d'attention a montré que OASIS redistribue efficacement l'attention loin du puits dominant du premier token, produisant des motifs d'attention causale plus lisses et plus locaux par rapport à la fois à Vanilla et aux références intermédiaires (comme AoS ou OutEffHop).
• Comparaison : OASIS a constamment surpassé les références, notamment Clip Attention, Gated Attention, OutEffHop et les normalisateurs épars (Sparsemax, Entmax15), tant sur les métriques de valeurs aberrantes que sur les performances des tâches en aval sous une quantification agressive.

Portée et affirmations
L'article affirme que OASIS fournit une solution fondée sur des principes aux limitations de robustesse des architectures à double normalisation. En introduisant une capacité nulle explicite, la méthode stabilise le routage résiduel, permettant aux modèles de gérer une compression agressive (quantification à faible nombre de bits) et une inférence sur de longs contextes sans la dégradation généralement causée par les puits d'attention et les valeurs aberrantes d'activation. Les auteurs soulignent que ces améliorations sont obtenues sans modifier l'objectif d'entraînement, suggérant que le routage nul explicite est un mécanisme fondamental pour stabiliser les architectures résiduelles complexes.

Limites
Les auteurs notent que leur évaluation est actuellement limitée à des modèles de base spécifiques (1B et 0,6B paramètres) et qu'une validation plus large sur des modèles plus grands est requise. Ils mettent également en garde contre le fait qu'une allocation nulle excessive pourrait potentiellement supprimer des signaux faibles utiles ou altérer l'interprétabilité des cartes d'attention, suggérant un travail futur sur le contrôle adaptatif de la force nulle.