MaskPro: Linear-Space Probabilistic Learning for Strict (N:M)-Sparsity on LLMs

Technische Samenvatting: MaskPro – Lineaire Ruimtelijke Probabilistische Lering voor Strikte (N:M)-Sparsiteit op LLM's

Probleemstelling
De snelle schaalvergroting van Large Language Models (LLM's) heeft een aanzienlijke knelpunt gecreëerd in de inferentie-efficiëntie, wat effectieve compressietechnieken noodzakelijk maakt. Semi-gestructureerde sparsiteit, specifiek het (N:M)-patroon waarbij $N$ elementen worden behouden uit elke $M$ opeenvolgende gewichten, biedt een hardware-vriendelijke oplossing voor versnelling en geheugenreductie. Het identificeren van het optimale (N:M)-masker is echter een formidabel combinatorisch optimalisatieprobleem. Bestaande benaderingen vallen in twee categorieën met onderscheidende beperkingen:

1. Regelgebaseerde methoden: Deze maken gebruik van greedy, laag-voor-laag zoekopdrachten met handgemaakte metrieken (bijvoorbeeld het minimaliseren van de reconstructiefout per laag). Hoewel ze computationeel goedkoop zijn, lijden ze aan inherente biases en aanzienlijke kloven met end-to-end taakprestaties.
2. Leringsgebaseerde methoden (bijvoorbeeld MaskLLM): Deze optimaliseren maskers direct via gradiëntafname op een generatiefout. Hoewel ze effectief zijn, brengen ze verbodsbepalende trainingskosten met zich mee. Specifiek vereist MaskLLM het opslaan van logits voor alle mogelijke maskers binnen een groep van $M$ gewichten, wat leidt tot een geheugencomplexiteit van $O(\binom{M}{N} \frac{d}{M})$. Voor grote $M$ schaalt dit exponentieel, waardoor het onhaalbaar wordt voor grote modellen. Bovendien induceert de enorme combinatorische ruimte een hoge variantie in beleidsgradiënten, wat leidt tot trainingsinstabiliteit.

Methodologie: MaskPro
De auteurs stellen MaskPro voor, een nieuw lineaire-ruimtelijk probabilistisch raamwerk ontworpen om (N:M)-sparsiteit efficiënt en stabiel te leren.

1. Lineaire Ruimtelijke Probabilistische Relaxatie:
   In plaats van een categorische verdeling te definiëren over alle $\binom{M}{N}$ mogelijke maskers voor een groep van $M$ gewichten, herformuleert MaskPro het probleem. Het stelt een categorische verdeling op over de $M$ individuele elementen binnen elke groep. Het (N:M)-masker wordt vervolgens gegenereerd door een N-weg steekproef zonder teruglegging uit deze verdeling uit te voeren.

   • Theoretische Basis: De auteurs bewijzen dat elk (N:M)-masker kan worden weergegeven als een probabilistische som van $N$ verschillende basisvectoren.
   • Geheugenefficiëntie: Deze herformulering reduceert het aantal leerbare parameters (logits) van $O(\binom{M}{N} \frac{d}{M})$ naar een lineaire $O(d)$, waarbij $d$ het totale aantal gewichten is. Dit elimineert de geheugenexplosie die geassocieerd is met eerdere leringsgebaseerde methoden.

2. Gefineerde Beleidsgradiënt Schatter (PGE):
   Om de logits te optimaliseren, maakt MaskPro gebruik van beleidsgradiënten. De auteurs identificeren dat standaard beleidsgradiënten die ruwe foutwaarden gebruiken, instabiel zijn in LLM's vanwege de hoge variantie veroorzaakt door minibatch-verschillen (waarbij de foutvariatie door datasteekproefname vaak de foutvariatie door maskersteekproefname overtreft).

   • Foutresiduen: In plaats van de ruwe fout $f(m \odot w, \xi)$ te gebruiken, update MaskPro de logits op basis van het foutresidu ten opzichte van een initiële basismasker $m_0$: $f(m \odot w, \xi) - f(m_0 \odot w, \xi)$. Dit isoleert de impact van de maskerwijziging.
   • Bewegend Gemiddelde Tracker: Om de training verder te stabiliseren en numerieke explosies te voorkomen, wordt een bewegend gemiddelde tracker $\delta$ geïntroduceerd om de basislijn te gladstrijken. De updateregel wordt:
     $$\pi_{t+1} = \pi_t - \eta (f(m_t \odot w, \xi) - f(m_0 \odot w, \xi) - \delta) \nabla \log p(m_t|\pi_t)$$
   • Theoretische Garanties: Het artikel levert bewijzen die aantonen dat deze verfijnde schatters onbevooroordeelde schatters blijven van de ware beleidsgradiënt. Bovendien is, onder de voorwaarde dat de fout van het gesteste masker minder dan twee keer de basisfout is (een voorwaarde die in de praktijk gemakkelijk wordt voldaan), de variantie van de verfijnde schatter strikt lager dan die van de standaard beleidsgradiënt.

Belangrijkste Bijdragen

1. Lineaire Ruimtelijk Raamwerk: MaskPro reduceert het geheugenvereiste voor het leren van (N:M)-sparsiteit van exponentieel $O(\binom{M}{N} \frac{d}{M})$ naar lineair $O(d)$ door het selectieproces te modelleren als N-weg steekproef zonder teruglegging uit een categorische verdeling over $M$ elementen.
2. Verbeterde Beleidsgradiënt: De auteurs introduceren een variantie-reducerend updatemechanisme met behulp van foutresiduen en een bewegend gemiddelde tracker, waarmee de trainingsinstabiliteit wordt aangepakt die wordt veroorzaakt door de hoge variantie van beleidsgradiënten in grote combinatorische ruimtes.
3. Theoretische Analyse: Uitgebreide bewijzen worden geleverd met betrekking tot de onbevooroordeeldheid van de schatters en de eigenschappen van variantiereductie van de voorgestelde updateregels.

Experimentele Resultaten
Uitgebreide experimenten werden uitgevoerd op meerdere LLM's (Gemma-7B, Vicuna-7B, LLaMA-2-7B, DeepSeek-7B) over diverse downstream-taken (bijvoorbeeld HellaSwag, RACE, PIQA) en sparsiteitspatronen ((2:4), (4:8), (8:16)).

• Prestaties: MaskPro presteert consequent beter dan state-of-the-art regelgebaseerde methoden (SparseGPT, Wanda, Pruner-Zero) en bereikt prestaties die vergelijkbaar zijn met het geheugenintensieve MaskLLM, vaak met een gemiddelde verbetering van meer dan 2% in nauwkeur.
• Geheugenschaalbaarheid: Waar MaskLLM Out-of-Memory (OOM) fouten ondervindt op 8x A100 GPU's voor (4:8)-sparsiteit op 7B-modellen, behoudt MaskPro stabiel geheugengebruik (~36-37 GB) ongeacht het sparsiteitspatroon.
• Data-efficiëntie: MaskPro toont opmerkelijke robuustheid tegen datavolume. Het bereikt stabiele prestaties zelfs wanneer getraind met slechts 1 data-staal, terwijl MaskLLM honderdduizenden stalen nodig heeft om te convergeren.
• Trainings-efficiëntie: De methode vereist aanzienlijk minder trainingsstalen en minder trainingstijd in vergelijking met MaskLLM, terwijl het geheugengebruik vergelijkbaar blijft met regelgebaseerde methoden.

Betekenis en Claims
Het artikel claimt dat MaskPro een haalbare weg biedt voor de praktische, schaalbare lokale implementatie van LLM's door de geheugen- en stabiliteitsknelpunten van leringsgebaseerde semi-gestructureerde sparsiteit op te lossen. De auteurs benadrukken dat hun benadering de kloof overbrugt tussen de hoge prestaties van leringsgebaseerde methoden en de efficiëntie van regelgebaseerde methoden. Door het geheugenvoetafdruk te reduceren naar lineaire ruimte en het trainingsproces te stabiliseren via verfijnde beleidsgradiënten, stelt MaskPro het leren van optimale (N:M)-maskers op schaalbare modellen mogelijk zonder de verbodsbepalende computationele overhead die voorheen met dergelijke taken geassocieerd was. De methode wordt gepresenteerd als een generaliseerbaar raamwerk voor probabilistisch maskerleren dat robuust is tegen datatekorten en modelschaal.