Résumé Technique : Mécanisme d'Attention Bayésienne (BAM)

Énoncé du Problème

Les modèles de langage basés sur les Transformers reposent sur le Codage Positionnel (PE) pour gérer l'ordre des jetons, pourtant les méthodes existantes font face à deux limitations critiques concernant l'extrapolation de la longueur du contexte (la capacité à généraliser à des séquences plus longues que celles vues pendant l'entraînement) :

1. Opacité Théorique : De nombreuses techniques de PE (par exemple, Sinusoidal, RoPE, ALiBi, NoPE) sont motivées empiriquement avec des fondements théoriques limités, rendant leur comportement lors de l'extrapolation mal compris.
2. Évaluation Inadéquate : La dépendance à la perplexité comme métrique principale est insuffisante. Les modèles peuvent atteindre une faible perplexité grâce à des motifs d'attention localisés (se concentrant uniquement sur les jetons récents) sans intégrer véritablement l'information à long terme, un défaut mis en évidence par les évaluations à fenêtre glissante.

Méthodologie : Le Mécanisme d'Attention Bayésienne (BAM)

Les auteurs proposent le Mécanisme d'Attention Bayésienne (BAM), un cadre théorique qui reformule l'auto-attention comme une espérance sur des valeurs calculées sous un modèle probabiliste conjoint du contenu des jetons et de leur position.

Cadre Théorique

• Formulation de Probabilité Conjointe : Dans le BAM, le poids d'attention $p_{ij}$ (la probabilité que le jeton $j$ s'attende au jeton $i$) est défini comme une distribution de probabilité conjointe :
  $$p_{ij} = p(f_{cont}(q_i, k_j) | g_{pos}(i, j)) \times p(g_{pos}(i, j))$$
  Ici, $f_{cont}$ représente la similarité de contenu, et $g_{pos}$ représente la relation positionnelle.
• Codage Positionnel comme Priors : Le cadre postule que les stratégies de codage positionnel correspondent à des distributions a priori spécifiques $p(g_{pos}(i, j))$ sur les positions des jetons.
  • NoPE est dérivé comme un prior Uniforme (restreint par le masque causal).
  • ALiBi est dérivé comme le produit d'un prior Uniforme et d'un prior Laplacien (décroissance exponentielle basée sur la distance).
• Le Scalaire de Normalisation ($Z$) : La dérivation introduit un scalaire $Z$ qui normalise le produit des probabilités de contenu et de position, représentant l'énergie libre d'interaction ou le couplage d'information entre le contenu et la position.

Solution Proposée : GGD-BAM

En tirant parti de ce cadre, les auteurs introduisent GGD-BAM, qui utilise une Distribution Gaussienne Généralisée (GGD) comme prior positionnel. La matrice de biais $B$ ajoutée aux scores d'attention est définie par :
$$b_{ij} = -\left| \frac{j - i - \mu}{\alpha} \right|^\beta$$
Où $\mu$, $\alpha$, et $\beta$ sont des paramètres apprenables (ou fixes pour des comportements spécifiques).

• $\beta \in (0, 1)$ : Produit une distribution à queue lourde qui décroît plus lentement que le prior Laplacien (ALiBi), permettant à l'attention de s'étendre plus loin qu'ALiBi tout en maintenant la localité.
• $\beta < 0$ : Une relaxation novatrice de la définition standard de la GGD. Cela crée un prior « inversé » où la masse de probabilité est supprimée près du jeton de requête et concentrée à l'extrémité lointaine de la fenêtre de contexte. Théoriquement, cela permet au modèle d'ignorer le contexte local et de se concentrer sur des dépendances arbitrairement longues.
• Softmax Évolutif (SSMax) : Les auteurs notent que le BAM est sujet à un « effacement de l'attention » (distribution uniforme à mesure que le contexte grandit). Ils combinent le BAM avec le SSMax, qui redimensionne dynamiquement les logits, pour préserver la netteté de la distribution.

Contributions Clés

1. Cadre Théorique Unifié : Le BAM fournit une unification probabiliste des méthodes de PE existantes, cartographiant explicitement NoPE et ALiBi vers des priors Uniforme et Laplacien, respectivement.
2. Prior Gaussien Généralisé : La proposition de GGD-BAM introduit une famille flexible de priors positionnels. En apprenant le paramètre de forme $\beta$, le modèle peut passer de manière adaptative entre les modes d'attention locale et de récupération à long terme.
3. Surcharge Minimale : La méthode introduit moins de 1 000 paramètres supplémentaires entraînables (spécifiquement 384 pour un modèle de 120M de paramètres) en apprenant $\theta_\alpha$ et $\theta_\beta$ pour chaque tête d'attention, tout en fixant $\theta_\mu = 0$.
4. Interprétabilité : Le cadre offre une manière fondée de visualiser et d'interpréter les motifs d'attention comme des distributions de probabilité explicites sur les positions.

Résultats Empiriques

Des expériences ont été menées sur des modèles allant de 120M à 1,1 milliard de paramètres, entraînés sur FineWeb 10B avec une longueur de contexte de 512 jetons.

• Récupération de Clé (Passkey Retrieval) : GGD-BAM (spécifiquement avec SSMax) a atteint une précision parfaite dans la récupération d'un jeton spécifique à partir de séquences allant jusqu'à 32 000 jetons (64x la longueur d'entraînement) et a maintenu une précision >80 % jusqu'à 512 000 jetons (1000x la longueur d'entraînement). En revanche, les bases de référence comme RoPE, Sinusoidal et NoPE se sont dégradées jusqu'à une précision proche du hasard au-delà de leur horizon d'entraînement. ALiBi a montré une meilleure extrapolation de perplexité mais a échoué de manière significative dans les tâches de récupération à des longueurs extrêmes.
• Benchmark RULER (NIAH) : GGD-BAM a surpassé toutes les bases de référence (y compris RoPE et ALiBi) dans le sous-ensemble « Needle In A Haystack », étant la seule méthode à maintenir une haute précision au-delà de 6 000 jetons dans la tâche Single 1.
• Perplexité : GGD-BAM a maintenu une perplexité comparable à ALiBi et RoPE sur des contextes longs, démontrant que l'amélioration de la récupération ne se fait pas au détriment de la fluidité du modèle de langage.
• Tâches en Aval : Sur les benchmarks MMLU et ARC, BAM SSMax a surpassé RoPE SSMax.
• Visualisation de l'Attention : Les visualisations ont confirmé que les têtes avec $\beta \le 0$ supprimaient efficacement le contexte local et se concentraient sur les jetons distants (la clé de récupération), tandis que les têtes avec $\beta \in (0, 1)$ maintenaient une attention à longue traîne.

Importance et Revendications

L'article revendique que le BAM sert un double objectif :

1. Théorique : Il fournit une base probabiliste rigoureuse pour le codage positionnel, clarifiant pourquoi certaines méthodes (comme ALiBi) fonctionnent pour la perplexité mais échouent pour la récupération (en raison de leur nature locale).
2. Pratique : Il offre une méthode hautement efficace pour l'extrapolation de la longueur du contexte. En permettant aux têtes d'attention d'apprendre des valeurs de $\beta$ négatives, le modèle peut explicitement supprimer la localité pour récupérer des informations de contextes arbitrairement longs sans nécessiter une exposition directe à de telles longueurs pendant l'entraînement.

Les auteurs soulignent que si la perplexité est une métrique standard, elle est insuffisante pour évaluer les véritables capacités de contexte long ; des tâches basées sur la récupération sont nécessaires pour vérifier que les modèles s'attendent à l'information distante plutôt que de simplement compter sur des motifs locaux. Le travail suggère que la capacité à modéliser explicitement les priors positionnels est un biais inductif crucial pour les Transformers de contexte long de nouvelle génération.