Two-Stage Regularization-Based Structured Pruning for LLMs

Probleemstelling

De implementatie van Large Language Models (LLMs) wordt gehinderd door hun enorme parameteraantallen, wat leidt tot hoge rekenkosten en geheugeneisen. Hoewel gestructureerd pruning (het verwijderen van hele lagen of kanalen) een veelbelovende oplossing is, hebben bestaande methoden voor laag-gebaseerd pruning (layer-wise pruning) een fundamenteel nadeel:

1. Kennisverlies: Bestaande methoden bepalen direct welke lagen "onbelangrijk" zijn op basis van een metric en verwijderen deze. Dit negeert het feit dat zelfs ogenschijnlijk minder belangrijke lagen waardevolle kennis kunnen bevatten.
2. Retraining-afhankelijkheid: Na het verwijderen van lagen daalt de prestatie aanzienlijk, wat uitgebreide hertraining (fine-tuning) vereist om de prestaties te herstellen. Dit veroorzaakt aanzienlijke rekenkosten en vertraagt de implementatie.

Methodologie: TRSP

De auteurs introduceren TRSP (Two-Stage Regularization-Based Structured Pruning), een nieuwe methode die eerst kennis herverdeelt voordat er gesneden wordt. Het proces verloopt in vier fasen:

1. Data Voorbereiding: Er wordt een zeer klein deel van de data (bijv. 128 voorbeelden) uit standaard benchmarks (zoals WikiText-2) geselecteerd. Dit minimaliseert de rekenkosten tijdens het trainingsproces.
2. Fase 1: Leren van Laag-Weights (Iteratief):
   • Aan elke transformer-laag wordt een leerbaar gewicht $S[i]$ toegekend (initialisatie op 1).
   • De output van elke laag wordt vermenigvuldigd met dit gewicht.
   • Een loss-functie wordt geoptimaliseerd die bestaat uit de taalmodellering-verlies ($L(W, X)$) en een $\ell_1$-norm regularisatie van de laag-weights.
   • Dit proces is iteratief: in elke iteratie wordt de laag met het laagste gewicht gemaskeerd (tijdelijk uitgeschakeld) en toegevoegd aan de set van te verwijderen lagen. Dit voorkomt dat opeenvolgende lagen (die samen cruciaal kunnen zijn) per ongeluk worden verwijderd, wat vaak gebeurt bij "one-shot" benaderingen.
3. Fase 2: Tweede Stadia Regularisatie (Kennisoverdracht):
   • Voordat de lagen daadwerkelijk worden verwijderd, wordt er een extra regularisatie toegepast op de lagen die als "te verwijderen" zijn gemarkeerd.
   • De regularisatie straft het verschil tussen de input en output van deze lagen af (met $\ell_1$- of $\ell_2$-norm).
   • Doel: Dit dwingt de modelkennis om zich te verplaatsen van de te verwijderen lagen naar de overgebleven lagen. De te verwijderen lagen worden effectief gedwongen om een "identiteitsfunctie" te worden (input $\approx$ output), waardoor ze minder kennis bevatten en het verwijderen ervan minder schadelijk is.
4. Pruning:
   • De lagen met de laagste weights worden fysiek verwijderd.
   • Het resultaat is een compact model dat geen hertraining vereist.

Belangrijkste Bijdragen

• Kennisbehoud: Door de twee-stadia regularisatie wordt kennis actief overgebracht van te verwijderen lagen naar de behouden lagen, wat het directe kennisverlies van traditionele pruning minimaliseert.
• Retraining-vrij: In tegenstelling tot bestaande methoden die uitgebreide fine-tuning nodig hebben, behoudt TRSP de prestaties direct na het pruning-proces zonder extra training.
• Efficiëntie: De methode vereist slechts een minimale dataset voor de regularisatie, wat de kosten voor het pruning-proces drastisch verlaagt.
• Iteratieve Strategie: Het iteratief leren van weights (in plaats van één keer alle weights te berekenen) voorkomt de verwijdering van opeenvolgende lagen, wat de stabiliteit van het model verbetert.

Resultaten

De auteurs hebben TRSP getest op diverse modellen (Phi-2, OPT, LLaMA2, LLaMA3) en taken (generatie en zero-shot).

• Prestatie: TRSP presteert significant beter dan sterke concurrenten zoals SLEB, ShortGPT en LaCo.
  • Bij een pruning-ratio van 25% op LLaMA2-7B is de perplexiteit 20% lager dan die van ShortGPT.
  • Op OPT-13B daalt de gemiddelde nauwkeurigheid slechts met 1% ten opzichte van het dichte model, terwijl andere methoden veel grotere verliezen lijden.
• Versnelling: Omdat het hele lagen worden verwijderd, is er sprake van aanzienlijke end-to-end versnelling.
  • Bij 50% pruning op OPT-13B resulteert dit in een 75% toename in throughput en een 46% reductie in latentie.
• Robuustheid: De methode werkt consistent goed over verschillende datasets (WikiText-2, Alpaca, PTB, C4) en verschillende pruning-ratio's (tot 60%), zelfs zonder hertraining.
• Schaalbaarheid: Experimenten op het grotere Llama-2-70B model bevestigen dat TRSP ook daar de beste prestaties levert vergeleken met andere gestructureerde en ongestructureerde pruning-methoden.

Betekenis en Impact

TRSP biedt een nieuwe paradigma voor het comprimeren van LLMs. Het lost het klassieke dilemma op tussen kwaliteit (kennisbehoud) en efficiëntie (snelheid/grootte) door kennis vooraf te migreren in plaats van deze te verliezen bij het snijden.

De belangrijkste implicaties zijn:

1. Kostenefficiëntie: Het elimineert de noodzaak voor dure hertraining, wat de drempel voor het deployen van kleine, snelle LLMs verlaagt.
2. Hardware-vriendelijkheid: Als een laag-gebaseerde methode biedt het echte versnelling op hardware (in tegenstelling tot ongestructureerde pruning die vaak speciale hardware vereist).
3. Toekomstige richting: Het paper suggereert dat het herverdelen van kennis via regularisatie een krachtige techniek is die verder kan worden toegepast op andere modelarchitecturen (zoals CNNs) en taken.

Kortom, TRSP is een doorbraak in gestructureerd pruning die LLMs efficiënter en praktischer maakt voor real-world toepassingen zonder in te leveren op intelligentie of taalvaardigheid.