Résumé Technique : Comment les Transformers Apprennent à Associer des Tokens : Les Termes Dominants du Gradient Apportent une Interprétabilité Mécaniste

Énoncé du Problème

Les grands modèles de langage (LLM) basés sur l'attention automatique ont démontré des capacités remarquables à capturer des connaissances factuelles et à générer du texte cohérent, principalement grâce à leur capacité à apprendre des associations sémantiques (par exemple, le lien entre « oiseau » et « a volé » ou l'interchangeabilité de « voiture » et « camion »). Bien que des études récentes aient identifié des mécanismes spécifiques tels que les têtes d'induction ou des relations sémantiques linéaires, une explication fondée sur des principes et ancrée dans la théorie, décrivant comment ces associations émergent lors de l'entraînement de transformers basés sur l'attention sur des données de langage naturel réelles, fait encore défaut.

Les travaux théoriques antérieurs reposent souvent sur des hypothèses irréalistes, telles qu'un langage structuré synthétique, des architectures simplifiées (sans encodages de position ni connexions résiduelles) ou des régimes d'entraînement non standard (par exemple, un entraînement séquentiel composant par composant). Ces écarts par rapport à la pratique limitent la généralisabilité de leurs insights aux LLM modernes. Cet article comble cette lacune en se demandant : Comment les associations sémantiques émergent-elles lors de l'entraînement de modèles de langage basés sur l'attention sur des données de langage naturel ?

Méthodologie

Les auteurs proposent un cadre théorique ancré dans la dynamique d'entraînement, en analysant spécifiquement l'approximation des termes dominants des gradients au cours des premières étapes de l'entraînement.

1. Configuration du Modèle : L'étude se concentre sur un transformer standard à $L$ couches basé sur l'attention, avec des encodages de position (relatifs), un masquage causal et des flux résiduels, entraîné par descente de gradient par lot complet standard sur des corpus de langage naturel.
2. Approximation des Termes Dominants : Reconnaissant que les transformers acquièrent des comportements fondamentaux (y compris les relations sémantiques) tôt dans l'entraînement, les auteurs dérivent des expressions fermées pour les matrices de poids (Sortie $W_O$, Valeur $V^{(l)}$, Requête-Clef $W^{(l)}$, et Encodage de Position $P^{(l)}$) en approxmant les mises à jour de gradient. Ils montrent que pour un nombre suffisant d'étapes $s$ (où $s \leq \eta^{-1} \min(\dots)$), les poids restent proches de leurs termes dominants de gradient.
3. Décomposition en Fonctions de Base : L'analyse théorique révèle que les poids appris peuvent être exprimés comme de simples compositions de trois fonctions de base fondamentales dérivées des statistiques du corpus :
   • Cartographie des Bigrammes ($\bar{B}$) : Capture les dépendances du token suivant ($P(e_j|e_i)$).
   • Cartographie de l'Interchangeabilité ($\Sigma \bar{B} = \bar{B}^\top \bar{B}$) : Capture la similarité fonctionnelle entre les tokens basée sur la similarité de leurs distributions de tokens précédents (par exemple, synonymes ou rôles grammaticaux partagés).
   • Cartographie du Contexte ($\bar{\Phi}$) : Encode la co-occurrence préfixe-suffixe à plus long terme, représentant le contexte moyen dans lequel un token apparaît.

Contributions Clés

L'article apporte trois contributions principales :

1. Caractérisation Explicite des Poids : Il présente la première caractérisation explicite et fermée des poids dans les transformers basés sur l'attention entraînés sur des corpus de texte réel sous la perte de prédiction du token suivant. Il est démontré que les poids sont des compositions des trois fonctions de base mentionnées ci-dessus.
2. Interprétation Mécaniste des Associations Sémantiques : Les auteurs interprètent les caractéristiques apprises dans les poids comme des compositions spécifiques de cartographies de bigrammes, d'interchangeabilité et de contexte. Ils démontrent comment ces compositions permettent au modèle de capturer des associations sémantiques (par exemple, comment « poisson » se rapporte à « étang » via la cartographie du contexte, ou comment « rouge » se rapporte à « camion » via la cartographie des bigrammes).
3. Validation Empirique : Les prédictions théoriques sont validées à la fois sur un modèle d'attention jouet à 3 couches entraîné sur le jeu de données TinyStories et sur un LLM pratique à grande échelle (Pythia-1.4B) entraîné sur OpenWebText.

Résultats

• Alignement Théorique : Les expériences sur un transformer à 3 couches entraîné sur TinyStories montrent que la similarité cosinus entre les termes dominants prédits théoriquement et les poids réellement appris reste extrêmement élevée (par exemple, >0,99 pour les matrices d'Attention et de Valeur) même après 100 époques d'entraînement.
• Structure Sémantique : L'analyse qualitative confirme que les fonctions de base capturent des structures linguistiques significatives. Par exemple, la cartographie des bigrammes $\bar{B}$ corrèle les adjectifs avec les noms qu'ils modifient, tandis que la cartographie du contexte $\bar{\Phi}$ corrèle les entités avec leurs environnements typiques (par exemple, « poisson » avec « étang »).
• Généralisation aux LLM Réels : Lorsqu'appliquée à Pythia-1.4B, les caractéristiques de termes dominants théoriques montrent un fort accord avec les embeddings appris et les corrélations d'attention, en particulier au cours des premières étapes de l'entraînement. L'analyse révèle que, bien que les poids dérivent vers des représentations de connaissances plus riches à mesure que l'entraînement progresse, ils conservent les caractéristiques associatives fondamentales prédites par la théorie pendant une durée significative.
• Coopération des Composantes : L'analyse du calcul complet du modèle (Éq. 12) illustre comment le flux résiduel fournit une prédiction moyenne du token suivant, tandis que le bloc d'attention automatique (piloté par $\bar{Q}$ et $\bar{\Phi}^\top \Sigma \bar{B}$) affine cette prédiction en se concentrant sélectivement sur les tokens les plus indicatifs du token suivant, basés sur les statistiques du corpus.

Importance

L'article prétend fournir une fondation mécaniste pour l'apprentissage de représentations dans les grands modèles de langage en reliant la dynamique d'entraînement du deep learning à la théorie linguistique. En ancrant l'analyse dans des distributions de texte naturalistes et des procédures d'entraînement standard, ce travail réduit considérablement l'écart entre l'analyse théorique formelle et l'utilisation pratique des LLM.

Les auteurs soutiennent que leurs résultats offrent un moyen rigoureux de contextualiser la fonction de chaque composant du transformer. Plutôt que de considérer le modèle comme une boîte noire, la théorie suggère que le comportement du modèle se développe en établissant d'abord ces caractéristiques associatives fondamentales (bigrammes, interchangeabilité, contexte), qui servent d'ancres pour apprendre des relations sémantiques plus complexes. Ce travail ouvre des voies pour la recherche en interprétabilité en suggérant que les matrices de poids à travers différents composants peuvent être décomposées en fonctions simples de statistiques de corpus partagées, permettant des hypothèses plus larges sur la manière dont les concepts émergent dans les modèles au-delà de mécanismes spécifiques.