Sparsity Induction for Accurate Post-Training Pruning of Large Language Models

1. Het Probleem

Grote Taalmodellen (LLMs) vertonen indrukwekkende prestaties, maar hun enorme parametergroottes leiden tot hoge rekentkosten, latentie en energieverbruik, wat implementatie beperkt. Post-Training Sparsity (PTS) is een veelbelovende compressietechniek die gewichten verwijdert zonder volledige hertraining.

Echter, bestaande PTS-methoden (zoals magnitude pruning, Wanda, SparseGPT) stuiten op fundamentele beperkingen:

• Gebrek aan natuurlijke sparsiteit: Dichte matrices in LLMs hebben van nature geen hoge sparsiteit.
• Verstoring van modelstatistieken: Het direct verwijderen van gewichten (naïef hard masking) verstoort de zelfconsistente verdelings- en schaalstructuur die tijdens het trainen is opgebouwd.
• Diminishing returns: Het verfijnen van "importance scores" (hoe belangrijk een gewicht is) levert steeds minder op, omdat de onderliggende verdeling van de data en schaalverhoudingen tussen kanalen niet wordt hersteld. Dit leidt tot significante prestatieverlies, zelfs bij matige sparsiteitsniveaus.

2. Methodologie: Sparsity Induction (SI)

De auteurs introduceren Sparsity Induction (SI), een raamwerk dat het model proactief voorbereidt op pruning door de gewichtsverdeling en feature-structuur aan te passen voordat de daadwerkelijke verwijdering plaatsvindt. Dit gebeurt op twee niveaus:

A. Verdelingsniveau (Distribution Level)

Het doel is om de verdeling van de gewichten te "scherpen" zodat belangrijke kanalen grotere waarden hebben en onbelangrijke kanalen kleiner worden, zonder de functionaliteit van het model te veranderen.

• Techniek: Er wordt gebruikgemaakt van functioneel equivalente herschalingen (reparameterization).
  • Voor lineaire lagen: Per-kanaal schaling ($s$) en verschuiving ($\delta$) worden geïntroduceerd. De transformatie $Y = WX + b$ wordt herschreven als $Y = \tilde{W}\tilde{X} + \tilde{b}$, waarbij de nieuwe parameters $\tilde{W}$ en $\tilde{X}$ de oorspronkelijke functionaliteit behouden maar de statistieken aanpassen.
  • Voor Attention-lagen: Inverse schalingen worden toegepast op Queries ($Q$) en Keys ($K$) zodat $QK^T$ onveranderd blijft.
• Implementatie: Deze schalingsfactoren worden getraind op een kleine calibratiedataset met een lichtgewicht objectief (output-matching). Na training worden de factoren volledig opgeslokt (absorbed) in de gewichtsmatrices. Dit betekent dat er geen extra parameters of inferentie-overhead zijn.

B. Feature-niveau (Feature Level)

Om de zoekrichting voor de optimalisatie te stabiliseren en feature-sparsiteit te bevorderen vanuit een laag-rang (low-rank) perspectief:

• Spectral Norm Loss: Er wordt een regularisatieterm toegevoegd aan het verlies (loss function) die de spectrale norm van de gewichten minimaliseert. Dit moedigt een laag-rang structuur aan, wat de feature-sparsiteit verhoogt.
• Robuuste Initialisatie: De schalingsfactoren worden niet willekeurig gestart, maar geïnitieerd op basis van robuuste statistieken (mediaan en gemiddelde) van de output-activaties.
• Doel: Dit zorgt voor snellere convergentie en betere prestaties zonder volledige parameter-finetuning.

C. Efficiëntie: Snelle Hessian Update

Om de optimalisatie snel te maken, vermijden de auteurs het opnieuw berekenen van activaties.

• Ze vervangen de activatievariatie door een diagonale Gauss-Newton/Hessian-proxy.
• Door de schalingsfactoren in de proxy te integreren, kunnen de importance scores (zoals bij de Wanda-methode) in constante tijd ($O(d_{in})$) worden ververst. Dit resulteert in een 22x versnelling in de update-tijd ten opzichte van klassieke methoden.

3. Belangrijkste Bijdragen

1. Concept van Sparsity Induction: Een nieuwe aanpak die het model "sparsiteitsvriendelijk" maakt voordat pruning plaatsvindt, in plaats van alleen te proberen de beste gewichten te selecteren.
2. Twee-dimensionaal Raamwerk: Een geïntegreerde oplossing die werkt op zowel het verdelingsniveau (via equivalente schaling) als het feature-niveau (via spectrale norm loss).
3. Geen Inference Overhead: Alle toegevoegde parameters worden vooraf in de bestaande gewichten opgeslokt, waardoor het model compatibel blijft met hardware-versnelling (zoals N:M semi-structured sparsity) zonder extra rekentijd.
4. Efficiëntie: Een geoptimaliseerd framework dat de berekeningstijd voor het bijwerken van importance scores drastisch reduceert.

4. Resultaten

De methode is getest op diverse modellen (OPT, LLaMA-1, LLaMA-2) en taken (taalmodellering, zero-shot benchmarks zoals ARC, HellaSwag).

• Perplexiteit (PPL): SI vermindert de PPL-degradatie aanzienlijk. Bij agressieve sparsiteit (bijv. 50% of 2:4 patroon) presteert SI+Baseline (bijv. Wanda+SI) veel beter dan de baselines alleen. Bijvoorbeeld, bij OPT-1.3B met 50% sparsiteit daalt de PPL van 1712 (Magnitude) naar 26.39 (Magnitude+SI).
• Zero-Shot Accuracy: Over zes benchmarks heen verbetert SI consistent de nauwkeurigheid van één-shot pruning methoden. Voor LLaMA-7B met 50% sparsiteit stijgt de gemiddelde nauwkeurigheid van 51.47% (Wanda) naar 52.50% (Wanda+SI).
• Snelheid:
  • De "Fast Hessian Update" verkort de tijd voor het bijwerken van de Wanda-metric op LLaMA-7B van 345s naar 15s (22.65x sneller).
  • De end-to-end latency voor inferentie blijft gelijk aan de baselines (geen extra overhead), maar levert wel een 1.24x snelheidswinst op ten opzichte van dichte modellen bij 2:4 sparsiteit.

5. Betekenis en Conclusie

Dit paper biedt een fundamentele verschuiving in hoe we naar post-training pruning kijken. In plaats van alleen te focussen op het vinden van de "beste" maskers, toont het aan dat het aanpassen van de modelverdeling voordat pruning plaatsvindt cruciaal is voor het behoud van prestaties.

De Sparsity Induction methode is uniek omdat:

• Het de "sparsiteitsvriendelijkheid" van het model verhoogt zonder de inferentie-architectuur te veranderen.
• Het werkt als een plug-and-play module die compatibel is met bestaande pruning-methoden (Magnitude, Wanda, SparseGPT).
• Het de kosten van compressie verlaagt door de noodzaak voor zware hertraining te elimineren en toch hoge nauwkeurigheid te behouden.

Dit opent de weg voor efficiëntere, snellere en minder geheugen-intensieve LLM's die makkelijker te deployen zijn in resource-beperkte omgevingen.