GHOST: Unmasking Phantom States in Mamba2 via Grouped Hidden-state Output-aware Selection & Truncation

Probleemstelling

De overgang van Mamba1 naar Mamba2 heeft de staatdimensie ($N$) van de State Space Model (SSM) vergroot van 16 naar 128 om de temporele modellering te verbeteren. Hoewel dit de prestaties verhoogt, creëert het een aanzienlijke inferentie-flesnek:

• Geheugenbandbreedte: De recurrente staat ($H_t$) van een 1.3B-parameter model groeit van ongeveer 12 MB naar 100 MB. Dit verzadigt de geheugenbandbreedte tijdens autoregressieve generatie en verslechtert de cache-localiteit.
• Beperkingen van bestaande methoden:
  • Unstructured pruning: Methoden zoals SparseGPT of Wanda verkleinen de gewichten, maar laten de activaties (en dus de recurrente staat) dicht (dense) achter, waardoor er geen bandbreedtebesparing optreedt.
  • Statische magnitude-pruning: Deze methoden baseren zich op vaste gewichtnormen. Het paper identificeert hier een kritiek falen: ze onderscheiden niet tussen statische capaciteit en dynamisch gebruik. Dit leidt tot het onbedoeld verwijderen van "Phantom States" (kanalen met lage gewichtnormen maar hoge dynamische activiteit) en het behouden van "Corporeal States" (kanalen met hoge gewichten maar lage activiteit).
  • Gradient-based pruning: Methoden zoals Taylor-pruning zijn nauwkeurig maar vereisen backpropagation en het opslaan van Hessians, wat leidt tot prohibitieve geheugeneisen (bijv. >45 GB VRAM voor een 1.3B model) en lastige "one-shot" schattingen die lijden onder distributieshift.

Methodologie: GHOST

Het paper introduceert GHOST (Grouped Hidden-state Output-aware Selection and Truncation), een gestructureerd pruning-framework dat de principes van balanced truncation uit de regeltheorie benadert, uitsluitend gebruikmakend van forward-pass statistieken.

Kernprincipes:

1. Controleerbaarheid en Observabiliteit: GHOST meet twee dynamische eigenschappen voor elke staat:
   • Controleerbaarheid: Hoe sterk de inputgeschiedenis een kanaal beïnvloedt (gemeten via de covariantie van de hidden state $H_t$).
   • Observabiliteit: Hoe sterk een kanaal de toekomstige output beïnvloedt (gemeten via de Hessian van de output-energie ten opzichte van de hidden state, benaderd door de projectie $C'_t$).
2. Saliëntie Score: De belangrijkheid van een staat wordt berekend als het product van deze twee maten (vergelijkbaar met Hankel-singuliere waarden). Dit combineert de "energie" van de input naar de staat en de staat naar de output.
3. Groepsgebaseerde Selectie: GHOST werkt binnen de Grouped Query Attention (GQA) structuur van Mamba2. Het poolt scores over groepen van koppen (heads) binnen een laag, waardoor adaptieve capaciteitsallocatie mogelijk is: complexe dynamische groepen behouden meer staten, terwijl redundante groepen agressief worden ingekort.
4. Implementatie:
   • Het vereist slechts twee forward passes over een kalibratie-set.
   • Er zijn geen gradients, Hessians of backpropagation nodig.
   • Het proces is sequentieel: na het prunen van laag $j$ worden de geactiveerde waarden gebruikt om laag $j+1$ te kalibreren, wat distributieshift minimaliseert.
   • De complexiteit is lineair ($O(N)$) in plaats van kwadratisch, en het geheugengebruik is minimaal ($O(G \cdot N)$).

Belangrijkste Bijdragen

• GHOST Framework: Een data-gedreven, gestructureerd pruning-framework voor Mamba2 dat de staatdimensie $N$ reduceert op basis van gemeten dynamische energie in plaats van statische gewichten.
• Oplossing voor "Proxy Failure": Het demonstreert dat GHOST het falen van magnitude-pruning oplost door "Corporeal States" te verwijderen en "Phantom States" te behouden, wat resulteert in een veel nauwkeurigere selectie van essentiële kanalen.
• Efficiëntie vs. Prestatie: GHOST bereikt prestaties die vergelijkbaar zijn met dure, gradient-based methoden (zoals Taylor-pruning), maar vereist geen backpropagation en werkt binnen de geheugenlimieten van standaard GPU's (bijv. 15 GB vs. 45 GB).

Resultaten

De experimenten zijn uitgevoerd op modellen variërend van 130M tot 2.7B parameters, met name op WikiText-2 en diverse downstream-taken.

• Verdunning (Sparsity): Bij een reductie van 50% in de staatdimensie bereikt GHOST een toename van slechts ~1 perplexity-punt op WikiText-2 (van 13.17 naar 14.23).
• Robuustheid bij hoge verdunning: Terwijl Taylor-pruning instabiel wordt en catastrofale ineenstorting vertoont bij hoge verdunningspercentages (>70%) en kleine modellen, behoudt GHOST een soepele degradatiecurve en blijft bruikbaar zelfs bij 70% verdunning.
• Sequentielengte Generalisatie: GHOST generaliseert uitstekend naar langere contexten dan de kalibratie-set. In tegenstelling tot Taylor-pruning (waarbij de perplexity explodeert bij langere sequenties), volgt GHOST het gedrag van het dichte model.
• Schalingswetten: GHOST is robuust over alle modelgroottes. Taylor-pruning faalt volledig op kleine modellen (130M, 370M) door ruis in de gradient-signalen, terwijl GHOST stabiel blijft.
• Out-of-Distribution (OOD): GHOST toont betere generalisatie naar OOD-taken (zoals code-generatie en wiskunde) dan Hessian-gebaseerde methoden, wat suggereert dat het fundamenteel belangrijke kanalen voor de SSM-dynamiek selecteert en niet alleen kanalen die specifiek zijn voor de kalibratie-data (Engelse tekst).
• Zero-Shot Prestaties: Op benchmarks zoals Lambada, PIQA en ARC presteert GHOST concurrerend met Taylor-pruning en significantly beter dan magnitude- of random-pruning.

Significantie

Dit werk is significant omdat het een praktische oplossing biedt voor de inferentie-flesnek van grote Mamba2-modellen zonder de kwaliteit drastisch te verlagen.

• Democratisering: Door de geheugenbandbreedte-eisen te verlagen en de VRAM-vereisten voor compressie te minimaliseren, maakt GHOST de deploy van grote SSM-modellen haalbaar op standaard hardware.
• Theoretische Integratie: Het paper slaat een brug tussen deep learning en regeltheorie (control theory), toont aan dat forward-pass statistieken voldoende zijn voor geavanceerde modelcompressie, en biedt een alternatief voor de computationally zware gradient-based benaderingen.
• Toekomstige Impact: Het biedt een pad voor efficiëntere inferentie van state-of-the-art sequentiemodellen en suggereert dat regeltheoretische principes waardevol kunnen zijn voor de compressie van andere architecturen.