Résumé Technique : Élagage de Nœuds Multi-Granulaire pour la Découverte de Circuits Causaux

1. Énoncé du Problème

La découverte de circuits causaux vise à identifier les sous-réseaux minimaux au sein des Grands Modèles de Langage (LLM) qui pilotent causalement des comportements spécifiques. Bien que les méthodes existantes aient progressé en interprétabilité mécaniste, elles font face à deux limitations principales :

1. Scalabilité Computationnelle : Les approches actuelles reposent principalement sur l'élagage d'arêtes (suppression des connexions entre les composants). Comme le nombre d'arêtes dans un transformeur croît de manière quadratique avec le nombre de nœuds ($O(|N|^2)$), l'élagage au niveau des arêtes devient computationnellement intraitable à mesure que la taille du modèle augmente. De plus, ces méthodes nécessitent souvent de stocker des représentations de flux résiduel et des activations intermédiaires désentrelacées, ce qui entraîne une utilisation prohibitive de la mémoire.
2. Limitations de Granularité : Les méthodes existantes opèrent généralement à une granularité grossière, traitant des blocs entiers de têtes d'attention ou de MLP comme des unités atomiques. Cela néglige la possibilité que des neurones (nœuds) spécifiques au sein de ces composants soient responsables de comportements cibles, ce qui aboutit à des circuits récupérés qui manquent de structures fines et conservent des composants inutiles.

2. Méthodologie

Les auteurs proposent un cadre d'élagage de nœuds multi-granulaire qui découvre des circuits minimaux à travers plusieurs niveaux de granularité au sein d'une étape d'optimisation unifiée.

2.1 Conception de la Granularité Hiérarchique

Au lieu d'élaguer les arêtes entre des composants grossiers, le cadre apprend des masques sur des nœuds à quatre niveaux distincts :

• Niveau Bloc : Blocs d'Attention entiers et Blocs MLP.
• Niveau Tête : Têtes d'Attention individuelles.
• Niveau Neurone : Neurones individuels au sein des couches d'Attention et des deux couches linéaires des blocs MLP (MLP1/Clés et MLP2/Valeurs).

2.2 Configuration d'Élagage Causal à Deux Flux

Pour identifier l'importance causale sans introduire de décalage de distribution (un risque avec l'ablation par zéro simple), la méthode utilise une passe directe à deux flux :

• Flux Propre (Clean Stream) : Le modèle traite une entrée conçue pour susciter un comportement spécifique (ex: prédire une année > 50).
• Flux Corrompu (Corrupted Stream) : Le modèle traite une version minimamente perturbée de l'entrée (ex: inverser les pronoms de genre ou changer l'année de départ) pour susciter une prédiction différente.
• Interpolation : Pour chaque nœud $n_i$, un masque apprenable $m_i \in [0, 1]$ interpole entre l'activation propre ($h^{clean}$) et l'activation corrompue ($h^{corrupted}$) :
  $$h^{node}_i = m_i \cdot h^{clean}_i + (1 - m_i) \cdot h^{corrupted}_i$$
  L'optimisation cherche à trouver des masques qui permettent au modèle de produire la sortie correcte (correspondant au comportement propre) tout en maximisant la parcimonie.

2.3 Optimisation et Paramétrage des Masques

• Distribution Hard Concrete : Les masques sont paramétrés à l'aide de la distribution Hard Concrete pour permettre une optimisation basée sur le gradient tout en induisant un comportement approximativement binaire (0 ou 1).
• Objectif Unifié : L'objectif d'entraînement combine trois composantes :
  1. Perte de Fidélité ($L_{faith}$) : Une divergence KL assurant que la distribution de sortie du circuit élagué correspond à celle du modèle de base.
  2. Perte de Tâche ($L_{task}$) : Une perte de classement basée sur une marge garantissant que le circuit distingue correctement les options de vérité terrain et les distracteurs.
  3. Régularisateur de Parcimonie ($L_0$) : Encourage les masques à être nuls, minimisant la taille du circuit.
• Cohérence Hiérarchique : Après l'entraînement, les masques sont binarisés avec une contrainte garantissant que si un nœud parent (ex: un bloc MLP) est élagué, tous ses nœuds enfants (ex: neurones au sein de ce bloc) sont également désactivés.

3. Contributions Clés

1. Cadre Multi-Granulaire : L'introduction d'un cadre d'optimisation unifié qui apprend des masques simultanément à travers les blocs de transformeurs, les têtes d'attention, les blocs MLP et les neurones individuels.
2. Scalabilité et Efficacité : La méthode évite la croissance quadratique de l'élagage d'arêtes et la nécessité de stocker les activations intermédiaires pour toutes les arêtes. Elle réduit l'utilisation de la mémoire GPU de 3× à 11× par rapport aux méthodes précédentes, permettant la découverte de circuits sur Llama 3.1-8B avec un seul GPU de 40 Go.
3. Découverte Fine : L'approche identifie des circuits nettement plus compacts que les lignes de base, atteignant une parcimonie élevée spécifiquement aux niveaux de bloc MLP et de neurone, là où les méthodes précédentes ont largement échoué.

4. Résultats Expérimentaux

La méthode a été évaluée sur trois tâches de découverte de circuits standards (Identification de l'Objet Indirect, Pronoms de Genre et "Greater Than") en utilisant des modèles GPT-2, GPT-2-XL et Llama.

4.1 Parcimonie et Compacité

• Élagage MLP et Neurone : La méthode proposée surpasse significativement les lignes de base d'élagage d'arêtes (EAP, EP) et l'élagage de poids non structuré (DiscoGP). Dans le cas le moins favorable, elle a élagué 33,34 % de blocs MLP en plus et 59,8 % de neurones en plus que les méthodes précédentes.
• Structures Dépendantes de la Tâche : Les circuits découverts présentent des topologies distinctes selon la tâche :
  • IOI : Dominé par un calcul MLP distribué ; l'attention est largement élaguée, sauf dans les dernières couches.
  • Pronoms de Genre (GP) : Hautement localisé et discontinu ; le calcul MLP est restreint à quelques couches, tandis que l'attention est parcimonieuse.
  • Greater Than (GT) : Présente une parcimonie extrême avec des blocs contigus de couches inactives et des motifs d'exécution de type "skip-like".

4.2 Fidélité et Cohérence

• Performance : Les circuits élagués maintiennent une haute précision de tâche (souvent >95 %) et une faible divergence KL par rapport au modèle complet.
• Cohérence : Les expériences sur cinq graines aléatoires montrent que les circuits découverts (métriques de fidélité et profils de parcimonie) sont hautement stables.
• Analyse Qualitative : Dans Llama 3.2-1B, les têtes d'attention survivantes dans la tâche IOI ont été identifiées comme remplissant des rôles fonctionnels spécifiques (IO-Movers, S-Inhibitors), reflétant les circuits découverts manuellement dans GPT-2, validant ainsi la capacité de la méthode à récupérer des intuitions mécanistes.

4.3 Efficacité de Calcul

• Mémoire : La méthode nécessite au moins 3× moins de mémoire que les lignes de base d'élagage d'arêtes (ex: 6,27 Go contre 72,79 Go pour EAP sur GPT-2).
• Temps : L'entraînement se termine en moins de 6 minutes, ce qui est nettement plus rapide que les méthodes d'élagage d'arêtes (environ 50 minutes), bien que légèrement plus lent que la méthode EAP à passage unique (qui sacrifie le temps pour une utilisation massive de la mémoire).

5. Signification et Revendications

L'article affirme que sa principale signification réside dans le dépassement des goulots d'étranglement de la scalabilité et de la granularité des méthodes actuelles. En passant de l'élagage d'arêtes à l'élagage de nœuds multi-granulaire, les auteurs démontrent que :

• De nombreux neurones conservés par les méthodes à granularité grossière sont inutiles et peuvent être supprimés avec un impact négligeable sur la performance.
• Les modèles plus grands (ex: Llama 8B) supportent des élagages aussi forts, voire plus forts, que les modèles plus petits, suggérant que l'augmentation de l'échelle introduit de la redondance plutôt que de nécessiter des circuits spécifiques à la tâche plus denses.
• La méthode offre une voie computationnellement réalisable pour découvrir des circuits causaux fins dans des modèles à grande échelle, révélant que différents comportements sont encodés par des motifs distincts d'activité d'attention, de MLP et de neurones.

Les auteurs reconnaissent une limite : le cadre identifie quels nœuds sont nécessaires mais ne récupère pas explicitement la structure d'interaction (arêtes) entre les nœuds survivants. Ils suggèrent que la combinaison de la sélection de nœuds avec la récupération d'arêtes constitue une direction future pour inférer les structures d'interaction.