Causal Evidence of Stack Representations in Modeling Counter Languages Using Transformers

Résumé Technique : Preuve Causale des Représentations de Pile dans la Modélisation des Langages de Comptage par les Transformers

Énoncé du Problème
Bien que l'interprétabilité mécaniste ait identifié avec succès des caractéristiques latentes dans les modèles de langage (LM) entraînés sur des langages formels, une lacune critique demeure dans l'établissement du rôle causal de ces caractéristiques. Des travaux antérieurs, tels que Tiwari et al. (2025), ont démontré que les transformers entraînés sur des langages de comptage (ex: Dyck-1, Shuffle-k) apprennent des représentations cohérentes avec une structure de pile sous-jacente. Cependant, ces études n'ont pas déterminé si ces représentations sont simplement corrélées à la tâche ou si elles sont causalement nécessaires à la computation du modèle. Ce document aborde la question suivante : les représentations de pile apprises sont-elles impliquées causalement dans la capacité du transformer à prédire le jeton suivant ?

Méthodologie
L'étude emploie un pipeline expérimental en trois étapes impliquant l'entraînement, le sondage (probing) et l'intervention causale :

1. Entraînement du Modèle : Un transformer de type encodeur est entraîné pour la prédiction du jeton suivant sur des langages Shuffle-k (où $k \in \{2, 4, 6, 8\}$). Le Shuffle-k est défini comme l'entrelacement de $k$ langages Dyck-1 indépendants, nécessitant au modèle de suivre la profondeur de $k$ piles de crochets distinctes simultanément. L'architecture du modèle consiste en une seule couche avec 4 têtes d'attention, une dimension d'embedding de 32 et une dimension cachée de 64. Le modèle atteint une précision de validation parfaite dès le 25ème époque, garantissant que toutes les analyses ultérieures sont effectuées sur un modèle entièrement convergé.
2. Sondage de la Profondeur de Pile : Des sondes linéaires sont entraînées pour prédire la profondeur de piles spécifiques ($c_i(t)$) à partir des états cachés du modèle ($h^{(l)}_t$) à la couche $l$. Pour s'assurer que les sondes capturent des informations structurelles authentiques plutôt que des corrélations spécieuses, une sonde de contrôle est entraînée sur des étiquettes perturbées. La métrique de Sélectivité est définie comme la différence de précision entre la sonde expérimentale et la sonde de contrôle. Une sélectivité élevée indique que la profondeur de la pile est véritablement encodée dans l'espace caché.
3. Intervention Causale (Ablation) :
   • Extraction de Direction : La matrice de poids $W$ de la sonde est décomposée via la Décomposition en Valeurs Singulières (SVD). Le vecteur singulier de droite principal, $w$, est identifié comme la direction principale correspondant à la représentation de la profondeur de la pile.
   • Ablation : Lors de l'inférence, la composante de l'état caché alignée avec $w$ est supprimée en utilisant l'opération $h' = h - \alpha(w \cdot h)w$, où $\alpha \in [0, 1]$ représente l'intensité de l'ablation.
   • Contrôle : Des directions aléatoires sont similairement ablées pour servir de base de référence.
   • Évaluation : Les états ablatés sont transmis directement à la tête de décodage (sans ré-exécuter les couches internes du transformer) pour mesurer l'impact sur la Précision Positionnelle ($A_{pos}$) et la Précision Séquentielle ($A_{seq}$).

Résultats Clés

• Existence de la Représentation : Les sondes linéaires extraient avec succès l'information de profondeur de pile avec une sélectivité élevée (ex: 85,8 % pour $k=6$, Pile 2), confirmant que le modèle apprend des représentations cohérentes avec une structure de pile.
• Nécessité Causale : L'ablation ciblée de la direction principale de la pile ($w$) provoque un effondrement dramatique des performances du modèle. À mesure que l'intensité de l'ablation $\alpha$ augmente de 0 à 1, la Précision Séquentielle chute à près de 0 %.
• Spécificité : L'ablation de directions aléatoires n'a aucun effet significatif sur la précision, confirmant que l'effondrement des performances est spécifique à la suppression de la représentation de la pile et non un artefact général de la manipulation de vecteurs.
• Sensibilité de la Métrique : L'étude observe que si la précision positionnelle baisse, la précision séquentielle s'effondre plus brutalement. Cela est attribué à la nature cumulative des erreurs dans la génération de séquences, où un seul jeton incorrect invalide l'intégralité de la chaîne.

Signification et Revendications
Le document affirme fournir une preuve empirique forte que les représentations de pile dans les transformers entraînés sur des langages de comptage ne sont pas seulement des artefacts émergents, mais sont causalement nécessaires à la computation interne du modèle et à la prédiction du jeton suivant.

En démontrant que la suppression d'une direction de représentation spécifique et apprise conduit à un échec quasi total, les auteurs soutiennent que le transformer utilise explicitement ces structures de type pile pour résoudre les tâches de raisonnement logique requises par les langages formels. Cela fait passer le domaine de l'analyse corrélationnelle (sondage) vers une compréhension mécaniste causale, validant que le comportement de la "boîte noire" du modèle repose sur des primitives computationnelles spécifiques et identifiables.

Directions Futures
Les auteurs suggèrent que les travaux futurs devraient explorer :

• L'intervention au milieu du réseau en faisant passer des états ablatés ou édités à travers le réseau de neurones à propagation avant (feed-forward) pour observer les changements dans les sorties suivantes.
• L'investigation du rôle de différentes architectures dans l'émergence, la redondance et la causalité de ces représentations.
• La compréhension de la manière dont les transformers gèrent la manipulation d'autres primitives computationnelles émergentes comme les compteurs.