Résumé technique : Similarité tensorielle pour l'interprétabilité mécaniste

Énoncé du problème

L'interprétabilité mécaniste vise à décomposer les réseaux de neurones en parties constitutives significatives pour comprendre leurs calculs internes. Un prérequis critique pour ce domaine est la capacité de vérifier rigoureusement si deux composants décomposés (ou deux modèles différents) implémentent la même fonction sous-jacente. Les méthodes d'évaluation actuelles souffrent de deux limitations principales :

1. Similarité empirique/comportementale : Les méthodes comparant les sorties ou les activations des modèles sur des jeux de données spécifiques sont aveugles aux mécanismes hors distribution (OOD). Elles ne peuvent pas détecter des problèmes tels que des portes dérobées qui ne se déclenchent que sur des entrées spécifiques, ni échouer à généraliser si la décomposition repose sur des heuristiques approximatives plutôt que sur la fonction véritable.
2. Similarité basée sur les paramètres : Les méthodes comparant les matrices de poids (par exemple, la similarité cosinus matricielle) sont hautement sensibles aux symétries de l'espace des poids (telles que la permutation et le mise à l'échelle). Elles échouent à reconnaître l'équivalence fonctionnelle lorsque les modèles sont entraînés différemment ou lorsque les mécanismes sont distribués sur plusieurs couches plutôt que concentrés dans des neurones individuels.

Il manque une métrique rigoureuse et indépendante du jeu de données, invariante à ces symétries et capable de capturer l'équivalence fonctionnelle globale.

Méthodologie

Les auteurs proposent la Similarité Tensorielle, une métrique basée sur les poids conçue pour les réseaux de neurones basés sur des tenseurs (spécifiquement les modèles multilinéaires comme les transformateurs bilinéaires). Cette approche remplace l'utilisation de fonctions d'activation non linéaires standard par des architectures polynomiales (multilinéaires) afin de permettre une analyse algébrique des poids.

Changement architectural fondamental

La méthode repose sur des modèles multilinéaires où les fonctions d'activation sont remplacées par un contrôle multiplicatif (par exemple, des couches bilinéaires et une attention bilinéaire).

• Couches bilinéaires : Elles transforment une entrée relevée $\tilde{x} = (1, x)$ en une sortie via un polynôme de degré 2 : $A(x) = D((L\tilde{x}) \odot (R\tilde{x}))$.
• Représentation multilinéaire : Ces couches peuvent être exprimées comme une contraction d'un tenseur de poids $A$ avec $n$ copies de l'entrée : $A(x) = A \cdot \text{in}^{\otimes n} \tilde{x}$.
• Modèles profonds : Les réseaux multilinéaires profonds sont construits comme des décompositions en arbre de tenseurs globaux, permettant un calcul efficace sans matérialiser le tenseur global exponentiellement grand.

La métrique de similarité tensorielle

La métrique est définie comme un produit scalaire généralisé entre deux tenseurs de poids $A$ et $B$, normalisé par leurs normes induites :
$$\text{sim}_M(A, B) = \frac{\langle A | M | B \rangle}{\|A\|_M \|B\|_M}$$
où $M$ est un tenseur métrique.

Pour garantir que la métrique satisfasse aux critères d'invariance par symétrie et d'indépendance du jeu de données, les auteurs dérivent une forme spécifique pour $M$ :

1. Symétrisation : Basée sur l'isomorphisme de polarisation, les auteurs établissent que deux modèles multilinéaires sont fonctionnellement équivalents si et seulement si leurs tenseurs de poids ont le même représentant dans le sous-espace symétrique. La métrique projette les tenseurs sur ce sous-espace symétrique pour éliminer les degrés de liberté liés aux symétries de permutation des entrées.
2. Métrique gaussienne : Pour affiner davantage la métrique et la relier à la similarité comportementale, les auteurs définissent $M$ en utilisant l'espérance du tenseur d'entrée sous une distribution gaussienne standard ($x \sim \mathcal{N}(0, I)$). En utilisant le théorème d'Isserlis, cela aboutit à une métrique sous forme fermée qui calcule le produit scalaire des activations sur tout l'espace d'entrée gaussien, plutôt que sur un jeu de données échantillonné.
3. Calcul efficace : Un algorithme récursif basé sur la matrice de Gram permet le calcul de cette similarité pour les réseaux profonds en traitant les tenseurs locaux couche par couche, évitant ainsi la nécessité d'instancier le tenseur global complet.

Contributions clés

1. Une métrique basée sur les poids fondée sur des principes : L'article introduit une métrique invariante aux symétries de l'espace des poids (permutations, mise à l'échelle) et indépendante du jeu de données d'entrée. Elle vérifie formellement l'équivalence fonctionnelle en examinant si deux modèles implémentent la même logique polynomiale.
2. Analyse sous forme fermée et sans données : Contrairement aux métriques comportementales, la similarité tensorielle ne nécessite pas d'échantillonner des données. Elle intègre sur tout l'espace d'entrée gaussien, ce qui la rend capable de détecter des mécanismes OOD (comme des portes dérobées) invisibles pour les métriques de performance standard.
3. Capacités de localisation : La métrique peut être « découpée » pour se concentrer sur des dimensions de sortie spécifiques ou « différencier » des tenseurs pour isoler les changements fonctionnels entre des points de contrôle, permettant ainsi d'attribuer des mécanismes spécifiques à des sorties spécifiques.
4. Évolutivité : L'algorithme récursif garantit que la métrique est calculable pour les réseaux profonds, évitant les bornes de convergence exponentielles ou quadratiques associées à l'échantillonnage de Monte Carlo.

Résultats

Les auteurs valident la similarité tensorielle sur quatre domaines distincts, démontrant sa supériorité par rapport aux métriques existantes (CKA, cosinus matriciel, similarité comportementale) :

1. Oubli catastrophique (Vision) : Dans un scénario d'entraînement progressif sur SVHN, la similarité tensorielle localise avec précision la perte de connaissances pour des chiffres spécifiques (par exemple, le chiffre '9') lorsque le jeu de données est réduit. Elle montre une chute brutale de similarité pour la classe affectée tout en restant stable pour les autres, alors que les métriques comportementales affichent des signaux ambigus.
2. Grokking (Arithmétique modulaire) : Dans des tâches d'addition modulaire, la similarité tensorielle suit la réorganisation continue du modèle vers des fréquences de Fourier spécifiques. Elle révèle que la réorganisation structurelle se poursuit même après que la précision a bondi vers la perfection, offrant une vision plus lisible de la dynamique d'entraînement que les courbes de précision seules.
3. Détection de portes dérobées (OOD) : Lorsqu'une porte dérobée est injectée dans SVHN (se déclenchant sur un patch de losange noir), les métriques comportementales standard sur des données propres échouent à détecter le changement. La similarité tensorielle, en revanche, signale la porte dérobée à partir de données propres uniquement et localise le changement à la classe attaquée spécifique ('9'), démontrant sa capacité à détecter des mécanismes OOD.
4. Modélisation du langage : Appliquée à un transformateur à attention bilinéaire à deux couches entraîné sur The Pile, la similarité tensorielle révèle des structures de blocs claires et des transitions nettes dans les régimes d'entraînement que les métriques existantes brouillent. Elle suit la formation de statistiques complexes de n-grammes plus efficacement que les proxies comportementaux.

Importance et affirmations

L'article positionne la similarité tensorielle comme un outil fondamental pour l'interprétabilité mécaniste. Son importance principale réside dans la réduction du problème de mesure de similarité et de vérification de la fidélité, passant d'un problème d'approximation empirique à un problème algébrique résolu.

• Fidélité : En s'appuyant sur les propriétés algébriques des modèles multilinéaires, la métrique garantit que les décompositions capturent la véritable fonction sous-jacente plutôt que des heuristiques approximatives.
• Robustesse : Elle fournit des affirmations de robustesse plus fortes en s'appliquant à toute la distribution, éliminant efficacement les mécanismes qui ne se déclenchent que sur des entrées spécifiques (comme les portes dérobées).
• Limites : Les auteurs reconnaissent modestement que la méthode repose actuellement sur des modèles multilinéaires, ce qui exclut la plupart des architectures de pointe utilisant des non-linéarités non polynomiales. Cependant, ils soutiennent que les modèles multilinéaires servent de terrain d'essai vital pour l'interprétabilité et que des preuves suggèrent qu'ils peuvent égaler les performances des architectures standard.
• Perspectives futures : Ce travail suggère que la restriction aux modèles multilinéaires est un compromis nécessaire pour une analyse algébrique traitable, mais qu'il ouvre la voie à de futures méthodes d'attribution (par exemple, calculer la similarité entre des chemins dans des transformateurs profonds) et soutient la vision selon laquelle les architectures multilinéaires sont plus transparentes que les réseaux profonds ordinaires.