Pruning via Causal Attribution Preserves Reasoning Performance in Large Language Models

Technische Samenvatting: Causal Attribution Pruning (CAP)

Probleemstelling
Large Language Models (LLM's) vertonen sterke vermogens voor meerstapsredeneren, maar veroorzaken aanzienlijke inferentiekosten, wat de implementatie in de echte wereld belemmert. Hoewel modelpruning een oplossing biedt door redundante parameters te verwijderen zonder hertraining, verslechteren bestaande methoden vaak de prestaties bij complexe redeneertaken. Huidige benaderingen vertrouwen op correlationele signalen, zoals gewichtsmagnitude (bijv. Wanda) of lokale gradiëntstatistieken. Deze methoden gaan ervan uit dat kleine gewichten onbelangrijk zijn; echter kunnen parameters met een lage magnitude cruciaal zijn voor meerstapsredeneringscircuits. Bovendien falen puur correlationele criteria in het vastleggen van cross-layer interacties die essentieel zijn voor chain-of-thought (CoT) inferentie, wat vaak leidt tot "layer collapse", waarbij volledige modules hun functie verliezen.

Methodologie: Causal Attribution Pruning (CAP)
De auteurs introduceren Causal Attribution Pruning (CAP), een training-vrije, one-shot pruning-methode die ontworpen is om redeneerprestaties te behouden. In tegen tegenstelling tot methoden die gewichten rangschikken op basis van statische eigenschappen, meet CAP de causale impact van specifieke modelcomponenten op de taakprestaties. Het framework werkt in drie stadia:

1. Kalibratie en Causale Scoring (Stadium 1):

   • Een kleine kalibratiedataset ($D_{cal}$) wordt geconstrueerd door te samplen uit redeneergerichte benchmarks: GSM8K (wiskunde), StrategyQA (multi-hop common sense) en ARC-Challenge (wetenschap).
   • Voor elke attention head $h_\ell$ in het model voert CAP een interventionele masking-experiment uit. De head wordt tijdelijk gemaskeerd tijdens de forward pass, en de resulterende toename in token-level cross-entropy loss ($\Delta(h_\ell)$) wordt gemeten ten opzichte van het intacte model.
   • Om robuustheid te garanderen tegen variantie door kleine steekproefgroottes, wordt de dataset verdeeld in $K$ subsets, en wordt de mediaan van de loss-toename berekend.
   • Dit proces identificeert heads die functioneel cruciaal zijn voor redeneren, waardoor ze worden onderscheiden van heads die slechts numeriek groot maar redundant zijn.

2. Gewichtsniveau Belangrijkheidsconversie (Stadium 2):

   • De head-level causale scores worden omgezet in genormaliseerde belangrijksheidsfactoren $I(h_\ell)$, die worden toegewezen aan een bereik van $[1.0, 10.0]$. Heads met een hoge impact ontvangen hogere factoren.
   • Deze factoren worden gepropageerd naar individuele gewichten binnen de corresponderende projectiematrices ($W_q, W_k, W_v, W_o$). De belang score voor een specifiek gewicht $W_{ij}$ wordt berekend als het product van zijn magnitude en de belangsfactor van zijn geassocieerde head-regio: $S_{ij} = |W_{ij}| \cdot I(h_{region})$.
   • Voor MLP-lagen, die geen natuurlijke head-gepartitioneerde structuur hebben, wijst de methode een gemiddelde belang per laag toe, afgeleid van de co-gelokaliseerde attention heads. Dit wordt genoteerd als een benadering.

3. Globale Pruning (Stadium 3):

   • Alle gewichten over het hele model worden gerangschikt op basis van hun belang-gewogen magnitude scores ($S_{ij}$).
   • Een globale drempelwaarde wordt bepaald via binary search om de gewenste sparsiteit te bereiken (bijv. 10%, 20%, 50%). Gewichten onder de drempelwaarde worden op nul gezet, wat een ongestructureerd sparsiteitspatroon creëert dat compatibel is met standaard inferentie-infrastructuur.

Belangrijkste Bijdragen

• Formalisering van Causale Attributie voor Pruning: De auteurs formaliseren een methode om de verwachte toename in taakspecifieke loss te meten wanneer individuele attention heads worden gemaskeerd, waarmee zij verder gaan dan correlationele proxies om functionele kritikaliteit te kwantificeren.
• Budget-bewuste Gewichtsniveau Pruning: Zij introduceren een procedure die head-level causale scores omzet in gewichts-niveau belangwaarden. Dit maakt fijnmazige magnitude-pruning mogelijk onder een globale sparsiteitsrestrictie, terwijl de structuur van kritieke redeneercircuits behouden blijft.
• Empirische Validatie: De methode wordt geëvalueerd op Llama-3-8B-Instruct en Mistral-7B-Instruct over de benchmarks GSM8K, StrategyQA en ARC-Challenge.

Resultaten

• Matige Sparsiteit (10–20%): CAP presteert over het algemeen beter dan de Wanda-baseline in het behouden van redeneerprestaties.
  • Op Llama-3-8B-Instruct bij 20% sparsiteit bereikt CAP een relatieve verbetering van 61% ten opzichte van Wanda op ARC-Challenge (70,8% vs. 43,9% nauwkeurigheid).
  • Op GSM8K bij 10% sparsiteit laat CAP een relatieve verbetering van 20% zien ten opzichte van Wanda (71,5% vs. 59,5%).
  • De auteurs merken op dat hoewel CAP soms achterblijft bij Wanda op specifieke taken (bijv. Llama-3 GSM8K bij 20%), het consistent de prestaties behoudt op meerstapsredeneringstaken waar de correlationele criteria van Wanda falen.
• Hoge Sparsiteit (50%): Prestaties verslechteren aanzienlijk voor beide methoden. CAP lijdt onder catastrofaal falen op Llama-3 (perplexity springt naar 428,2), wat wordt toegeschreven aan de grofmazige behandeling van MLP-lagen. Op dit niveau verwijdert de globale drempelwaarde essentiële MLP-gewichten die een lage belangscore kregen op basis van de attention head-benadering.
• Perplexity vs. Redeneren: Bij matige sparsiteit bereikt CAP vaak een betere redeneernauwkeurigheid ondanks een iets hogere perplexiteit vergeleken met Wanda, wat suggereert dat er een ontkoppeling bestaat tussen algemene taalmodellering-kwaliteit en taakspecifieke redeneerbehoud.
• Architecturale Beperkingen: Voorlopige experimenten op Mixture-of-Experts (MoE) modellen suggereren dat CAP minder effectief is, omdat dynamische routing de head-belangrijkheid niet-stationair maakt.

Betekenis en Claims
Het artikel claimt dat attention-head-level causale attributie een superieur criterium is voor het behoud van redeneerprestaties vergeleken met correlationele pruning-methoden bij equivalente sparsiteitsniveaus. Door expliciet de causale bijdrage van heads aan redeneertaken te meten, kan CAP parameters identificeren en beschermen die "numeriek klein maar functioneel kritiek" zijn, welke door magnitude-gebaseerde methoden worden weggegooid.

De auteurs benadrukken dat hun werk bescheiden is in omvang:

• Het is een training-vrije methode, wat het computationeel efficiënter maakt vergeleken met benaderingen die gewicht-updates vereisen (zoals SparseGPT).
• Het is taak-gealigneerd, specifiek gekalibreerd op redeneerdata in plaats van generieke taalmodellering.
• Het erkent beperkingen, met name het gebrek aan fijnmazige causale attributie voor MLP-neuronen en de moeilijkheid om de methode toe te passen op dynamisch gerouteerde architecturen (MoE).

De kernbijdrage is de demonstratie dat interventionele, taak-gealigneerde pruning beter in staat is om specifieke capaciteiten (redeneren) te behouden dan methoden die geoptimaliseerd zijn voor algemene perplexiteit of statische gewichtstatistieken, mits het sparsiteitsniveau matig blijft.