Technische Samenvatting: Multi-Granulaire Node Pruning voor Causal Circuit Discovery

1. Probleemstelling

Causal circuit discovery heeft als doel om minimale sub-netwerken binnen Large Language Models (LLM's) te identificeren die specifiek gedrag causaal aansturen. Hoewel bestaande methoden de mechanistische interpreteerbaarheid hebben gevorderd, kampen ze met twee primaire beperkingen:

1. Computationele Schaalbaarheid: Huidige benaderingen vertrouwen voornamelijk op edge pruning (het verwijderen van verbindingen tussen componenten). Aangezien het aantal edges in een transformer kwadratisch groeit met het aantal nodes ($O(|N|^2)$), wordt edge-level pruning computationeel onuitvoerbaar naarmate de modelgrootte toeneemt. Bovendien vereisen deze methoden vaak het opslaan van gedisentangledged residual-stream representaties en intermediaire activaties, wat leidt tot een verbijsterend hoog geheugengebruik.
2. Granulariteitsbeperkingen: Bestaande methoden opereren doorgaans op een grove granulariteit, waarbij volledige attention heads of MLP-blocks als atomaire eenheden worden behandeld. Dit negeert de mogelijkheid dat specifieke neuronen (nodes) binnen deze componenten verantwoordelijk zijn voor doelgerichte gedragingen, wat resulteert in herstelde circuits die fijne structuren missen en onnodige componenten behouden.

2. Methodologie

De auteurs stellen een multi-granulair node-pruning framework voor dat minimale circuits ontdekt over meerdere niveaus van granulariteit binnen één enkele verenigde optimalisatiestap.

2.1 Hiërarchisch Granulariteitsontwerp

In plaats van edges tussen grove componenten te prunen, leert het framework masks over nodes op vier verschillende niveaus:

• Block Level: Volledige Attention Blocks en MLP Blocks.
• Head Level: Individuele Attention Heads.
• Neuron Level: Individuele neuronen binnen Attention lagen en de twee lineaire lagen van MLP blocks (MLP1/Keys en MLP2/Values).

2.2 Two-Stream Causal Pruning Setup

Om causale belangrijkheid te identificeren zonder distributieverschuivingen te introduceren (een risico bij eenvoudige zero-ablatie), gebruikt de methode een two-stream forward pass:

• Clean Stream: Het model verwerkt een input die ontworpen is om een specifiek gedrag op te roepen (bijv. het voorspellen van een jaar > 50).
• Corrupted Stream: Het model verwerkt een minimaal geperturbeerde versie van de input (bijv. het inverteren van gender-voornaamwoorden of het veranderen van het startjaar) om een ander voorspelling op te roepen.
• Interpolatie: Voor elke node $n_i$ interpoleert een leerbaar mask $m_i \in [0, 1]$ tussen de clean activatie ($h^{clean}$) en de corrupted activatie ($h^{corrupted}$):
  $$h^{node}_i = m_i \cdot h^{clean}_i + (1 - m_i) \cdot h^{corrupted}_i$$
  De optimalisatie streeft ernaar om masks te vinden die het model in staat stellen de correcte output te produceren (overeenkomend met het clean gedrag) terwijl de sparsity wordt gemaximaliseerd.

2.3 Optimalisatie en Mask Parameterisatie

• Hard Concrete Distribution: Masks worden geparametriseerd met de Hard Concrete distributie om gradiënt-gebaseerde optimalisatie mogelijk te maken terwijl een ongeveer binaire gedraging (0 of 1) wordt geïnduceerd.
• Unified Objective: De trainingsdoelstelling combineert drie componenten:
  1. Faithfulness Loss ($L_{faith}$): Een KL-divergentie die ervoor zorgt dat de geprunede circuit output-distributie overeenkomt met het basismodel.
  2. Task Loss ($L_{task}$): Een margin-based ranking loss die ervoor zorgt dat de circuit correct onderscheid maakt tussen ground-truth en distractor opties.
  3. Sparsity Regularizer ($L_0$): Moedigt de masks aan om nul te zijn, waardoor de circuit grootte wordt geminimaliseerd.
• Hiërarchische Consistentie: Na de training worden de masks gebinariseerd met een constraint die garandeert dat als een parent node (bijv. een MLP block) wordt gepruned, alle child nodes (bijv. neuronen binnen dat block) ook worden gedeactiveerd.

3. Belangrijkste Bijdragen

1. Multi-Granulair Framework: De introductie van een verenigd optimalisatieframework dat gelijktijdig masks leert over transformer blocks, attention heads, MLP blocks en individuele neuronen.
2. Schaalbaarheid en Efficiëntie: De methode vermijdt de kwadratische schaling van edge pruning en de noodzaak om intermediaire activaties voor alle edges op te slaan. Het vermindert het GPU-geheugengebruik met 3× tot 11× vergeleken met eerdere methoden, wat circuit discovery op Llama 3.1-8B mogelijk maakt met een enkele 40GB GPU.
3. Fijnmazige Ontdekking: De aanpak identificeert aanzienlijk compactere circuits dan baselines, waarbij specifiek een hoge sparsity wordt bereikt op het MLP block en neuron niveau, waar voorheen de meeste methoden falen.

4. Experimentele Resultaten

De methode is geëvalueerd op drie standaard circuit discovery taken (Indirect Object Identification, Gendered Pronouns, en Greater Than) met behulp van GPT-2, GPT-2-XL en Llama modellen.

4.1 Sparsity en Compactheid

• MLP en Neuron Pruning: De voorgestelde methode presteert significant beter dan edge-pruning baselines (EAP, EP) en ongestructureerde weight pruning (DiscoGP). In de minst gunstige setting pruneerde het 33,34% meer MLP blocks en 59,8% meer neuronen dan voorgaande methoden.
• Taak-afhankelijke Structuren: De ontdekte circuits vertonen verschillende topologieën op basis van de taak:
  • IOI: Gedomineerd door gedistribueerde MLP computation; attention wordt grotendeels gepruned, behalve in de laatste lagen.
  • Gendered Pronouns (GP): Sterk gelokaliseerd en discontinu; MLP computation is beperkt tot enkele lagen, terwijl attention sparse is.
  • Greater Than (GT): Vertoont extreme sparsity met aaneengesloten blokken van inactieve lagen en "skip-like" executiepatronen.

4.2 Fidelity en Consistentie

• Prestaties: De geprunede circuits behouden een hoge taakaccuratesse (vaak >95%) en een lage KL-divergentie ten opzichte van het volledige model.
• Consistentie: Experimenten over vijf random seeds tonen aan dat de ontdekte circuits (fidelity metrics en sparsity profielen) zeer stabiel zijn.
• Kwalitatieve Analyse: In Llama 3.2-1B werden de overlevende attention heads in de IOI taak gevonden die specifieke functionele rollen vervullen (IO-Movers, S-Inhibitors), wat de handmatige ontdekte circuits in GPT-2 weerspiegelt en de capaciteit van de methode valideert om mechanistische inzichten te herstellen.

4.3 Computationele Efficiëntie

• Geheugen: De methode vereist minstens 3× minder geheugen dan edge-pruning baselines (bijv. 6,27 GB vs. 72,79 GB voor EAP op GPT-2).
• Tijd: Training voltooit in minder dan 6 minuten, wat aanzienlijk sneller is dan edge-pruning methoden (ongeveer 50 minuten), hoewel het iets langzamer is dan de single-pass EAP methode (die tijd inruilt voor massaal geheugengebruik).

5. Betekenis en Claims

Het artikel claimt dat de primaire betekenis ligt in het overwinnen van de schaalbaarheid en granulariteits-bottlenecks van huidige circuit discovery. Door over te stappen van edge pruning naar multi-granulair node pruning, demonstreren de auteurs dat:

• Veel neuronen die door grove methoden worden behouden, onnodig zijn en kunnen worden verwijderd met een verwaarloosbare impact op de prestaties.
• Grotere modellen (bijv. Llama 8B) even sterke of zelfs sterkere pruning ondersteunen dan kleinere modellen, wat suggereert dat verhoogde schaal redundantie introduceert in plaats van dat het dichtere taak-specifieke circuits vereist.
• De methode een computationeel haalbaar pad biedt voor het ontdekken van fijnmazige causale circuits in grootschalige modellen, waarbij wordt onthuld dat verschillende gedragingen worden gecodeerd door onderscheidende patronen van attention, MLP en neuron activiteit.

De auteurs erkennen een beperking: het framework identificeert welke nodes noodzakelijk zijn, maar herstelt niet expliciet de interactiestructuur (edges) tussen de overlevende nodes. Zij suggereren een toekomstige richting waarbij node selectie wordt gecombineerd met edge recovery om interactiestructuren af te leiden.