Technische Samenvatting: Het corrigeren van gradiëntgebaseerde circuitlokalisatie via interactiebewuste backpropagation

Probleemstelling

Circuitlokalisatie heeft als doel de specifieke subset van modelcomponenten (bijv. attention heads, neuronen, verbindingen) te identificeren die verantwoordelijk zijn voor specifieke gedragingen in Large Language Models (LLMs) om mechanistische interpreteerbaarheid mogelijk te maken. De meeste bestaande methoden gaan uit van de aanname dat componenten onafhankelijk opereren, waarbij de belangrijkheid wordt geschat door elke component in isolatie te verstoren.

Echter, componenten in neurale netwerken interageren, en het negeren van deze interacties leidt tot een systematische miscalculatie van de belangrijkheid van componenten. De auteurs identificeren een specifieke, problematische interactie genaamd attention self-repair (zelfherstel van aandacht). Bij dit fenomeen, wanneer meerdere posities grote attention weights bezitten en vergelijkbare informatie (value vectors) bevatten, heeft het ableren van een enkele attention score minimaal effect op de output. Dit gebeurt omdat de softmax-activatiefunctie de gewichten herverdeelt naar de resterende redundante posities om te compenseren. Bijgevolg onderschatten standaard perturbatiegebaseerde en gradiëntgebaseerde verklarende methoden de belangrijkheid van componenten die bijdragen aan deze attention scores, wat vaak resulteert in verdwijnende gradiënten voor invloedrijke scores.

Methodologie: Gradient Interaction Modifications (GIM)

Het artikel introduceert Gradient Interaction Modifications (GIM), een gradiëntgebaseerde attributiemethode die is ontworpen om de impact van het gelijktijdig verstoren van meerdere interagerende componenten te schatten. GIM modificeert standaard backpropagation via drie complementaire technieken:

1. Temperature-adjusted Softmax Gradients (TSG):

   • Motivatie: Het aanpakken van attention self-repair, waarbij standaard gradiënten verdwijnen omdat de winst van het verhogen van één softmax-gewicht wordt gecompenseerd door het verlies van het onderdrukken van anderen wanneer de waarden vergelijkbaar zijn.
   • Mechanisme: Tijdens backpropagation berekent TSG de softmax opnieuw met een verhoogde temperatuur ($T > 1$) voordat de gradiënten worden berekend. Dit vlakt de waarschijnlijkheidsverdeling af, waardoor grote gewichten krimpen en kleine gewichten groeien. Door de set van "significante" posities te vergroten, wordt de aanname dat bijdragen identiek zijn over deze posities versoepeld. Dit zorgt ervoor dat de gradiënten ($g_j - g_k$) niet verdwijnen, waardoor het gradiëntsignaal kan overleven en de effect van gelijktijdige perturbatie kan benaderen.

2. Layernorm Freeze:

   • Motivatie: Layer normalization introduceert een zero-sum dynamiek waarbij het ableren van één component de activatiemagnitude vermindert, waardoor de normalisatiefactor de resterende componenten versterkt (een vorm van self-repair).
   • Mechanisme: De normalisatiefactor ($\sigma$) wordt tijdens backpropagation als een constante behandeld. Dit meet de directe bijdrage van elke component voordat de normalisatie-rescaling plaatsvindt, wat voorkomt dat de belangrijkheid wordt onderschat door compensatoire amplificatie.

3. Gradient Norm:

   • Motivatie: Multiplicatieve operaties (bijv. $A = QK^T$ waarbij zowel $Q$ als $K$ afhangen van input $X$) leiden tot dubbeltellingen in standaard eerste-orde gradiënten. Het verstoren van $X$ beïnvloedt zowel $Q$ als $K simultaan, maar de standaard kettingregels tellen hun individuele effecten op, waardoor de verandering wordt overschat.
   • Mechanisme: De gradiënten van multiplicatieve interacties worden gedeeld door het aantal multiplicatieve inputs (bijv. delen door 2 voor $Q \times K$). Dit komt overeen met de "uniform rule" en corrigeert de overtelling van bijdragen in attention-value multiplicatie, query-key multiplicatie en MLP gate-projection multiplicatie.

Belangrijkste Bijdragen

1. GIM Framework: De introductie van een gradiëntgebaseerde attributiemethode die expliciet rekening houdt met feature-interacties, waarmee state-of-the-art prestaties wordt behaald op de circuit localization track van de Mechanistic Interpretability Benchmark (MIB).
2. Identificatie van Attention Self-Repair: De formele identificatie en bewijsvoering dat attention self-repair (een OR-gate-achtige interactie) ervoor zorgt dat gradiënten verdwijnen voor invloedrijke attention scores, wat leidt tot een systematische onderschatting van de belangrijkheid van componenten.
3. Superioriteit in Feature Attribution: Demonstratie dat GIM bestaande gradiëntgebaseerde methoden overtreft in feature attribution over diverse taken en LLM-architecturen.

Resultaten

• Circuit Localization: Op de MIB presteert GIM significant beter dan bestaande methoden (inclusief EAP, EAP-IG, Hybrid-ens, en IFR) over meerdere modellen (GPT-2, Qwen-2.5, Gemma-2, Llama-3.1) en taken (Indirect Object Identification, Rekenen, MCQA, ARC). Het bereikt de hoogste gemiddelde Circuit Performance Ratio (CPR), met name op taken die stapsgewijze redenering vereisen.
• Feature Attribution: Over zeven instructie-getunede LLM's en zes datasets (dekking van QA, feitelijke verificatie, sentiment en haatzaaiende taal detectie), behaalt GIM over het algemeen hogere scores voor comprehensiveness en sufficiency vergeleken met GradientXInput, Integrated Gradients, DeepLIFT en TransformerLRP.
• Validatie van Self-Repair: Empirische analyse bevestigt dat attention self-repair wijdverbreid is in echte LLM's, waarbij het gemiddeld 65–1200 keer per input voorkomt. Het gezamenlijk ableren van twee grote attention scores levert een aanzienlijk groter effect op dan de som van hun individuele ablations, wat de theoretische analyse valideert.
• Ablatie-studie: Het verwijderen van een enkele modificatie (TSG, Layernorm Freeze, of Gradient Norm) uit GIM vermindert de prestaties over de meeste model-dataset paren, wat bevestigt dat elk onderdeel betekenisvol bijdraagt aan de getrouwheid van de methode.

Betekenis

Het artikel claimt dat GIM een getrouwere mechanistische analyse van LLM's mogelijk maakt door de systematische onderschatting van de belangrijkheid van componenten te corrigeren die wordt veroorzaakt door feature-interacties. Door expliciet te modelleren hoe componenten interageren—met name via softmax-redistributie en normalisatie-rescaling—biedt GIM een nauwkeuriger beeld van de interne mechanismen die het gedrag van het model drijven. De auteurs positioneren GIM niet slechts als een incrementele verbetering, maar als een noodzakelijke correctie op de "onafhankelijkheidsaanname" die gangbaar is in huidige interpreteerbaarheidstools, waardoor onderzoekers circuits kunnen identificeren die voorheen werden verborgen door verdwijnende gradiënten. De methode wordt vrijgegeven als een open-source Python-package om bredere adoptie te faciliteren.