Attention Rollout

Introduction

Attention Rollout est une méthode d'interprétation des modèles basés sur transformeurs, principalement les grands modèles de langage, en calculant l'attention cumulative qui circule de chaque jeton d'entrée vers chaque autre jeton à travers toutes les couches du réseau. Introduite en réponse au constat que les poids d'attention dans les couches individuelles sont souvent bruités et ne reflètent pas directement la décision finale du modèle, Attention Rollout agrège les matrices d'attention à travers les couches pour produire une carte d'importance cohérente et unique. Cette technique est largement utilisée dans le domaine de l'IA explicable pour visualiser quelles parties d'une entrée contribuent le plus à la prédiction d'un modèle, facilitant ainsi le débogage, la détection de biais et la compréhension du comportement du modèle.

L'idée centrale repose sur le fait que, dans un transformeur, chaque couche calcule une somme pondérée des représentations de la couche précédente via son mécanisme d'attention. En multipliant les matrices d'attention des couches successives, on peut retracer comment l'information provenant des premiers jetons se propage jusqu'à la sortie finale. Cela s'apparente à suivre un chemin dans un graphe orienté, où les nœuds sont les jetons et les arêtes sont les poids d'attention. La matrice de rollout qui en résulte offre une vue globale de l'influence des jetons, souvent plus interprétable que l'examen d'une seule couche d'attention.

Contexte et motivation

Les transformeurs, introduits dans l'article de 2017 « Attention Is All You Need » par des chercheurs de Google DeepMind et University of Toronto, reposent sur des mécanismes d'auto-attention pour capturer les dépendances entre les jetons. Chaque couche calcule des scores d'attention indiquant dans quelle mesure chaque jeton prête attention aux autres. Cependant, ces scores ne sont pas directement interprétables comme une mesure d'importance. Par exemple, un jeton peut prêter attention à un signe de ponctuation qui a peu de poids sémantique, ou l'attention peut être répartie de manière diffuse sur de nombreux jetons. De plus, l'attention dans les couches ultérieures est calculée sur des représentations déjà transformées par les couches précédentes, de sorte que les poids bruts ne tiennent pas compte de cet effet cumulatif.

Les premières tentatives d'explicabilité se sont concentrées sur la visualisation des cartes d'attention de têtes ou de couches individuelles, mais celles-ci produisaient souvent des résultats contradictoires ou déroutants. Des chercheurs comme Anima Anandkumar et Jakob Uszkoreit ont noté que les schémas d'attention peuvent varier considérablement d'une couche à l'autre et d'une tête à l'autre, ce qui rend difficile d'en tirer des conclusions. Cela a motivé le développement de méthodes qui agrègent l'information à travers l'ensemble du réseau.

L'algorithme

Attention Rollout fonctionne sous l'hypothèse que les poids d'attention représentent une forme de routage de l'information. L'algorithme procède comme suit :

1. Pour chaque couche \( l \), obtenir la matrice d'attention \( A_l \) de forme (longueur_de_séquence, longueur_de_séquence), où \( A_l[i][j] \) est le poids d'attention du jeton \( i \) vers le jeton \( j \). Cette matrice est généralement moyennée sur toutes les têtes d'attention de cette couche.
2. Ajouter la matrice identité à chaque matrice d'attention pour tenir compte de la connexion résiduelle, ce qui permet à chaque jeton de conserver sa propre information. La matrice résultante est \( \tilde{A}_l = A_l + I \).
3. Normaliser chaque ligne de \( \tilde{A}_l \) pour que sa somme soit égale à 1, garantissant ainsi que les matrices restent stochastiques.
4. Calculer la matrice de rollout \( R \) en multipliant les matrices normalisées de toutes les couches : \( R = \tilde{A}_1 \cdot \tilde{A}_2 \cdot ... \cdot \tilde{A}_L \), où \( L \) est le nombre de couches.

La matrice résultante \( R \) donne l'attention totale qui circule de chaque jeton vers chaque autre jeton, en considérant tous les chemins possibles à travers le réseau. L'entrée \( R[i][j] \) peut être interprétée comme la proportion d'information provenant du jeton \( j \) qui influence la représentation du jeton \( i \) à la sortie.

Cette méthode a été proposée pour la première fois dans un article de 2020 par des chercheurs de OpenAI et Stanford AI Lab, qui ont démontré son efficacité sur GPT-2 et d'autres modèles. Ils ont montré qu'Attention Rollout pouvait identifier les jetons saillants dans des tâches telles que l'analyse de sentiments et la résolution de coréférences, surpassant souvent l'attention d'une seule couche en termes d'alignement avec les jugements humains.

Applications dans l'interprétation des modèles

Attention Rollout a été appliqué à une variété de tâches en traitement du langage naturel et au-delà. Dans la recherche en apprentissage automatique, il est utilisé pour :

• Identifier les jetons clés : En examinant la matrice de rollout, les praticiens peuvent voir quels mots ou sous-mots de l'entrée influencent le plus la sortie du modèle. Par exemple, dans une tâche de classification de sentiments, le rollout peut mettre en évidence des mots comme « terrible » ou « incroyable » comme étant très influents.
• Détecter les biais : Si un modèle prête systématiquement attention à des termes démographiques d'une manière qui conduit à des prédictions biaisées, le rollout peut révéler ces schémas, aidant ainsi les chercheurs à atténuer les comportements injustes.
• Déboguer les erreurs du modèle : Lorsqu'un modèle fait une prédiction incorrecte, le rollout peut montrer s'il s'est concentré sur des parties non pertinentes de l'entrée, guidant ainsi l'amélioration des données d'entraînement ou de l'architecture.
• Comparer les modèles : En calculant les matrices de rollout pour différents modèles, on peut comparer leur flux d'information interne, ce qui est utile pour la sélection de modèles et la compréhension des différences architecturales.

Au-delà du texte, Attention Rollout a été adapté pour les transformeurs de vision (ViT) utilisés dans les tâches de vision par ordinateur. Par exemple, des chercheurs de Google Cloud et Amazon Web Services l'ont appliqué à la classification d'images pour visualiser quelles régions d'une image contribuent à une prédiction, facilitant ainsi l'analyse d'imagerie médicale et les systèmes de conduite autonome.

Limites et critiques

Malgré sa popularité, Attention Rollout présente plusieurs limitations. Premièrement, l'hypothèse selon laquelle les poids d'attention peuvent être traités comme des distributions de probabilité du flux d'information n'est pas toujours valide. Les poids d'attention sont calculés sur la base de représentations apprises et ne reflètent pas nécessairement une relation causale. Des critiques comme Melanie Mitchell et Aleksander Madry ont soutenu que l'attention n'est pas une explication et que des méthodes comme le rollout peuvent produire des interprétations trompeuses.

Deuxièmement, la multiplication des matrices d'attention à travers les couches peut entraîner des problèmes numériques, en particulier dans les réseaux profonds avec de nombreuses couches. La matrice de rollout peut devenir trop diffuse, avec tous les jetons ayant une importance similaire, ou au contraire trop concentrée sur quelques jetons en raison de la multiplication répétée. Cela peut rendre les résultats moins utiles en pratique.

Troisièmement, Attention Rollout ne tient pas compte des transformations non linéaires (réseaux feed-forward, normalisation de couche) qui se produisent entre les couches d'attention. Ces transformations peuvent modifier considérablement le flux d'information, de sorte que le simple modèle multiplicatif peut manquer des effets importants.

Quatrièmement, la méthode moyenne l'attention sur toutes les têtes, ce qui peut masquer les rôles distincts des différentes têtes. Certaines têtes peuvent capturer des relations syntaxiques, tandis que d'autres capturent des relations sémantiques ; les moyenner fait perdre cette nuance.

Variantes et extensions

Pour répondre à certaines de ces limitations, les chercheurs ont proposé plusieurs variantes :

• Attention Flow : Cette méthode, introduite par des chercheurs de MIT CSAIL, traite l'attention comme un flux et utilise des algorithmes de flot maximal pour calculer l'importance des jetons, évitant ainsi l'hypothèse multiplicative.
• Rollout avec poids résiduels : Certaines implémentations pondèrent la matrice identité par un facteur qui reflète la force de la connexion résiduelle, plutôt que de l'ajouter uniformément.
• Propagation de pertinence par couches (LRP) : Cette technique, issue de la communauté du apprentissage profond, propage les scores de pertinence en arrière à travers le réseau, offrant une alternative aux méthodes basées sur l'attention.
• Gradients intégrés et SHAP : Ces méthodes d'attribution de caractéristiques ne reposent pas du tout sur les poids d'attention, mais plutôt sur les gradients ou des concepts de théorie des jeux. Elles sont souvent utilisées comme vérification croisée des résultats obtenus par Attention Rollout.

Malgré ces alternatives, Attention Rollout reste une référence populaire en raison de sa simplicité et de son efficacité computationnelle. Il ne nécessite aucun entraînement supplémentaire ni calcul de gradient, ce qui le rend facile à appliquer à n'importe quel transformeur pré-entraîné.

Considérations d'implémentation

Implémenter Attention Rollout est simple avec les frameworks d'apprentissage profond modernes. Dans PyTorch ou TensorFlow, on peut utiliser des hooks pour capturer les matrices d'attention pendant la passe avant. Les étapes clés sont :

1. Enregistrer des hooks sur chaque couche d'attention pour stocker les poids d'attention.
2. Après la passe avant, moyenner les poids d'attention sur toutes les têtes.
3. Ajouter la matrice identité et normaliser les lignes.
4. Multiplier les matrices séquentiellement.

Une considération pratique est que les matrices d'attention peuvent être volumineuses pour de longues séquences, ce qui entraîne une surcharge mémoire. Pour une séquence de longueur 1024, chaque matrice est de taille 1024x1024, et multiplier de nombreuses matrices de ce type peut être coûteux en calcul. Cependant, pour des entrées typiques, cela reste gérable.

Une autre considération est que la méthode suppose une architecture de transformeur standard. Pour les modèles avec attention croisée (comme les modèles encodeur-décodeur), le rollout doit être adapté pour gérer l'interaction entre l'attention de l'encodeur et celle du décodeur. Dans de tels cas, on peut calculer des rollouts séparés pour l'encodeur et le décodeur, ou les combiner de manière plus complexe.

Relation avec d'autres méthodes d'explicabilité

Attention Rollout s'inscrit dans un paysage plus large de techniques d'explicabilité pour les systèmes d'intelligence artificielle. Il est souvent comparé à :

• Méthodes basées sur les gradients : Ces méthodes calculent le gradient de la sortie par rapport aux embeddings d'entrée, fournissant une mesure de sensibilité. Elles sont plus fondées théoriquement mais nécessitent une rétropropagation.
• Méthodes par perturbation : Ces méthodes consistent à modifier les jetons d'entrée (par exemple, en les supprimant ou en les masquant) et à observer le changement dans la sortie. Elles sont indépendantes du modèle mais coûteuses en calcul.
• Modèles de substitution : Des techniques comme LIME ou SHAP entraînent des modèles interprétables localement pour approximer le modèle boîte noire. Elles sont utiles pour les données tabulaires mais moins adaptées au texte.

Attention Rollout est unique en ce sens qu'il est purement basé sur les poids d'attention internes du modèle, ne nécessitant aucun calcul supplémentaire au-delà d'une passe avant. Cela le rend particulièrement attrayant pour des applications en temps réel, comme les outils de débogage interactifs.

Directions futures

À mesure que les modèles de transformeurs continuent de croître en taille et en complexité, le besoin de méthodes d'explicabilité robustes augmente. Attention Rollout est susceptible d'évoluer de plusieurs manières :

• Intégration avec des méthodes causales : Combiner le rollout avec des techniques d'inférence causale pourrait fournir des attributions plus précises.
• Gestion des modèles multimodaux : Étendre le rollout aux modèles qui traitent simultanément du texte, des images et de l'audio, comme ceux développés par OpenAI et Anthropic.
• Analyse automatisée : Utiliser les sorties du rollout pour générer automatiquement des explications en langage naturel des décisions du modèle, ce qui pourrait être utile pour la conformité et l'audit.
• Évaluation comparative : Développer des références standardisées pour évaluer la fidélité des méthodes d'explicabilité, comme proposé par des chercheurs de Carnegie Mellon University et BAIR (Berkeley AI Research).

Malgré ses limites, Attention Rollout est devenu un outil fondamental dans la boîte à outils de l'explicabilité. Sa simplicité et son efficacité en ont fait une référence standard dans de nombreux articles de recherche et applications pratiques. À mesure que le domaine progresse, il sera probablement affiné et combiné avec d'autres techniques pour fournir des aperçus plus profonds sur le fonctionnement interne des réseaux de neurones.

Conclusion

Attention Rollout offre un moyen pratique et intuitif de comprendre comment les transformeurs traitent l'information. En propageant les poids d'attention à travers les couches, il fournit une perspective globale sur l'influence des jetons qui est souvent plus utile que l'inspection des couches individuelles. Bien qu'il ne soit pas sans défauts, sa facilité d'utilisation et son interprétabilité ont consolidé sa place dans la littérature sur l'explicabilité. Pour quiconque travaille avec des modèles de transformeurs, comprendre Attention Rollout est une étape précieuse pour démystifier ces systèmes puissants mais opaques.