From Layers to Submodules: Rethinking Granularity in Replacement-Based LLM Compression

Technische Samenvatting: Van Lagen naar Submodules (SUBFIT)

Probleemstelling
Post-training compressie van Large Language Models (LLM's) beoogt de implementatiekosten te verlagen door architecturale componenten te verwijderen. Hoewel methoden die volledige lagen of aaneengesloten blokken verwijderen effectief zijn, verstoren ze de residustroom, wat vaak leidt tot een scherpe kwaliteitsdegradatie. Recente "vervangingsgebaseerde" benaderingen pakken dit aan door verwijderde componenten te vervangen door lichtgewicht, aangepaste modules. Echter, bestaande methoden delen twee restrictieve ontwerpbeperkingen:

1. Full-layer granulariteit: Ze behandelen de volledige Transformer-laag als de atomaire eenheid van verwijdering.
2. Contigüe selectie: Ze verwijderen componenten in opeenvolgende dieptebereiken.

De auteurs stellen dat deze beperkingen te restrictief zijn. Redundantie in pre-trained transformers is niet beperkt tot aaneengesloten regio's, noch is het gelijkmatig verdeeld tussen Attention- en FeedForward (FFN) outputs. Bijgevolg zijn verschillende strategieën nodig om verschillende submodule-typen te benaderen, en hoewel verwijderbare componenten niet noodzakelijkerwijs geclusterd hoeven te zijn binnen aaneengesloten dieptebereiken.

Methodologie: SUBFIT
Het artikel introduceert SUBFIT (Submodule-level Fitted residual replacement), een post-training compressiemethode die opereert op het granulariteitsniveau van individuele Attention- en FFN-submodules. Het proces omvat twee hoofdfasen:

1. Niet-contigüe Selectie:

   • SUBFIT selecteert een doel aantal Attention- en FFN-submodules onafhankelijk over de modeldiepte op basis van een sparsity budget $s$.
   • De selectie is niet-contigüe; modules worden onafhankelijk op basis van een score gekozen zonder een opeenvolgend blok te vormen.
   • Scoringstrategie:
     • Attention: Modules worden gerangschikt op een "Impact score" berekend op kalibratie-activaties, die meet hoeveel een submodule de residustroom roteert en de relatieve magnitude ervan.
     • FFN: Voor base-modellen wordt een vervangingsbewuste score gebruikt die verwijderingseffecten en surrogate-fouten combineert. Voor instructie-getunede modellen wordt een cosine-only score gebruikt om ruis door slecht gekalibreerde residu-magnitudes te vermijden.
   • Volgorde: Attention-submodules worden eerst geselecteerd en vervangen. Vervolgens worden FFN-submodules gescoord en geselecteerd op het gedeeltelijk gecomprimeerde model om rekening te houden met de verschuiving in de residustroom die door de Attention-vervanging is veroorzaakt.

2. Fitted Residual Replacement:

   • Elke geselecteerde submodule wordt vervangen door een lichtgewicht parametrische surrogaat $\hat{F}_\ell$ die de oorspronkelijke submodule-bijdrage ($y = F_\ell(x)$) benadert, in plaats van de interne berekening.
   • Surrogaatstructuur: De surrogaat gebruikt een affine correctie (elementgewijze gate, bias en centrering) plus een gestructureerde interactieterm gedefinieerd door een low-rank operator $U_\ell = V_\ell^\top W_\ell$.
   • Asymmetrische Parameterisatie:
     • Attention: Elke verwijderde Attention-submodule ontvangt zijn eigen onafhankelijke low-rank basis ($V_\ell$) en parameters, wat de heterogene rollen over verschillende dieptes weerspiegelt.
     • FFN: Een gedeelde input-basis $V$ wordt gebruikt voor alle geselecteerde FFN-lagen. Alleen de output-factoren ($W_\ell$) en affine termen blijven laag-specifiek. Deze cross-layer parameterdeling vermindert de implementatiekosten aanzienlijk.
   • Fitting: De parameters worden analytisch teruggewonnen via closed-form fitting op kalibratiestatistieken (ridge-regularized least squares en eigenwaardedecompositie), wat geen gradiënt-gebaseerde optimalisatie of full-model fine-tuning vereist.

Belangrijkste Bijdragen

1. Identificatie van Beperkingen: De auteurs identificeren full-layer granulariteit en contigüe selectie als limiterende factoren in bestaande vervangingsgebaseerde compressie.
2. SUBFIT Framework: Een nieuwe methode die gezamenlijke selectie en vervanging uitvoert op het niveau van submodules (Attention en FFN) met niet-contigüe selectie.
3. Asymmetrisch Ontwerp: Een specifieke architectuur waarbij Attention onafhankelijke bases gebruikt, terwijl FFN-bypasses een gemeenschappelijke input-basis delen om de parameter-overhead te minimaliseren.
4. Uitgebreide Evaluatie: Uitgebreide testen over 10 LLM's (5 base, 5 instruction-tuned) en 5 sparsity-niveaus (12,5% tot 37,5%), waarbij een verbeterde trade-off wordt aangetoond ten opzichte van vier baselines (Streamline, ReplaceMe).

Experimentele Resultaten

• Perplexity (Base Modellen): SUBFIT bereikt de laagste perplexiteit-degradatie in 9 van de 10 model-sparsity instellingen. Het gat met de sterkste baseline wordt aanzienlijk groter onder agressieve compressie (bijv. bij 25% sparsity vertoont SUBFIT een 2,42× perplexiteit-degradatie versus 4,34× voor de beste baseline).
• Downstream Nauwkeurigheid (Instruction-Tuned): Bij 25% sparsity behoudt SUBFIT gemiddeld 84,6% van de dense downstream nauwkeurigheid, waarmee het de sterkste baseline (81,6%) verslaat. Het is de enige methode die boven de 80% nauwkeurigheid behoudt bij 25% sparsity en boven de 73% bij 37,5% sparsity.
• Stabiliteit: SUBFIT vertoont een significant lagere variantie in perplexiteit-degradatie vergeleken met baselines, die instabiliteit vertonen (bijv. standaarddeviatie die orden van grootte groeit) bij hoge sparsity.
• Efficiëntie: De methode levert meetbare versnellingen in inferentie (1,18×–1,40× in time-to-first-token) en KV-cache besparingen (proportioneel aan sparsity) zonder dat sparse kernels of custom hardware nodig zijn.
• Ablatie Studies:
  • Residu Compensatie: Het vervangen van FFN's is de dominante factor voor stabiliteit; het vervangen van zowel Attention als FFN's levert de beste resultaten op.
  • Shared Basis: De gedeelde FFN-basis bereikt een prestatie die vergelijkbaar is met per-laag bases, maar met aanzienlijk minder parameters.
  • Rank Sensitiviteit: FFN-kwaliteit verbetert tot rank 4096, terwijl Attention verzadigt bij rank 256, wat de asymmetrische rank-instellingen rechtvaardigt.

Significantie en Claims
Het artikel claimt dat granulariteit een onderbelicht aspect is in post-training compressie. Door over te stappen van laag-niveau naar submodule-niveau vervanging, herstelt SUBFIT meer van het gedrag van het dense model dan contigüe full-layer vervanging onder hetzelfde parameterbudget. De resultaten suggereren dat niet-contigüe, submodule-specifieke recovery een sterkere geaggregeerde kwaliteit–efficiëntie trade-off biedt, met name onder agressieve compressie. De methode is ontworpen om goedkoop te implementeren (closed-form fitting) en compatibel te zijn met bestaande inference stacks. De auteurs suggereren bescheiden dat deze aanpak kan generaliseren naar andere removal units, zoals attention heads of individuele neuronen, maar claimen niet directe generalisatie naar die specifieke units in dit werk.