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

Probleemstelling

Het begrijpen van de interne mechanismen van transformer-gebaseerde Large Language Models (LLMs) is cruciaal voor hun betrouwbare inzetbaarheid. Een centrale uitdaging is attributie: het identificeren van welke invoertokens of interne componenten (zoals attention heads of neuronen) verantwoordelijk zijn voor een specifieke voorspelling van het model.

Bestaande methoden kampen met ernstige beperkingen:

• Berekeningskosten: Methoden zoals Activation Patching vereisen een aparte forward-pass voor elke component, wat bij dichte attributie (voor alle componenten) prohibitief duur is ($O(M)$ complexiteit, waarbij $M$ het aantal componenten is).
• Onnauwkeurigheid: Gradient-based methoden lijden vaak onder saturatie en ruis, terwijl attention-based methoden geen betrouwbare indicatoren zijn voor causaliteit.
• Benaderingsfouten: Bestaande decompositiemethoden negeren vaak belangrijke niet-lineaire interacties of vereisen contrafactuele voorbeelden.

Er is dus een dringende behoefte aan een methode die zowel trouw (faithfulness) als rekenkundige efficiëntie biedt, zonder contrafactuele voorbeelden te gebruiken.

Methodologie: Dual Path Attribution (DPA)

De auteurs introduceren Dual Path Attribution (DPA), een decompositie-gebaseerd framework dat informatieflow traceert door een bevroren transformer in slechts één forward-pass en één backward-pass.

Kernprincipes:

1. Doelgerichte Propagatie (Target Propagation): In plaats van elke component naar voren te laten passeren, wordt het unembedding vector van het doel-token (de target) achterwaarts door het netwerk gepropageerd. Hierdoor wordt voor elke residuele positie een "effectief doel" ($t^{(0)}$) berekend.
2. Decompositie in SwiGLU: De methode maakt gebruik van de bilineaire structuur van SwiGLU-architecturen (zoals LLaMA, Mistral, Qwen). De auteurs ontleden de computationele graaf in twee distincte paden:
   • Content-paden: Propageren representaties via de value en up projecties (informatiedragend).
   • Control-paden: Beheren de routing en transformatie via attention scores en gate scores (beslissend over hoe informatie wordt verwerkt).
3. Linearisatie en Inversie: Om de backward-propagatie mogelijk te maken, worden de niet-lineaire operaties (zoals SiLU/Swish en Softmax) lokaal gelineariseerd.
   • Bij GLU (Gated Linear Units) wordt de inverse transformatie benaderd als een gewogen som van het up-projection pad en het gate-projection pad.
   • Bij MHSA (Multi-Head Self-Attention) wordt de inverse opgesplitst in bijdragen van de query, key en value paden, waarbij de softmax-activatie wordt benaderd via een Taylor-ontwikkeling voor stabiliteit.
4. Efficiëntie: Omdat de attributie wordt berekend door een enkele vector terug te laten lopen door het netwerk en te dot-producten met de gecachede activaties, is de tijdscomplexiteit ten opzichte van het aantal componenten $O(1)$. Dit maakt dichte attributie mogelijk voor lange sequenties.

Belangrijkste Bijdragen

1. Nieuw Framework: Introductie van DPA, een decompositie-gebaseerd framework dat logit-scores toewijst zonder contrafactuele voorbeelden en met een enkele forward/backward pass.
2. Theoretische Decompositie: Het aantonen dat de computationele graaf van SwiGLU-Transformers kan worden ontbonden in specifieke control en content paden, wat een nauwkeurigere analyse van niet-lineaire interacties mogelijk maakt dan eerdere lineaire benaderingen.
3. Unprecedented Efficiëntie: Het bereiken van $O(1)$ complexiteit voor dichte attributie, wat een enorme snelheidswinst oplevert ten opzichte van methoden zoals Activation Patching.
4. Empirisch Bewijs: Uitgebreide experimenten die aantonen dat DPA state-of-the-art resultaten bereikt op zowel invoer- als component-niveau attributie.

Resultaten

De auteurs hebben DPA getest op diverse modellen (Llama-3.1-8B, Llama-2-7B, Qwen3-4B, Mistral-7B, Qwen2.5-32B) en datasets (Known 1000, SQuAD v2.0, IMDb, Indirect Object Identification).

• Truwigheid (Faithfulness): DPA presteert consistent beter dan bestaande baselines (zoals Integrated Gradients, Attention Rollout, en DePass) op de metrics Comprehensiveness (disruptie bij verwijderen van belangrijke tokens) en Sufficiency (herstel van prestaties bij behoud).
  • Op de Known 1000 en SQuAD datasets behaalde DPA de hoogste totale scores.
  • Bij Component Attributie (zoals het vinden van circuits voor indirecte objecten in IOI) isoleerde DPA de causale mechanismen met een disruptie van 0.00 en een zeer snelle herstelcurve, wat aangeeft dat het de belangrijkste componenten zeer vroeg in de ranglijst plaatst.
• Efficiëntie:
  • DPA is 40x sneller dan Activation Patching en DePass op lange contexttaken (bijv. SQuAD v2.0).
  • Waar AP en DePass enkele uren nodig hadden voor een taak, deed DPA dit in enkele minuten, terwijl het geheugengebruik vergelijkbaar bleef met de baselines.
• Sensitiviteit: Een analyse van de hyperparameters ($\mu$) toont aan dat een evenwicht tussen content- en control-paden (ongeveer 50-50) de meest stabiele resultaten geeft, hoewel fijne componentattributie meer afhankelijk is van de control-paden (gates en keys).

Betekenis en Impact

DPA is een doorbraak in mechanistische interpretatie van LLMs. Het lost het fundamentele dilemma op tussen trouw (nauwkeurige causaliteit) en schaalbaarheid.

• Het maakt het mogelijk om dichte attributie uit te voeren op grote modellen en lange contexten, wat eerder onhaalbaar was vanwege de rekenkosten.
• Het biedt een robuust hulpmiddel voor onderzoekers om de interne redeneerpaden van moderne LLMs te ontrafelen, tot op het niveau van individuele attention heads en MLP-neuronen.
• Hoewel de methode specifiek is geoptimaliseerd voor SwiGLU-architecturen, is de onderliggende filosofie van doelgerichte propagatie architectonisch agnostisch, wat de weg vrijmaakt voor toepassing op andere modellen (zoals MoE of State Space Models) na het afleiden van de juiste inverse functies.

Kortom, DPA biedt een praktische, schaalbare en uiterst nauwkeurige oplossing voor het begrijpen van "waarom" een LLM een bepaald antwoord geeft.