Résumé Technique : Assign et Add : Une Étude Mécaniste de l'Arithmétique Compositionnelle

Énoncé du Problème

Les grands modèles de langage (LLM) démontrent une capacité à composer des compétences pour accomplir des tâches complexes non vues durant l'entraînement, pourtant les mécanismes spécifiques permettant cette « généralisation compositionnelle » restent élusifs. Bien que la généralisation soit souvent attribuée à la composition de mécanismes atomiques, l'interaction entre des mécanismes distincts — spécifiquement la liaison de variables (rappel associatif via des têtes d'induction) et l'arithmétique modulaire (représentations structurées via des caractéristiques de Fourier) — n'a pas été étudiée systématiquement dans un cadre contrôlé.

Ce document répond à la question suivante : Comment les transformers apprennent-ils à composer la liaison de variables et les mécanismes d'addition modulaire pour généraliser au-delà de la distribution d'entraînement, et comment ces capacités émergent-elles durant l'entraînement ?

Méthodologie

Les auteurs emploient un dispositif expérimental contrôlé impliquant un petit transformer mécanistiquement interprétable, entraîné sur une tâche synthétique combinant l'assignation de variables et l'addition modulaire.

• Architecture du Modèle : Un transformer à 2 couches et une seule tête avec $d_{model} = d_{head} = 128$ et $d_{mlp} = 512$. L'architecture omet délibérément les biais, les couches de normalisation et le MLP de la Couche 1 afin de simplifier la découverte de circuits. Il utilise des plongements de position absolus de type GPT-2 et un masquage causal.
• Définition de la Tâche : Le modèle est entraîné pour effectuer une addition modulaire ($N=59$) où les opérandes peuvent être des constantes ou des variables assignées dynamiquement dans le contexte (ex: c=1, b=17, a=42, b+19=?).
• Distribution des Données : Les séquences sont tokenisées avec un remplissage aléatoire pour empêcher la dépendance aux positions absolues. L'ensemble d'entraînement est partitionné pour évaluer la généralisation :
  • 0-variable (0-var) : Addition de constantes pures.
  • 1-variable (1-var) & 2-variables (2-var) : Addition impliquant des variables assignées.
  • Ensembles de test (Held-out) : Environ 30 % des paires d'addition sont retirées de l'entraînement (par exemple, si 1+3 est inédit, toutes les séquences contenant cette opération sont exclues). De plus, des variables spécifiques sont restreintes de certaines positions dans l'ensemble d'entraînement pour tester la généralisation positionnelle.
• Techniques d'Analyse : L'étude utilise des outils d'interprétabilité mécaniste (via la bibliothèque TransformerLens) pour analyser les motifs d'attention, les vecteurs du flux résiduel et les matrices de poids. Elle emploie également des sondes linéaires et une analyse théorique du gradient pour suivre la dynamique d'entraînement.

Contributions Clés

1. Modèle Traitable pour l'Arithmétique Compositionnelle : Les auteurs présentent un modèle simplifié qui permet l'interprétation mécaniste de la manière dont l'assignation de variables se compose avec l'arithmétique modulaire.
2. Dynamiques d'Entraînement Découplées : L'analyse révèle que l'arithmétique modulaire est apprise préalablement et indépendamment de l'assignation de variables. Le modèle atteint une précision quasi parfaite sur les tâches 0-var tout en apprenant à effectuer la liaison de variables.
3. Explication Mécaniste de la Généralisation : Le papier identifie que la généralisation provient de la composition de deux circuits spécifiques :
   • Une tête d'induction modifiée dans la Couche 1 qui récupère les variables assignées et route leurs valeurs.
   • Un circuit d'addition modulaire dans la Couche 2 (spécifiquement la matrice Output-Value et le MLP de la Couche 2) qui effectue l'arithmétique sur les valeurs récupérées.
4. Cadre Théorique : Les auteurs fournissent une analyse théorique démontrant comment les dynamiques de gradient facilitent la composition des circuits de liaison de variables et d'addition modulaire, montrant que le gradient négatif sur la matrice Query-Key de la Couche 2 conduit naturellement à la formation d'un bloc de mémoire associative.

Résultats

Dynamiques et Phases d'Entraînement

Le processus d'entraînement présente trois phases distinctes :

1. Apprentissage de l'Addition Modulaire : Le modèle apprend d'abord à effectuer l'addition sur des séquences sans variables (0-var), atteignant une précision élevée avant que la liaison de variables ne soit pleinement établie.
2. Apprentissage de l'Assignation de Variables : Le modèle apprend ensuite la structure requise pour l'assignation de variables. Un "pic" brutal de précision (autour de l'étape 5 300–5 400) coïncide avec l'émergence du circuit de routage d'assignation de variables, où l'attention de la Couche 1 route correctement les constantes vers leurs variables assignées.
3. Phase de Raffinement : Une seconde phase (étapes 14 500–16 000) implique le raffinement du routage pour les séquences "difficiles" (ex: variables spécifiques dans des positions tenues en réserve), corrigeant les échecs initiaux d'attention à certaines positions de variables.

Capacités de Généralisation

• Généralisation Positionnelle : Le modèle se généralise à de nouveaux ordonnancements de variables (ex: prédire b+14 après avoir vu 14+b).
• Généralisation Hors-Distribution (OOD) : Le modèle prédit avec succès des paires d'addition inédites (ex: 1+3) qui n'ont jamais été vues sous aucune forme durant l'entraînement, atteignant une précision de 99,4 % sur des ensembles de test curatés.
• Exigences en Données : La généralisation à des paires de constantes inédites nécessite un entraînement sur au moins ~25 % de la distribution totale ($f \approx 0,25$). La généralisation aux séquences à 2 variables nécessite une fréquence relative d'environ 20 % ($r \approx 0,2$) de séquences à 2 variables dans les données d'entraînement.

Découvertes Mécanistes

• Représentation Partagée : Les vecteurs du flux résiduel pour l'addition avec assignation de variables et l'addition de constantes pures sont presque identiques (similarité cosinus $\approx 1$) avant le MLP de la Couche 2, indiquant un mécanisme de calcul partagé.
• Caractéristiques de Fourier : Les préactivations des neurones du MLP de la Couche 2 correspondent à la forme théorique des caractéristiques de Fourier ($\cos(\frac{2\pi k}{N}n + \alpha_k) + \cos(\frac{2\pi k}{N}m + \beta_k)$), confirmant l'utilisation de mécanismes périodiques établis pour l'arithmétique modulaire.
• Modification de la Tête d'Induction : L'attention de la Couche 1 agit comme une tête "précédent-token" pour les constantes, portant attention à la variable précédant immédiatement celles-ci. La matrice Query-Key de la Couche 2 forme ensuite un motif d'identité fort sur les embeddings de variables, permettant au token = d'être directement en attention avec les constantes associées aux variables.

Signification

Le papier soutient que la généralisation compositionnelle dans les transformers peut être une conséquence naturelle de la compositionnalité des mécanismes internes. En démontrant que des mécanismes distincts (têtes d'induction et arithmétique basée sur Fourier) peuvent être appris indépendamment puis composés via le routage par attention, ce travail offre une base mécaniste pour comprendre comment les transformers généralisent à des tâches de raisonnement complexes.

Les auteurs font preuve de modestie quant à leur portée, notant que l'étude est limitée à un cadre restreint et contrôlé pour faciliter l'analyse. Ils ne prétendent pas que ces résultats expliquent directement tous les comportements des LLM, mais suggèrent qu'ils offrent un cadre pour comprendre comment les capacités de généralisation émergent dans les transformers, ce qui pourrait motiver de futures investigations sur des tâches plus complexes et des lois d'échelle.