LEAP: Learnable End-to-End Adaptive Pruning of Large Language Models

Technische Samenvatting: LEAP: Leerbaar End-to-End Adaptief Snoeien van Grote Taalmodellen

Probleemstelling
De implementatie van Grote Taalmodellen (LLM's) wordt steeds vaker geblokkeerd door geheugen- en rekenkrachtbeperkingen. Hoewel gewichtssparse een oplossing biedt, is het landschap verdeeld in gestructureerde, semi-gestructureerde (bijv. 2:4) en ongestructureerde regimes. Recente vooruitgang in GPU-kernen (bijv. FlashLLM, SpInfer) en dataflow-hardware heeft de bottleneck verschoven van het uitvoeren van ongestructureerde sparse naar het induceren ervan met minimaal verlies aan nauwkeurigheid.

Huidige state-of-the-art methoden voor ongestructureerd snoeien, zoals Wanda, SparseGPT, Thanos en ADMM, vertrouwen op laag-voor-laag surrogate die zijn afgeleid van het Optimal Brain Surgeon (OBS)-principe. Deze methoden minimaliseren lokale reconstructiefouten in plaats van het end-to-end modelverlies. Bijgevolg accumuleren ze lokale fouten die zich versterken in diepe netwerken, wat leidt tot significante degradatie van de nauwkeurigheid, vooral onder agressieve sparse (50–60%).

Omgekeerd optimaliseren end-to-end leerbare maskermethoden zoals MaskLLM en PATCH maskers direct tegen taalmodelleringverlies, maar zijn beperkt tot semi-gestructureerde patronen. Hun parametrisatie rust op een categorische verdeling over geldige patronen binnen een groep. Hoewel dit hanteerbaar is voor kleine groepen (bijv. $\binom{4}{2}=6$ patronen voor 2:4 sparse), wordt deze aanpak combinatorisch onhanteerbaar voor ongestructureerde sparse. Voor een rij met breedte $d=4096$ bij 50% sparse is het aantal geldige maskers $\binom{4096}{2048} \approx 10^{1229}$, waardoor het opslaan en indexeren van patroon-logits onmogelijk wordt.

Methodologie: LEAP
De auteurs introduceren LEAP (Learnable End-to-End Adaptive Pruning), een raamwerk ontworpen om end-to-end ongestructureerd maskerleren hanteerbaar te maken door de onhanteerbare parametrisatie van categorieën-over-patronen te vervangen door een per-gewicht benadering.

• Per-Gewicht Bernoulli Relaxatie: In plaats van een logit toe te wijzen aan elk geldig patroon, wijst LEAP een leerbare logit toe aan elk individueel gewicht. Voor een gewichtsmatrix $W \in \mathbb{R}^{m \times n}$ wordt een parallelle logit-matrix $P \in \mathbb{R}^{m \times n}$ geïntroduceerd. Het stochastische masker $M$ wordt gegenereerd via de Gumbel-sigmoid relaxatie:
  $$M = \sigma\left(\frac{\alpha P + g}{\tau}\right)$$
  waarbij $g$ Gumbel-ruis is, $\sigma$ de sigmoidfunctie is, $\alpha$ een schalingsfactor is en $\tau$ een temperatuurparameter. Deze formulering schaalt lineair met het aantal gewichten ($O(mn)$) in plaats van het aantal geldige patronen.
• Optimalisatiestrategie:
  • Bevroren Gewichten: De vooraf getrainde gewichten $W$ worden bevroren gehouden. Alleen de masker-logits $P$ worden getraind. Dit behoudt de kalibratie van het vooraf getrainde model en vereenvoudigt de implementatie.
  • Initialisatie: $P$ wordt geïnitieerd met een eenmalig Wanda-masker (door geselecteerde invoerwaarden op $+s$ te zetten en anderen op $-s$) om een redelijk startpunt te bieden en problemen met koude start te voorkomen.
  • Schema's: Een schaal-schema ($\alpha$) wordt opgevoerd om de sigmoid naar binaire waarden te duwen, terwijl een temperatuur-schema ($\tau$) afneemt om het masker scherper te maken.
  • Verliesfunctie: Het doel combineert het taalmodelleringverlies op een kalibratiestroom met twee regularisatietermen:
    1. Global Sparsity Regularizer: Handhaaft de doel-dichtheid globaal over alle lagen, waardoor individuele lagen hun dichtheid kunnen aanpassen op basis van end-to-end belangrijkheid.
    2. Magnitude-Aware Term: Beïnvloedt de optimalisatie om gewichten met hogere magnitude te behouden, om degeneratieve minima te voorkomen waarin kritieke gewichten worden verwijderd.
• TrainingsEfficiëntie: De methode traint ongeveer 2.000 iteraties op een kleine kalibratiestroom met algemeen tekst (SlimPajama).

Belangrijkste Bijdragen

1. Identificatie van Combinatorische Obstructie: Het artikel identificeert formeel dat parametrisaties van categorieën-over-patronen (gebruikt door MaskLLM/PATCH) niet schalen naar ongestructureerde sparse vanwege de exponentiële groei van geldige patronen.
2. Hanteerbare Hervorming: De auteurs stellen een per-gewicht Bernoulli-via-Gumbel-sigmoid hervorming voor die end-to-end differentieerbaar leren voor ongestructureerde maskers op LLM-schaal mogelijk maakt.
3. LEAP Raamwerk: Een lichtgewicht snoeiraamwerk dat werkt op bevroren gewichten, waarbij slechts een kleine kalibratiedataset en ongeveer 2.000 trainingsstappen nodig zijn.
4. Empirische Validatie: Uitgebreide evaluatie over vijf LLM-families (Qwen-2.5, Gemma-3, LLaMA-3.2, LLaMA-3.1) variërend van 0,5B tot 8B parameters bij 50% en 60% sparse.

Resultaten
LEAP werd geëvalueerd tegenover laag-voor-laag baselines (Wanda, SparseGPT, Thanos, ADMM) en de semi-gestructureerde leerbare baseline MaskLLM.

• Nauwkeurigheidswinst: Over tien (model, sparse) instellingen verbeterde LEAP het gemiddelde zero-shot-nauwkeurigheidsresultaat over zes taken met +2,59 punten ten opzichte van ADMM, de beste laag-voor-laag baseline.
• Specifieke Verbeteringen: De grootste winst was +5,40 punten op LLaMA-3.2 1B bij 60% sparse (50,39 versus 44,99).
• Perplexiteit: LEAP behaalde consistent een lagere perplexiteit (PPL) dan baselines. Bijvoorbeeld, op LLaMA-3.1 8B bij 50% sparse behaalde LEAP 7,66 PPL tegenover 9,12 voor ADMM.
• Sparse Allocatie: Experimenten toonden aan dat onder end-to-end optimalisatie met een globaal budget, LEAP convergeert naar een bijna uniforme per-blok dichtheidsallocatie, in tegenstelling tot de niet-uniforme variaties die worden gezien bij semi-gestructureerde methoden.
• Rekenkosten: Hoewel duurder dan eenmalige methoden (uren op H100's versus minuten), is LEAP aanzienlijk goedkoper dan MaskLLM (bijv. 276 GPU-uren voor LLaMA-3.1 8B versus duizenden A100-uren voor MaskLLM op vergelijkbare schalen).

Betekenis en Claims
Het artikel claimt dat LEAP de fundamentele obstructie oplost die end-to-end maskerleren in ongestructureerde sparse verhindert. Door maskerleren te ontkoppelen van gewichtsupdates, biedt het een praktische weg tot het comprimeren van foundation-modellen zonder de kalibratie van vooraf getrainde gewichten te offeren.

De auteurs benadrukken dat LEAP maskers produceert in het 50%–60% ongestructureerde regime, die direct compatibel zijn met bestaande versnelde kernen (bijv. FlashLLM, SpInfer, MACKO). De betekenis ligt niet in het creëren van nieuwe inferentie-kernen, maar in het mogelijk maken van de implementatie van deze kernen op hoogwaardige, gesnoeide modellen die eerder onbereikbaar waren via laag-voor-laag methoden. Het werk suggereert dat ongestructureerde sparse, wanneer geïnduceerd via end-to-end leren, een gunstige afweging kan bieden tussen compressie en nauwkeurigheid voor implementatie op hardware van algemene aard.

Beperkingen en Scope

• Bevroren Gewichten: De huidige scope houdt gewichten bevroren. Hoewel gezamenlijke gewicht-en-masker optimalisatie wordt genoemd als een natuurlijke uitbreiding, is deze niet geïmplementeerd in dit werk.
• Uniforme Allocatie: De observatie van uniforme per-blok dichtheid wordt gepresenteerd als een regiem-specifieke bevinding voor de bestudeerde modellen en sparse-niveaus, niet als een algemene regel tegen niet-uniforme allocatie.
• Veiligheid: De auteurs waarschuwen dat gecomprimeerde modellen verschillende faalmodi kunnen vertonen op lange-staart invoer en raden her-evaluatie aan voor veiligheidskritieke implementaties.