Résumé Technique : La structure de la tâche inverse l'encodage de l'état par couche dans les modèles de séquence

Énoncé du problème

L'interprétabilité mécaniste des modèles de séquence traite souvent les encodages d'état par couche comme des traits architecturaux fixes. L'hypothèse prédominante est que les modèles récurrents (ex. LSTM, GRU) et les modèles d'espace d'état (SSM comme Mamba) concentrent l'état lisible dans les couches tardives, tandis que les modèles basés sur l'attention (Transformers) distribuent l'accumulation de l'état à travers la profondeur. Cet article conteste la vision selon laquelle ces stratégies d'encodage sont des invariants architecturaux stables. Il postule que la même architecture inverse son profil d'encodage d'état par couche lorsque la structure computationnelle de la tâche change. De plus, la littérature confond souvent deux moteurs potentiels de ces schémas : la structure algébrique (spécifiquement la commutativité) versus la structure computationnelle (mises à jour de préfixes vs suivi de pile). Les auteurs visent à démêler ces facteurs et à déterminer si l'encodage de l'état est une propriété de l'architecture seule ou une interaction entre l'architecture et la tâche.

Méthodologie

L'étude emploie une suite expérimentale contrôlée à travers cinq familles d'architectures : Transformers, Mamba, Mamba-2, LSTM et GRU.

Tâches :
Les auteurs utilisent trois tâches de suivi d'état formelles et un benchmark semi-réel :

1. Parité : Une tâche de mise à jour de préfixe à bit unique et commutative.
2. Composition de permutations $S_3$ : Une tâche de mise à jour de préfixe à 6 classes et non commutative. Elle sert de discriminant critique : si la commutativité algébrique conduit le profil d'encodage, $S_3$ devrait se comporter comme Dyck ; si la structure computationnelle conduit le profil, $S_3$ devrait se comporter comme la Parité.
3. Dyck-$k$ : Une tâche de suivi de pile à profondeur bornée (correspondance de parenthèses) avec une profondeur maximale de 10.
4. Code-Depth : Une tâche de profondeur d'imbrication de jetons dérivée de Python dans CodeSearchNet, servant de validation semi-réelle.

Configuration expérimentale :

• Tâches formelles : Les expériences utilisent des modèles à 4 couches avec $d_{model}=128$. Les architectures sont appariées en termes de paramètres (environ 830K paramètres) lorsque cela est possible. L'entraînement se déroule sur une longueur de séquence $L=40$, avec une évaluation à $L \in \{40, 60, 80, 100, 120\}$.
• Échelle pré-entraînée : Le fine-tuning est effectué sur Mamba-130M (24 couches) et Pythia-160M/410M (12/24 couches).
• Mesures :
  • Probing linéaire : Entraînement de classifieurs de régression logistique sur les états cachés gelés pour mesurer la lisibilité de l'état à chaque couche.
  • Interventions causales : Ablation zéro (projection des directions de sondage) et ablation de sous-espace (suppression des principales directions PCA ou singulières) pour tester la nécessité fonctionnelle.
  • Généralisation OOD : Application d'ablations entraînées à $L=40$ à des longueurs non vues ($L=80, 100, 120$).
  • Attribution de composant : Ablation du composant Conv1D de Mamba sous un backend pur PyTorch équivalent pour isoler sa contribution selon les tâches.
  • Patching d'activation : Dans le Mamba pré-entraîné, injection d'activations de positions intermédiaires d'entrées propres dans des entrées corrompues pour mesurer la récupération du logit.

Contributions clés

1. Inversion dépendante de la tâche : L'article démontre que le schéma « concentration tardive » vs « construction distribuée » n'est pas un invariant architectural. Il s'inverse selon la tâche :
   • Tâches de mise à jour de préfixe (Parité, $S_3$) : Les modèles récurrents et les SSM concentrent l'état lisible dans la dernière couche, tandis que les Transformers le construisent progressivement à travers les couches.
   • Suivi de pile borné (Dyck) : Le schéma s'inverse. Les Transformers atteignent une lisibilité quasi parfaite immédiatement (couches précoces), tandis que les modèles récurrents et les SSM construisent la représentation progressivement.
2. Démêlage de l'algèbre et de la computation : En introduisant la tâche non commutative $S_3$, les auteurs montrent que le profil d'encodage suit la structure computationnelle (mise à jour de préfixe vs pile) plutôt que la structure algébrique (commutativité). $S_3$ se groupe avec la Parité à travers toutes les architectures, réfutant l'hypothèse que la commutativité est le principal moteur du profil d'encodage.
3. Découplage Lisibilité–Causalité : L'étude révèle que la couche où l'état est le plus linéairement lisible (meilleure couche de sonde) ne coïncide pas toujours avec la couche qui est causalement nécessaire pour la tâche. Ce découplage est particulièrement prononcé dans les modèles pré-entraînés, où les représentations sont plus redondantes et distribuées.
4. Dépendance de la tâche au niveau du composant : Même au sein d'une seule architecture (Mamba), l'importance de composants spécifiques varie selon la tâche. Par exemple, le composant Conv1D est critique pour $S_3$ et la Parité mais n'a aucun effet significatif sur la performance de Dyck, contestant la notion d'un « mécanisme » unique pour une architecture.

Résultats clés

• Niveau comportemental : Les bases récurrentes généralisent parfaitement sur la Parité et $S_3$ et surpassent les SSM et les Transformers sur Dyck. L'ordre de performance de $S_3$ (Récurrent > SSM > Transformer) reflète celui de la Parité, et non celui de Dyck.
• Profils de sondage (Probing) :
  • Parité/$S_3$ : Mamba et les LSTM montrent un saut brusque de la précision de sondage à la couche finale ($L_3$ dans les modèles à 4 couches). Les Transformers montrent une augmentation graduelle de $L_1$ à $L_3$.
  • Dyck : Les Transformers atteignent une précision élevée à $L_0$ ou $L_1$. Mamba et les LSTM montrent une augmentation graduelle, n'atteignant la saturation qu'aux couches ultérieures.
  • Confirmation de $S_3$ : $S_3$ suit le profil de la Parité (concentration tardive pour les RNN/SSM, progression graduelle pour les Transformers), confirmant l'hypothèse de la structure computationnelle.
• Interventions causales :
  • Dans les modèles formels à 4 couches, l'ablation de la meilleure direction de sonde cause des chutes significatives de précision, et ces directions restent nécessaires à des longueurs de séquence plus longues (OOD).
  • Dans les modèles pré-entraînés, une scission se produit :
    • Pythia (Dyck) : Un goulot d'étranglement marqué au milieu des couches existe (l'ablation à $L_6$–$L_7$ fait chuter la précision de ~81 % à 160M). La meilleure couche de sonde ($L_{11}$) n'est pas la plus critique causalement.
    • Mamba (Parité/Dyck/$S_3$) : La couche finale est hautement lisible, mais l'ablation d'une direction unique cause des chutes négligeables (<0,2 %). Cependant, le patching d'activation en position médiane à la couche finale récupère ~97–98 % de l'écart de logit, indiquant que l'état est suffisant mais distribué sur de nombreuses directions, et n'est pas capturé par une sonde linéaire unique.
• Attribution Conv1D : Supprimer Conv1D de Mamba cause une chute importante sur $S_3$ (-12,64 pp) et la Parité (-6,83 pp), mais aucun effet significatif sur Dyck (+4,04 pp, non significatif). Cela indique que l'importance des composants dépend de la tâche.

Signification et affirmations

L'article soutient que demander « comment l'architecture X encode l'état ? » est incomplet sans spécifier la structure computationnelle de la tâche. Les auteurs affirment que :

1. Les signatures mécanistes sont des propriétés conjointes : Les stratégies d'encodage de l'état ne sont pas des traits architecturaux fixes, mais émergent de l'interaction entre les biais inductifs du modèle et les exigences computationnelles de la tâche (mise à jour de préfixe vs pile).
2. Prudence méthodologique : Extrapoler des découvertes mécanistes d'un régime de tâche (ex. mises à jour de préfixes) à un autre (ex. suivi de pile bornée) risque de confondre des signatures spécifiques à la tâche avec des invariants architecturaux.
3. Lisibilité vs Causalité : Le sondage identifie où l'information est linéairement disponible, mais cela ne correspond pas toujours aux endroits où la computation est limitée, surtout dans les grands modèles pré-entraînés où la redondance est élevée.
4. Algèbre vs Computation : Dans le contexte de ces tâches formelles, la structure computationnelle (mise à jour de préfixe séquentielle vs pile hiérarchique) est un meilleur prédicteur des profils d'encodage par couche que les propriétés algébriques comme la commutativité.

Les auteurs concluent que leurs conclusions s'appliquent aux tâches formelles, aux échelles pré-entraînées (jusqu'à 410 millions de paramètres) et aux benchmarks de code semi-réels, suggérant une interaction robuste entre la structure de la tâche et l'architecture dans la modélisation de séquences.