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

1. Problemstellung

Das Verständnis der internen Mechanismen von Transformer-basierten Large Language Models (LLMs) ist entscheidend für deren zuverlässigen Einsatz. Ein zentrales Problem ist die Attribution (Zuordnung): Die Identifizierung, welche Eingabetokens oder interne Komponenten (wie Attention-Heads oder Neuronen) für eine spezifische Vorhersage verantwortlich sind.

Bestehende Methoden leiden unter erheblichen Nachteilen:

• Rechenkosten: Dichte Attribution (die Analyse aller Komponenten) ist oft prohibitiv teuer. Methoden wie Activation Patching erfordern separate Vorwärtsdurchläufe für jede Komponente, was bei langen Sequenzen und großen Modellen nicht skalierbar ist.
• Genauigkeit vs. Effizienz: Gradientenbasierte Methoden leiden unter Sättigungseffekten und Rauschen, während einfache Attention-Metriken oft unzuverlässig sind.
• Approximationsfehler: Viele Ansätze vernachlässigen wichtige nichtlineare Interaktionen oder ignorieren interne Komponenten.

Das Ziel ist es, eine Methode zu entwickeln, die sowohl hochverlässig (faithful) als auch extrem effizient ist, ohne auf kontrafaktische Beispiele (Counterfactuals) angewiesen zu sein.

2. Methodik: Dual Path Attribution (DPA)

Die Autoren stellen Dual Path Attribution (DPA) vor, ein neues Framework, das die Informationsfluss-Spur auf einem eingefrorenen Transformer in einem Vorwärts- und einem Rückwärtsdurchlauf nachzeichnet.

Kernkonzept

Anstatt jeden Eingabeteil vorwärts durch das Netz zu schicken, propagiert DPA den Unembedding-Vektor des Ziel-Tokens rückwärts durch das Netzwerk. Dabei wird für jede Position im Residual-Stream ein „effektives Ziel" (effective target) berechnet, das angibt, wie stark diese Position zur Vorhersage beiträgt.

Technische Details

DPA nutzt die spezifische Architektur von SwiGLU-Transformern (wie Llama, Mistral, Qwen), um den Berechnungsgraphen analytisch zu zerlegen. Der Schlüssel liegt in der Trennung in zwei Pfade:

1. Content-Pfad (Inhalt): Propagiert die Repräsentationen durch die Value- und Up-Projektionen (bei Attention) bzw. Up- und Down-Projektionen (bei GLU/MLP). Dieser Pfad ist linear, wenn die Aktivierungen eingefroren sind.
2. Control-Pfad (Steuerung): Steuert das Routing und die Transformation durch Attention-Scores und Gate-Werte. Da diese Pfade nichtlinear sind (z. B. Softmax, SiLU-Aktivierung), werden sie durch Linearisierung (basierend auf lokalen Gradienten oder Taylor-Approximationen) invertiert.

Der Algorithmus:

• Vorwärtsdurchlauf: Das Modell verarbeitet die Eingabe einmal. Alle notwendigen Aktivierungen (Attention-Weights, Gate-Werte, Normalisierungsparameter) werden zwischengespeichert.
• Rückwärtsdurchlauf (Target Propagation):
  • Startend mit dem Unembedding-Vektor des Ziel-Tokens wird dieser rekursiv durch die Schichten nach oben (Layer $L$ zu $0$) propagiert.
  • Für GLU-Module wird der Gradient in einen Up-Pfad (Inhalt) und einen Gate-Pfad (Steuerung) aufgeteilt und gewichtet summiert.
  • Für Multi-Head Self-Attention (MHSA) wird der Gradient in Value-, Query- und Key-Pfade zerlegt. Die Softmax-Nichtlinearität wird durch eine Taylor-basierte Zerlegung stabil approximiert.
• Attributions-Score: Am Ende wird der Score für jede Komponente durch das Skalarprodukt ihrer gespeicherten Aktivierung mit dem propagierten Zielvektor berechnet.

Komplexität

Ein entscheidender Vorteil ist die Zeitkomplexität. DPA hat eine Komplexität von $O(1)$ bezüglich der Anzahl der zu analysierenden Komponenten. Im Gegensatz dazu skalieren Methoden wie Activation Patching linear ($O(M)$) mit der Anzahl der Komponenten $M$.

3. Hauptbeiträge

1. Neues Framework: Einführung von DPA, das logit-basierte Attribution in einem einzigen Vorwärts- und Rückwärtsdurchlauf ermöglicht, ohne kontrafaktische Eingaben zu benötigen.
2. Architekturspezifische Zerlegung: Demonstration, dass der Berechnungsgraph von SwiGLU-Transformern in diskrete Control- und Content-Pfade zerlegt werden kann, was eine präzise Invertierung erlaubt.
3. Skalierbarkeit: Die Methode ermöglicht dichte Attribution (Analyse jedes Neurons und Heads) mit konstantem Rechenaufwand, unabhängig von der Modellgröße oder der Anzahl der Komponenten.
4. Empirische Validierung: Umfassende Experimente zeigen, dass DPA den State-of-the-Art in Bezug auf Zuverlässigkeit und Effizienz erreicht.

4. Ergebnisse

Die Evaluation erfolgte auf Standard-Benchmarks (Known 1000, SQuAD v2.0, IMDb, Indirect Object Identification) mit verschiedenen Modellen (Llama-3, Llama-2, Qwen, Mistral).

• Zuverlässigkeit (Faithfulness): DPA übertrifft alle Baselines (einschließlich Gradientenmethoden, Attention-Rollout, DePass und IFR) in den Metriken Comprehensiveness (Störung bei Entfernung wichtiger Teile) und Sufficiency (Wiederherstellung der Vorhersage bei Beibehaltung).
  • Auf dem Known-1000-Dataset erreichte DPA eine Gesamt-Zuverlässigkeit von 28,09 (vs. 16,28 für DePass).
  • Bei der Komponenten-Attribution (IOI) zeigte DPA eine nahezu perfekte Wiederherstellung (Recovery > 120%), was auf die präzise Identifizierung kausaler Schaltkreise hindeutet.
• Effizienz:
  • Geschwindigkeit: DPA ist massiv schneller. Auf dem SQuAD-Dataset benötigte DPA nur 129 Sekunden, während Activation Patching und DePass jeweils über 1,5 Stunden (ca. 5400 Sekunden) benötigten. Das entspricht einem 40-fachen Geschwindigkeitsvorteil.
  • Speicher: Der Peak-Speicherverbrauch ist mit den Baselines vergleichbar, da DPA keine zusätzlichen Vorwärtsdurchläufe benötigt, die den Speicher durch wiederholte Berechnungen belasten würden.
• Sensitivitätsanalyse: Die Ergebnisse zeigen, dass für eine optimale Token-Attribution eine Balance zwischen Content- und Control-Pfaden nötig ist (ca. 50-70% Content), während für die feingranulare Komponenten-Attribution eine stärkere Gewichtung der Steuerungsmechanismen (Gates, Query/Keys) erforderlich ist.

5. Bedeutung und Fazit

DPA schließt die Lücke zwischen mechanistischer Interpretierbarkeit und praktischer Anwendbarkeit.

• Praktische Relevanz: Durch die $O(1)$-Komplexität wird es erstmals möglich, dichte Attributionen auf großen Modellen und langen Kontexten in akzeptabler Zeit durchzuführen. Dies ist ein Durchbruch für die Analyse von „Black-Box"-Modellen.
• Mechanistische Einsichten: Die Methode erlaubt es, kausale Mechanismen bis hin zu einzelnen Attention-Heads und MLP-Neuronen präzise zu isolieren, was für das Debuggen und Verstehen von LLMs unerlässlich ist.
• Einschränkungen: Die Methode ist derzeit spezifisch für SwiGLU-Architekturen optimiert. Die Anpassung an andere Architekturen (z. B. Mixture-of-Experts oder State-Space-Modelle) erfordert manuelle Herleitungen der inversen Transformationen. Zudem basiert sie auf Linearisierungen, was in stark nichtlinearen Regimen zu Approximationsfehlern führen kann.

Zusammenfassend bietet DPA ein robustes, effizientes und hochgenaues Werkzeug, um die internen Denkwege moderner Large Language Models zu entschlüsseln.