Dual Path Attribution: Efficient Attribution for SwiGLU-Transformers through Layer-Wise Target Propagation

1. Problématique

L'interprétabilité des grands modèles de langage (LLM) basés sur l'architecture Transformer est cruciale pour leur déploiement fiable. Le défi central réside dans l'attribution : identifier quels tokens d'entrée ou quels composants internes (têtes d'attention, neurones) sont responsables d'une prédiction spécifique.

Les méthodes existantes souffrent de limitations majeures :

• Coût computationnel prohibitif : Les méthodes d'intervention causale (comme le Activation Patching) nécessitent un passage avant (forward pass) séparé pour chaque composant à attribuer. Pour une attribution dense sur de longues séquences, cela devient ingérable.
• Approximations imprécises : Les méthodes basées sur les gradients (ex: Integrated Gradients) ou les poids d'attention souffrent souvent de saturation, de bruit ou d'erreurs d'approximation, manquant de fidélité (faithfulness).
• Ignorance des interactions complexes : Beaucoup de méthodes ne capturent pas correctement les effets indirects ou les interactions non linéaires au sein des modules du modèle.

2. Méthodologie : Dual Path Attribution (DPA)

Les auteurs proposent DPA, un cadre novateur qui permet une attribution fidèle et efficace en un seul passage avant et un seul passage arrière, sans nécessiter d'exemples contrefactuels.

Principe Fondamental : Propagation Cible Vers le Haut

Au lieu de faire passer chaque composant vers l'avant pour voir son impact, DPA inverse la perspective :

1. On part du vecteur de désembedding (unembedding vector) du token cible (la prédiction souhaitée).
2. On propage ce vecteur récursivement vers l'arrière à travers le Transformer gelé.
3. À chaque position résiduelle, on calcule un "cible effective" qui représente l'impact nécessaire de ce niveau pour atteindre le résultat final.

Décomposition du SwiGLU en Deux Chemins

La méthode exploite la structure bilinéaire spécifique des Transformers SwiGLU (utilisés par Llama, Mistral, Qwen, etc.). Elle décompose mathématiquement le graphe de calcul en deux chemins distincts :

1. Le Chemin de Contenu (Content Pathway) :
   • Transmet les représentations via les projections Up et Value.
   • Il s'agit d'un chemin linéaire (une fois les activations gelées).
2. Le Chemin de Contrôle (Control Pathway) :
   • Gère le routage et la transformation via les scores d'attention et les portes (gates).
   • Nécessite une linéarisation locale pour être inversé (ex: approximation du gradient de la fonction d'activation SiLU/Swish et de la fonction Softmax).

Algorithme de Propagation

Pour chaque couche $l$, la cible effective $t^{(l)}$ est mise à jour en inversant les modules MHSA (Self-Attention) et GLU :

• GLU : La cible est décomposée en contributions pondérées des chemins Up (contenu) et Gate (contrôle).
• MHSA : La cible est décomposée en contributions des chemins Value, Query et Key. La propagation à travers le Softmax est gérée par une décomposition de Taylor pour assurer la stabilité.

Ces chemins sont combinés avec des paramètres de pondération ($\mu$) pour équilibrer l'influence du contenu et du routage.

3. Contributions Clés

1. Efficacité Inédite ($O(1)$) : DPA atteint une complexité temporelle constante par rapport au nombre de composants ($M$). Contrairement aux méthodes d'intervention qui sont en $O(M)$, DPA ne nécessite aucun passage avant supplémentaire après le passage initial de cache. Cela permet une attribution dense même sur des contextes très longs.
2. Décomposition Analytique : L'article fournit une dérivation mathématique rigoureuse pour inverser les transformations bilinéaires des architectures SwiGLU, séparant explicitement les flux d'information (contenu) et de décision (contrôle).
3. Fidélité Supérieure : En évitant les approximations grossières des gradients et en capturant les effets directs et indirects via la propagation cible, DPA offre une localisation plus précise des mécanismes causaux.

4. Résultats Expérimentaux

Les auteurs ont évalué DPA sur plusieurs modèles (Llama-3.1-8B, Llama-2-7B, Qwen3-4B, Mistral-7B) et benchmarks (Known 1000, SQuAD v2.0, IMDb, Indirect Object Identification).

• Fidélité (Attribution d'Entrée) :
  • Sur les tâches de complétion factuelle (Known 1000) et de compréhension (SQuAD), DPA surpasse systématiquement les méthodes de base (gradients, attention, DePass) en termes de score total de fidélité (combinaison de disruption et de recovery).
  • DPA identifie mieux les tokens critiques nécessaires à la prédiction.
• Attribution de Composants (Interne) :
  • DPA excelle dans l'identification des têtes d'attention et des neurones MLP responsables de circuits spécifiques (ex: tâche IOI).
  • Il atteint un score de disruption proche de zéro et une recovery très élevée, surpassant les méthodes comme Attn-only ou Gradient.
• Performance Computationnelle :
  • Vitesse : DPA est environ 40 fois plus rapide que l'Activation Patching et DePass sur des tâches de contexte long (SQuAD). Là où les méthodes concurrentes prennent ~1,5 heure, DPA termine en ~2 minutes.
  • Mémoire : L'empreinte mémoire de pointe est comparable à celle des méthodes de base, car elle repose sur le cache des activations d'un seul passage avant.

5. Signification et Limites

Signification :
DPA comble le fossé historique entre la fidélité (précision de l'explication) et l'efficacité computationnelle. Il permet d'analyser des mécanismes causaux complexes dans des modèles de grande taille avec une densité d'attribution impossible à réaliser avec les méthodes d'intervention actuelles. C'est un outil pratique pour la recherche en interprétabilité mécaniste à grande échelle.

Limites :

• Spécificité Architecturale : La décomposition mathématique est optimisée pour les Transformers denses basés sur SwiGLU. L'adaptation aux architectures à routage discret (MoE) ou aux modèles sans attention (State Space Models) nécessite de dériver de nouvelles fonctions inverses.
• Approximation Linéaire : La méthode repose sur une linéarisation locale (développement de Taylor d'ordre 1). Elle peut perdre en précision dans des régimes de forte courbure ou lorsque les interactions non linéaires sont fortement couplées.
• Portée Logit : DPA attribue directement aux logits cibles et non aux probabilités post-softmax, ce qui peut rendre difficile la capture de mécanismes où la probabilité d'un token augmente par la suppression active de tokens concurrents.

En conclusion, Dual Path Attribution représente une avancée majeure pour l'analyse interne des LLM modernes, offrant un compromis optimal entre précision théorique et viabilité pratique.