RCPU: Rotation-Constrained Error Compensation for Structured Pruning of Large Language Models

Probleemstelling

Grote Taalmodellen (LLMs) accumuleren rijke semantische kennis in hun representatieruimte tijdens training op enorme datasets. Structured pruning (het verwijderen van volledige rijen of kolommen van gewichtsmatrices) is een efficiënte methode om de modelgrootte en inferentiekosten te verlagen. Echter, pruning wordt meestal uitgevoerd met een zeer kleine kalibratiedataset, wat onvermijdelijk leidt tot fouten in de uitvoer.

Bestaande methoden hebben beperkingen:

• Directe least-squares fitting: Kan de fouten op de kalibratiedataset wel verkleinen, maar neigt tot overfitting. Dit leidt tot destructieve wijzigingen van de vooraf getrainde gewichten en vervormt de geometrie van de uitvoerrepresentaties (hoeken en normen), wat de generalisatie schaadt.
• Bias-correctie (zoals FLAP): Compenseert alleen het gemiddelde van de fout, maar kan niet goed omgaan met input-afhankelijke directionele mismatches.
• Geen correctie: Simpelweg kolommen verwijderen (bijv. met WANDA-sp) introduceert significante prestatieverlies.

Het kernprobleem is dus hoe men de fouten na pruning kan compenseren zonder de geometrische structuur van de vooraf getrainde kennis te vernietigen, gezien de beperkte hoeveelheid kalibratiegegevens.

Methodologie: RCPU

De auteurs stellen RCPU (Rotation-Constrained Parameter Update) voor, een methode die bestaat uit twee hoofdbestanddelen:

1. Rotatie-geconstrueerde compensatie (Orthogonal Procrustes)

In plaats van een vrij lineaire aanpassing (least-squares), beperkt RCPU de update van de overgebleven parameters tot een rotatie (orthogonale transformatie).

• Formulering: Het probleem wordt geformuleerd als een Orthogonal Procrustes-probleem. Gegeven de oorspronkelijke uitvoer $Y$ en de gepreinde uitvoer $Z$, wordt een rotatiematrix $Q$ gezocht die $Z$ zo dicht mogelijk bij $Y$ brengt:
  $$Q^* = \arg \min_{Q^T Q = I} \| Y - QZ \|_F^2$$
• Oplossing: De optimale rotatie wordt gevonden via Singular Value Decomposition (SVD) van het product $YZ^T$.
• Voordeel: Door de update te beperken tot rotaties (en optioneel een isotrope schaling), blijven de hoeken en relatieve normen van de uitvoerrepresentaties behouden. Dit voorkomt geometrische vervorming (zoals shear of ongewenste schaling) die vaak optreedt bij onbeperkte least-squares updates, vooral bij kleine datasets. Dit verhoogt de statistische stabiliteit en vermindert overfitting.

2. Variance-aware Importance Score

De effectiviteit van de rotatie-geconstrueerde compensatie hangt af van welke kolommen worden verwijderd. Als kolommen die sterk bijdragen aan de hoofdrichtingen van de uitvoer worden verwijderd, is herstel moeilijk.

• Score-functie: De auteurs introduceren een nieuwe importance score $\gamma_j$ voor elke kolom $j$:
  $$\gamma_j = \|W_{[:,j]}\| \cdot \|X_{[j,:]}\| \cdot \text{Var}(X_{[j,:]})$$
  Deze score combineert de magnitude van de gewichten, de schaal van de input-activaties, en de variantie van de input-activaties.
• Doel: Kolommen met hoge variantie in de kalibratie-data worden voorkeur gegeven, omdat deze waarschijnlijk sterk correleren met de dominante uitgangsrichtingen. Dit zorgt ervoor dat de belangrijkste componenten voor de oriëntatie van het model behouden blijven.

Belangrijkste Bijdragen

1. Geometrisch behoudende correctie: Het formuleren van post-pruning herstel als een Orthogonal Procrustes-probleem, waardoor de geometrie (normen en inproducten) van de uitvoerrepresentaties wordt behouden.
2. Variance-aware scoring: Een nieuwe pruning-score die input-variantie integreert om kolommen die cruciaal zijn voor de richting van de uitvoer te behouden.
3. Efficiëntie en Implementatie: De methode vereist geen extra architecturale wijzigingen, geen fine-tuning, en voegt slechts een bescheiden rekenkosten toe (voornamelijk een SVD per laag). Het kan direct worden ingevoegd na bestaande column-pruning methoden zoals WANDA-sp.

Resultaten

De methode is geëvalueerd op Llama-7B en Llama-2-13B met verschillende pruning-ratio's (10%, 20%, 30%) en kalibratie-groottes (128 en 512 tokens).

• Perplexity (PPL): RCPU presteert consistent beter dan de baselines WANDA-sp en FLAP op de WikiText-2 dataset. Bij hogere pruning-ratio's (30%) is het voordeel het grootst.
• Taakprestaties: Op een reeks taalbegrijpings-benchmarks (BoolQ, PIQA, HellaSwag, WinoGrande, ARC, etc.) behaalt RCPU de hoogste gemiddelde nauwkeurigheid.
  • In vergelijking met Least Squares + Ridge (LS+Ridge): Hoewel LS+Ridge soms beter presteert op PPL bij grote kalibratiedatasets, presteert RCPU vaak beter op downstream-taken. Dit wordt toegeschreven aan het lagere aantal vrijheidsgraden van RCPU (door de orthogonale beperking), wat beter is voor het behoud van de vooraf getrainde kennis.
  • In vergelijking met SliceGPT: RCPU presteert beter bij hoge pruning-ratio's, waarschijnlijk omdat SliceGPT een globale transformatie toepast die leidt tot mismatches tussen lagen, terwijl RCPU laag-voor-laag optimaliseert.
• Efficiëntie: De pruning-tijd per laag blijft onder de 10 seconden, wat de methode praktisch toepasbaar maakt voor grote modellen.

Betekenis en Conclusie

RCPU biedt een robuust kader voor het comprimeren van LLMs zonder de noodzaak van kostbare fine-tuning. De kerninzicht is dat het handhaven van de geometrische structuur van de representatieruimte tijdens de correctie van pruning-fouten essentieel is voor het behoud van de modelkwaliteit.

De studie toont aan dat het combineren van een variantie-gevoelige selectie (om belangrijke richtingen te behouden) met een rotatie-geconstrueerde update (om de geometrie te herstellen) superieure resultaten oplevert ten opzichte van bestaande methoden die ofwel te agressief zijn (overfitting) of te beperkt (alleen bias-correctie). Dit opent nieuwe richtingen voor pruning-aware model-updates die statistische stabiliteit combineren met rekenkundige efficiëntie.