Efficient Discovery of Approximate Causal Abstractions via Neural Mechanism Sparsification

Dit artikel introduceert een efficiënte methode om interpreteerbare causale abstracties van getrainde neurale netwerken te ontdekken door structurele pruning te gebruiken als zoekstrategie, waarbij een nieuwe interventierisico-objectieve functie wordt afgeleid om eenheden te vervangen of te samenvoegen zonder opnieuw te hoeven trainen.

De kern

De Kern: Het Ontmaskeren van het "Gedachteproces" van een AI

Stel je voor dat een kunstmatige intelligentie (zoals een AI die foto's herkent) een enorme, ingewikkelde fabriek is. Deze fabriek heeft duizenden werknemers (de neuronen in het netwerk) die allemaal samenwerken om een product (een voorspelling) te maken. We weten dat de fabriek goed werkt, maar we hebben geen idee hoe precies. Welke werknemers doen echt iets belangrijks, en wie loopt alleen maar mee in de stroom?

De onderzoekers van dit paper willen deze fabriek niet alleen begrijpen, maar ook versimpelen. Ze willen een kleinere, overzichtelijkere versie van de fabriek bouwen die precies hetzelfde doet, maar zonder de overbodige werknemers. Dit noemen ze een "causale abstractie".

Het Probleem: Waarom is dit zo moeilijk?

Normaal gesproken proberen mensen AI-modellen te verkleinen door simpelweg te kijken wie er "stil" is (wie weinig doet). Dat is alsof je een werknemer ontslaat omdat hij vandaag niet hard heeft gelopen. Maar dat is gevaarlijk! Misschien is die werknemer juist cruciaal op momenten dat de machine een moeilijke beslissing moet nemen.

De oude methoden om te testen of een versimpelde machine nog goed werkt, zijn extreem duur en traag. Het is alsof je elke werknemer een voor een moet testen door de fabriek volledig stil te leggen, een nieuwe opdracht te geven en te kijken wat er gebeurt. Dat kan je niet elke dag doen.

De Oplossing: Een Slimme "Wat-zou-er-gebeuren"-Rekenmachine

De onderzoekers hebben een slimme truc bedacht. In plaats van de hele fabriek te testen, kijken ze naar de kracht en de rol van elke werknemer.

1. De "Wat-zou-er-gebeuren"-test (Interventies):
   Stel je voor dat je een werknemer vasthoudt en zegt: "Jij doet vandaag niets, je staat stil." Of: "Jij doet precies wat je buurman doet."

   • Als de fabriek daarna nog steeds perfect werkt, was die werknemer misschien niet zo belangrijk.
   • Als de fabriek in paniek raakt en fouten maakt, was die werknemer essentieel.

2. De Slimme Rekentruc (De Tweede Orde):
   Het probleem is dat je dit niet voor elke werknemer echt kunt doen (te langzaam). De onderzoekers hebben een wiskundige formule bedacht die voorspelt wat er zou gebeuren zonder het echt te doen.

   • Ze kijken naar twee dingen:
     • Hoeveel "kracht" (gradient) heeft de werknemer? (Hoe hard duwt hij?)
     • Hoeveel "stijfheid" (curvature) heeft de werknemer? (Hoeveel doet zijn werk tellen als hij verandert?)
   • Door deze twee te combineren, krijgen ze een score. Een lage score betekent: "Deze werknemer kan worden vervangen door een simpele robot of zelfs weggegooid zonder dat de fabriek crasht."

De Grote Ontdekking: Waarom de oude methoden faalden

Een van de belangrijkste ontdekkingen in dit paper is een waarschuwing voor een populaire oude methode.

• De oude methode (Variance-based): Deze keek alleen naar wie er het meest bewoog. "Wie beweegt het meest, is het belangrijkst."
• Het probleem: Stel je voor dat je een werknemer hebt die heel stil is, maar die een heel zware machine bedient. Als hij stopt, valt de hele fabriek stil. De oude methode zou zeggen: "Die beweegt niet, dus weg met hem!" en zou de fabriek kapot maken.
• De nieuwe methode: Deze kijkt niet alleen naar beweging, maar ook naar de impact. Ze houden rekening met hoe de machine is "geschaald" (vermenigvuldigd). Het is alsof ze kijken naar de echte waarde van de werknemer, niet alleen naar hoe hard hij zwaait.

Wat is het Resultaat?

De onderzoekers hebben een proces ontwikkeld dat:

1. Snel is: Het rekent in één keer door welke werknemers overbodig zijn.
2. Betrouwbaar is: Het maakt een kleinere versie van de AI die echt begrijpt waarom hij iets doet, niet alleen wat hij doet.
3. Robuust is: Zelfs als je de interne instellingen van de AI verandert (zoals de grootte van de werknemers), blijft de nieuwe methode dezelfde belangrijke werknemers kiezen. De oude methode zou hierdoor in de war raken en de verkeerde mensen ontslaan.

Samenvattend in een Metafoor

Stel je een orkest voor dat een symfonie speelt.

• De oude manier: Je kijkt naar wie het hardst blaast of het meest beweegt. Als een cellist stil zit, denk je: "Die is niet nodig." Maar misschien houdt die cellist de hele harmonie bij elkaar.
• De nieuwe manier (van dit paper): Je luistert naar de rol van elk instrument. Je simuleert: "Wat gebeurt er als we de cello vervangen door een simpele toon?" Als de muziek nog steeds mooi klinkt, is de cello misschien overbodig. Als de muziek instort, is hij essentieel.

Dit paper leert ons hoe we die simuleer-test snel en slim kunnen doen, zodat we van een groot, rommelig orkest een strak, efficiënt ensemble kunnen maken dat precies hetzelfde klinkt, maar veel minder mensen nodig heeft.

Technische samenvatting

1. Probleemstelling

Diepe neurale netwerken bereiken indrukwekkende voorspellingsnauwkeurigheid, maar zijn moeilijk mechanistisch te interpreteren. Een kernprobleem is dat observatiegedrag alleen niet onderscheidt tussen modellen die stabiele, generaliseerbare algoritmen implementeren en die welke gebruikmaken van spurious (schijnbare) patronen in de trainingsdata.

Om dit op te lossen, is er interesse in causale abstractie: het extraheren van een betrouwbaar, hoog-niveau causaal model (Structural Causal Model of SCM) uit een complex, laag-niveau neuraal netwerk. De uitdaging ligt in het ontdekken van zo'n abstractie voor bestaande, voorgetrainde netwerken.

• De moeilijkheid: Het vinden van een abstractie vereist doorgaans brute-force "interchange interventions" (het uitwisselen van interne activaties tussen inputs) of hertraining, wat computationeel zeer duur is.
• Het doel: Een efficiënte methode vinden om interne variabelen te identificeren die vervangen kunnen worden door eenvoudigere structuren (zoals constanten of lineaire combinaties) zonder de causale betrouwbaarheid onder interventies te verliezen.

2. Methodologie

De auteur reframed het probleem door gestructureerde pruning (het verwijderen van neurale eenheden) te zien als een zoektocht naar benaderende causale abstracties. Het netwerk wordt behandeld als een deterministisch SCM.

Kernconcepten:

1. Mechanism Replacement: In plaats van een kandidaat-hoog-niveau model te veronderstellen, wordt gezocht naar abstracties door mechanismen te vervangen:
   • Harde interventie: Vervangen van een eenheid door een constante ($a_j := c$).
   • Zachte interventie: Vervangen van een eenheid door een affiene functie van behouden eenheden ($a_j := \beta + \sum w_k a_k$).
2. Commutativiteit: Een geldige abstractie moet voldoen aan commutativiteit: interveniëren op het hoog-niveau en terugrekenen moet overeenkomen met interveniëren op het laag-niveau en dan abstraheren. Dit wordt gemeten via Interchange Intervention Accuracy (IIA).
3. Tractabele Surrogaat (Second-Order Taylor Expansie):
   • Directe optimalisatie van IIA is te duur. De auteur leidt een kwadratische benadering af voor de verandering in taakverlies ($\Delta L$) veroorzaakt door mechanismevervanging.
   • Dit resulteert in een gesloten vorm voor de optimale vervangingsconstante ($c^*$) en een per-eenheid score ($s_j$) die de minimale kost van vervanging kwantificeert.
   • De formule voor de optimale constante combineert een kromming-gewogen gemiddelde met een gradiëntcorrectie:
     $$c^*_j = \frac{\sum h_s A_{s,j}}{\sum h_s} - \frac{\sum g_s}{\sum h_s}$$
     Waarbij $g_s$ de gradiënt en $h_s$ de kromming (curvature) is.
4. Exacte Compilatie: Vervangingen kunnen worden omgezet in een kleiner, dicht netwerk zonder runtime masking:
   • Constante vervangingen worden verwerkt via bias folding (aanpassen van de bias in de volgende laag).
   • Affiene vervangingen worden verwerkt via weight folding (herverdelen van gewichten).

3. Belangrijkste Bijdragen

1. Constructieve Abstractie-ontdekking: Formalisering van het probleem als het zoeken naar een gereduceerd SCM door mechanismevervanging, in plaats van het verifiëren van een bestaand model.
2. Efficiënte Surrogaat: Een tweede-orde benadering die per-eenheid scores berekent in één autodiff-pas, waardoor brute-force zoektochten overbodig worden.
3. Exacte Compilatie: Bewijs dat deze abstracties direct kunnen worden gecompileerd naar standaard dense netwerken, waardoor ze operationeel en efficiënt zijn.
4. Verbinding met Variance-Based Pruning (VBP): De methode toont aan dat VBP (het verwijderen van eenheden met lage activatievariantie) een speciaal geval is van deze causale methode, geldig onder de aanname van uniforme kromming en stationariteit. De nieuwe methode verklaart wanneer VBP faalt (bijv. bij herparameterisatie).
5. Empirische Validatie: Demonstreert dat de ontdekte abstracties een hogere IIA behalen dan traditionele methoden, vooral onder sterke interventies.

4. Resultaten

De methode werd getest op MNIST (handgeschreven cijfers) en een synthetisch Booleaans circuit.

• Fideliteit vs. Complexiteit: De voorgestelde methode (Logit-MSE, gebaseerd op de afgeleide scores) behaalde een hogere Interchange Intervention Accuracy (IIA) en lagere KL-divergentie dan Variance-Based Pruning (VBP), vooral bij agressieve pruning (bijv. 256 eenheden behouden in plaats van 512).
• Invariantie Test (Stress Test): Een cruciaal experiment hanteerde exacte, functiewaarderende herparameterisatie (het herschalen van eenheden en inverse schalen van uitgaande gewichten).
  • VBP faalde: De selectie van eenheden veranderde drastisch (Jaccard-similairiteit daalde naar ~0.4), wat leidde tot minder betrouwbare abstracties.
  • De voorgestelde methode slaagde: De selectie bleef exact hetzelfde (Jaccard = 1.0) en de fideliteit bleef hoog. Dit bewijst dat de methode afhankelijk is van de causale structuur en niet van de coördinaten van de activaties.
• Affiene Vervanging: Het gebruik van affiene vervangingen (in plaats van alleen constanten) verbeterde de IIA bij zeer agressieve pruning, hoewel dit ten koste ging van de KL-fideliteit, wat een afweging (trade-off) biedt.

5. Betekenis en Conclusie

Dit paper biedt een fundamentele link tussen gestructureerde pruning en causale abstractie. Het toont aan dat het verwijderen van neurale eenheden niet slechts een compressietechniek is, maar een manier om een expliciet, interpreteerbaar causaal model te construeren.

De belangrijkste inzichten zijn:

• Traditionele heuristieken zoals variatie-pruning zijn kwetsbaar voor herparameterisatie en missen de causale nuance.
• Door gebruik te maken van kromming (curvature) en gradiënten, kan men robuuste abstracties vinden die bestand zijn tegen functioneel equivalente veranderingen in het netwerk.
• De methode maakt het mogelijk om grote, voorgetrainde netwerken efficiënt te reduceren tot kleinere, causaal betrouwbare modellen zonder hertraining, wat een grote stap is in de richting van mechanistische interpretatie van AI-systemen.