Résumé Technique : Co-Ablation Conditionnelle (COAX)

1. Énoncé du Problème

L'interprétabilité mécaniste repose sur l'attribution de scores d'importance à des composants individuels du modèle (par exemple, les têtes d'attention) pour comprendre comment un modèle produit des comportements spécifiques. L'approche standard utilise le score de premier ordre : abler une unité de manière isolée, mesurer le changement de sortie, et agréger ces effets. Cela suppose que l'importance d'un composant est additive.

Cependant, cette hypothèse échoue en présence d'auto-réparation (l'effet "Hydre"), où un modèle possède des composants de secours dormants qui ne s'activent que lorsqu'un composant primaire est supprimé.

• L'angle mort : Sur un modèle intact, un composant de secours semble non pertinent (faible importance) car il est silencieux. Lorsqu'un composant primaire est ablé, le composant de secours prend le relais, réparant le dommage et atténuant l'effet mesuré du primaire.
• Conséquence : Les scores de premier ordre interprètent mal les deux côtés de la redondance : ils sous-estiment le primaire (car le secours répare l'erreur) et manquent complètement le composant de secours (car il est dormant). Cela conduit à une attribution erronée, à l'échec de l'élimination de capacités et à un élagage (pruning) suboptimal du modèle.

2. Méthodologie : Co-Ablation Conditionnelle (COAX)

Les auteurs recadrent le problème comme une complétion de circuit conditionnelle : étant donné un ensemble de composants primaires $S$, identifier les unités restantes dont l'importance croît significativement une fois $S$ retiré.

Le Score COAX

COAX est un score sans étiquette (label-free) et ancré sur la sortie, qui mesure la croissance conditionnelle de l'effet d'ablation d'une unité.

1. Géométrie de Fisher : La méthode mesure les effets d'ablation dans la métrique induite par la distribution de sortie (information de Fisher). Pour une unité $u$, l'effet d'ablation est la perturbation du logit $\delta z_u = z_0 - z_{\{u\}}$. L'énergie de cet effet est la moyenne de la distance KL, approximée par la norme pondérée par la Fisher de la caractéristique centrée $\tilde{\delta z}_u$.
2. Croissance Conditionnelle :
   • Soit $S$ un ensemble de semence de composants primaires (identifiés par une méthode de premier ordre).
   • Calculer l'énergie de l'ablation d'une unité candidate $u$ en isolation : $E(\delta z_u | \emptyset)$.
   • Calculer l'énergie de l'ablation de $u$ après que $S$ a été ablé : $E(\delta z_u | S)$.
   • Le score COAX est la différence :
     $$\text{comp}_u(S) = E(\delta z_u | S) - E(\delta z_u | \emptyset)$$
3. Base Théorique :
   • Proposition 1 : Établit l'identité de Fisher, reliant le produit scalaire des caractéristiques centrées à la matrice de Gram de Fisher.
   • Proposition 2 (Angle Mort de l'Additivité) : Prouve que tout score basé uniquement sur des statistiques d'état propre (ex: activation propre ou gradients propres) ne peut pas distinguer un pur composant de secours (dormant sur le modèle propre) d'une unité inerte.
   • Proposition 3 (La Croissance Conditionnelle comme Synergie) : Montre que pour une unité dormante, le score de croissance conditionnelle approxime l'énergie de sa synergie par paire avec l'ensemble de semence $S$. Ainsi, COAX récupère l'interaction de second ordre que les scores de premier ordre rejettent.
4. Efficacité : Contrairement aux calculs explicites de synergie par paire qui nécessitent $O(|U|^2)$ passes avant, COAX ne nécessite que $O(|U|)$ passes (un seul balayage d'ablation simple et un balayage conditionnel par semence). C'est une méthode uniquement par passage avant (forward-only) et ne nécessite aucun passage arrière (backward pass) ni étiquettes de tâche.

3. Principales Contributions

1. Formalisation de l'Auto-Réparation : L'article formalise l'auto-réparation comme un échec de l'additivité de l'importance des composants et introduit la "complétion de circuit conditionnelle" comme la tâche de récupération des compensateurs dormants.
2. Score COAX : Introduit un score sans étiquette, ancré sur la géométrie de Fisher, qui expose les sauvegardes dormantes en mesurant la croissance de l'énergie d'ablation après conditionnement sur un ensemble primaire.
3. Validation Causale : Démontre que les têtes récupérées sont de véritables sauvegardes via des vérifications mécanistes (courbes de réveil, transfert d'attribution directe au logit) et un patch de contre-factuel (le gel des sauvegardes supprime l'auto-réparation).
4. Utilité Aval : Montre que la récupération de ces sauvegardes corrige l'attribution masquée, permet l'élimination complète d'une capacité et améliore l'élagage structuré.

4. Résultats Expérimentaux

Découverte de Sauvegarde (Circuit IOI de GPT-2-small)

• Tâche : Identifier 8 têtes de mouvement de nom (name-mover) documentées étant donné les 3 têtes primaires.
• Performance :
  • Baselines de premier ordre : Saliance par ablation simple (0.33 AUC), AtP (0.60), EAP-IG (0.70), et AtP*GradDrop sensible à l'auto-réparation (0.82).
  • COAX : Atteint 0.91 ROC-AUC.
  • Signification : COAX place 6 des 8 sauvegardes documentées dans ses 20 meilleurs candidats. L'amélioration par rapport à la meilleure baseline basée sur le gradient (AtP*) est statistiquement significative ($p < 0.05$) en contrôlant l'ensemble de semence.
• Benchmark Synthétique : Sur un benchmark synthétique avec 100 sauvegardes implantées, les scores en état propre performent au niveau du hasard ou en dessous (0.42–0.51), tandis que COAX atteint 0.90, confirmant l'angle mort théorique des méthodes de premier ordre.

Généralisation

• Échelle : Le phénomène de "réveil" (augmentation de la norme de sortie lors de l'ablation du primaire) est observé sur les modèles GPT-2-small, medium et large.
• Architecture : Appliqué au circuit d'induction sur 8 modèles (Pythia, GPT-Neo, Gemma, Llama, Qwen, OLMo). COAX récupère des ensembles compensateurs qui font chuter la log-probabilité d'induction (2.1x à 12x plus que l'aléatoire), démontant la transférabilité sans étiquette.

Applications Aval

1. Récupération d'Attribution : L'ablation des primaires seuls fait chuter la différence de logit IOI de seulement 0.22 (masquée par l'auto-réparation). L'ajout des sauvegardes récupérées par COAX rétablit la chute à 1.76, révélant le véritable effet causal.
2. Élimination de Capacité (Knockout) : Pour désactiver le comportement IOI, il faut supprimer les sauvegardes.
   • Supprimer uniquement les primaires : L'exactitude reste à ~0.97 (auto-réparation active).
   • Supprimer les primaires + les sauvegardes COAX : L'exactitude chute à 0.70, correspondant à l'oracle documenté (0.72).
   • Supprimer les primaires + les têtes suivantes classées par premier ordre : L'exactitude chute à 0.24, indiquant un sur-élagage dans les composants de base.
3. Élagage Structuré : Une stratégie d'élagage séquentiel qui remesure les scores COAX après chaque suppression (pour tenir compte des nouvelles sauvegardes exposées) surpasse toutes les baselines (Random, Magnitude, Wanda, Taylor) à travers différentes échelles de 124M à 7B de paramètres. À 50 % de sparsité, COAX maintient une perplexité nettement plus proche du modèle dense que n'importe quel concurrent.

5. Signification et Revendications

L'article soutient que l'importance d'un composant n'est pas une propriété intrinsèque d'une unité mais est conditionnelle à l'état du reste du circuit.

• Motivation : Dans les circuits robustes, les composants les plus importants sont souvent ceux qui restent silencieux jusqu'à ce que leur primaire échoue. Les méthodes de premier ordre, qui évaluent les unités en isolation, manquent systématiquement ces "héros cachés".
• Changement Opérationnel : COAX déplace le paradigme de "scorer les unités en isolation" vers "scorer les unités par rapport à ce qui a été retiré".
• Limites et Portée :
  • COAX est un outil de complétion, pas un outil de découverte autonome ; il nécessite un ensemble de semence primaire (qui peut être trouvé par des méthodes de premier ordre).
  • Il est plus efficace dans les régimes de substitution dormante (comme les sauvegardes IOI). Dans les régimes où la redondance est partagée entre des têtes homogènes co-activées (comme l'induction), COAX trouve des ensembles porteurs de charge mais n'identifie pas de classe de "sauvegarde" distincte.
  • La méthode est actuellement instanciée à la granularité de la tête d'attention ; elle ne récupère pas de signaux forts dans les circuits dominés par les MLP (ex: le circuit "greater-than"), suggérant que le mécanisme du circuit (MLP vs Attention) dicte la récupérabilité de l'auto-réparation.

L'article conclut que l'interprétabilité fidèle des modèles robustes doit être conditionnelle : la redondance n'est pas un bruit à moyenner, mais une structure qui doit être conditionnée pour révéler le circuit complet.