GRASPrune: Global Gating for Budgeted Structured Pruning of Large Language Models

1. Het Probleem

Grote Taalmodellen (LLMs) zijn duur om te draaien vanwege de hoge geheugeneisen en latentie. Deze kosten worden voornamelijk veroorzaakt door:

• De enorme hoeveelheid modelparameters.
• De berekening van de attention-mechanismen.
• De KV-cache (Key-Value cache), die groeit met de lengte van de inputsequentie en aanzienlijk geheugen verbruikt.

Bestaande methoden voor gestructureerde pruning (het verwijderen van kanalen of hoofdgroepen) hebben vaak drie tekortkomingen:

1. Gebrek aan globale coördinatie: Ze behandelen FFN-kanalen (Feed-Forward Networks) en KV-heads (Key-Value heads) vaak met aparte criteria, terwijl ze concurreren om hetzelfde implementatiebudget.
2. Vooraf bepaalde verdeling: Veel methoden gebruiken vaste verhoudingen per laag of diepte-afhankelijke schema's in plaats van een enkele globale toewijzing te leren.
3. Loskoppeling van budget en training: Veel pipelines leren eerst belangsscores zonder beperkingen en passen het budget pas achteraf toe. Dit verzwakt de link tussen de uiteindelijke architectuur en het trainingsdoel en vereist vaak zware reconstructie of fine-tuning.

2. Methodologie: GRASPrune

GRASPrune is een gestructureerd pruning-framework dat wordt toegepast na het pre-training, zonder de achterliggende gewichten (backbone weights) aan te passen. Het doel is om zowel FFN-kanalen als KV-head-groepen te verwijderen onder één globaal budget.

De kern van de methode bestaat uit drie fasen:

A. Globale Budget-gedreven Formulering

Het probleem wordt geformuleerd als een discrete selectieprobleem onder een budgetbeperking.

• Eenheden: Alle pruneerbare eenheden (FFN-kanalen en KV-head-groepen) over alle lagen worden verzameld in één set.
• Kostenmodel: Elke eenheid heeft een kostenfactor ($c_i$). FFN-kanalen hebben een eenheidskosten van 1, terwijl KV-head-groepen een hogere kostenfactor hebben ($\alpha$), gebaseerd op hun parameteraantal (inclusief GQA-structuur).
• Doel: Minimaliseer de taalmodelleringstijdsverlies ($L$) onder de voorwaarde dat de totale kosten van de bewaarde eenheden binnen het budget $B$ blijven.

B. Projected STE Training (Straight-Through Estimator)

In plaats van eerst scores te leren en daarna te selecteren, enforceert GRASPrune het budget tijdens elke trainingsstap.

1. Gate Scores: Er worden continue scores ($p$) geleerd voor elke eenheid.
2. Projectie: Bij elke stap worden deze scores geprojecteerd naar een harde binaire mask ($m \in \{0, 1\}$) die het budget respecteert. Dit gebeurt door een globale sortering van de eenheden op basis van hun score en het selecteren van de hoogst scorende eenheden totdat het budget wordt overschreden.
3. Forward/Backward:
   • Forward: Het model gebruikt de harde mask ($m$) voor inferentie.
   • Backward: Gradiënten worden berekend via de zachte scores ($p$) met behulp van een Straight-Through Estimator (STE). Dit stelt de optimizer in staat de scores te updaten alsof de mask zacht was, terwijl het model daadwerkelijk met de harde mask werkt.
4. Voordeel: De scores worden geleerd onder dezelfde beperkingen als de uiteindelijke implementatie, wat voorkomt dat het systeem neigt naar "goedkope" eenheden die slecht presteren.

C. Calibratie en Scaling

Na het vaststellen van de mask:

• Er wordt een lichte skalingscalibratie uitgevoerd op de bewaarde eenheden. Er worden schalingsfactoren ($\gamma$) geleerd om schaalverschillen te compenseren die ontstaan door het verwijderen van eenheden.
• Deze factoren worden in de gewichten "gefold" (samengevoegd), wat resulteert in een kleiner, dicht checkpoint zonder extra parameters tijdens de inferentie.

3. Belangrijkste Bijdragen

1. Eén Globale Formulering: Voor het eerst wordt FFN- en KV-pruning gecombineerd in één enkel optimalisatieprobleem met een globaal budget, rekening houdend met verschillende kosten per eenheid.
2. Budget-Feasible Training: De methode gebruikt een projectie-gebaseerde STE-lus die het budget bij elke stap garandeert. Dit voorkomt dat de geleerde scores loskoppelen van de implementatiebeperkingen.
3. Efficiëntie: Het proces vereist geen fine-tuning van het grote model. Het trainen van de gate-scores op een kleine, ongelabelde dataset (512 sequenties) duurt slechts ongeveer 6 minuten op één NVIDIA A100 GPU.
4. Resultaat: Het produceert een compact, dicht model dat direct kan worden gedeployd met standaard stacks.

4. Resultaten

De methode is geëvalueerd op modellen zoals LLaMA-2-7B, LLaMA-2-13B, Vicuna, en nieuwere families zoals LLaMA-3.1 en Qwen3.

• Prestaties (Perplexity): Op LLaMA-2-7B bereikt GRASPrune bij een behoud van 50% van de parameters een perplexiteit van 12.18 op WikiText-2. Dit is aanzienlijk beter dan concurrenten zoals LLM-Pruner, SliceGPT en FLAP (die respectievelijk 312, 26 en 24 scoren bij dezelfde ratio).
• Zero-Shot Accuracy: Het behoudt competitieve prestaties op vijf benchmarks (ARC, HellaSwag, PIQA, etc.), zelfs bij hoge compressie (40% behoud).
• Geheugen en Doorvoer:
  • Door KV-heads te verwijderen, neemt de KV-cache en het piekgeheugen bij lange contexten aanzienlijk af (bijna evenredig met de behoudsratio).
  • De doorvoer (tokens/sec) verbetert aanzienlijk, vooral bij grotere batchgroottes, dankzij de kleinere modelgrootte en minder geheugentransacties.
• Stabiliteit: De methode convergeert snel (binnen 4 epochs) en de mask is stabiel, wat betekent dat de selectie van eenheden niet willekeurig fluctueert.

5. Betekenis en Conclusie

GRASPrune demonstreert dat globale, budget-gedreven gating een eenvoudige maar krachtige methode is om LLMs te comprimeren zonder de kerncapaciteiten van taalmodellering te verliezen.

• Efficiëntie: Het elimineert de noodzaak voor dure fine-tuning of complexe reconstructie-algoritmen.
• Flexibiliteit: Het kan worden toegepast op verschillende modelarchitecturen en biedt inzicht in de trade-off tussen FFN- en KV-capaciteit.
• Praktische Toepassing: Het maakt het mogelijk om grotere contextvensters of hogere batchgroottes te draaien binnen hetzelfde geheugensbudget, wat cruciaal is voor de productiedeployments van LLMs.

Kortom, GRASPrune biedt een nieuwe state-of-the-art aanpak voor post-training compressie die zowel de modelgrootte als de inferentiekosten effectief verlaagt terwijl de kwaliteit behouden blijft.