Technische Samenvatting: RAP (Runtime Adaptive Pruning voor LLM Inference)

Probleemstelling

Grote Taalmodellen (LLMs) staan voor aanzienlijke implementatiebarrières vanwege hun enorme reken- en geheugenvoetafdruk. Hoewel modelcompressietechnieken (pruning, kwantisatie, distillatie) oplossingen bieden, lijden bestaande methoden aan kritieke beperkingen:

1. Statische Heuristieken: De meeste huidige pruning-strategieën vertrouwen op vaste, eenmalige heuristieken die uitgaan van statische werklasten. Ze falen in het aanpassen aan runtime-geheugenschommelingen veroorzaakt door fluctuerende gebruikersverzoeken (bijv. variërende batchgroottes en sequentielengtes).
2. Verwaarlozing van KV Cache: Traditionele pruning richt zich voornamelijk op het verminderen van statische modelgewichten (parameters). Echter, bij autoregressieve inferentie groeit het Key-Value (KV) cache-geheugen lineair met de batchgrootte en sequentielengte, en overstijgt dit vaak het statische parametergeheugen. Vaste pruning-verhoudingen kunnen de dynamische verschuiving van parameter-gedomineerde naar KV-cache-gedomineerde geheugenknelpunten niet opvangen.
3. Systeemniveau Variatie: Real-world edge- en cloud-omgevingen ervaren stochastische geheugentijdbeschikbaarheid door gelijktijdig draaiende applicaties en interferentie, waardoor vaste begrotingsaannames ongeldig worden.

De kernuitdaging die wordt aangepakt, is hoe een optimale LLM-pruningbeleid te selecteren dat dynamisch aanpast aan heterogene, tijdsvariërende verzoekwerklasten, terwijl het fluctuerende geheugenbudgetten respecteert zonder de outputkwaliteit te verslechteren.

Methodologie: Het RAP Framework

De auteurs stellen RAP (Runtime Adaptive Pruning) voor, een elastisch pruning-framework gedreven door Reinforcement Learning (RL) dat compressiestrategieën dynamisch aanpast op basis van realtime systeemtoestanden.

1. Greedy Sequential Importance (GSI)

Om de beperkingen van "eenmalige" pruning (die inter-lagen afhankelijkheden negeert en leidt tot prestatiedegradatie) aan te pakken, hanteert RAP Greedy Sequential Importance (GSI).

• Mechanisme: In plaats van alle blokken één keer te scoren, verwijdert GSI iteratief het blok waarvan de uitsluiting resulteert in de minimale verslechtering van modelprestaties (gemeten door perplexiteit).
• Hercalibratie: Na elke verwijdering worden de belangrijkheidsscores van de resterende blokken opnieuw geëvalueerd binnen de nieuwe, gecontracteerde architectuur. Dit stapsgewijze proces controleert foutaccumulatie en houdt rekening met de dynamische gevoeligheid van lagen naarmate de modelstructuur verandert.

2. Reinforcement Learning Formulering

RAP formuleert de selectie van een pruningbeleid als een sequentieel besluitvormingsprobleem (Markov Besluitproces) waarbij een RL-agent leert om geheugenefficiëntie en modeltrouw in evenwicht te brengen.

• Toestand ($S_t$): De agent observeert een samengestelde toestand bij elke inferentiestap $t$:
  • Verzoekkenmerken: Batchgrootte ($R_{bs}$) en sequentielengte ($R_{sql}$).
  • Modelconfiguratie: Belangrijkheidsscores van alle Multi-Head Attention (MHA) en Feed-Forward Network (FFN) blokken afgeleid van GSI.
  • Systeemtoestand: Huidig beschikbaar systeemgeheugen ($Sys_{avail}$) en geschatte geheugenoverhead ($Sys_{req}$) voor het voorgestelde actie.
• Actie ($A_t$): De agent selecteert een pruningbeleid door te beslissen of specifieke MHA- en FFN-blokken behouden of verwijderd moeten worden. Om de combinatorische explosie van de actieruimte ($2^{2N}$) te beheersen, wordt de beslissing opgesplitst in sequentiële selecties van één blok.
• Beloning ($R_t$): Een verenigde beloningsfunctie balanceert twee doelen:
  • Perplexiteit ($R_{ppl}$): Moedigt het behoud van taalmodelleringprestaties aan.
  • Geheugenstraf ($R_{mem}$): Straft piekgeheugengebruik tijdens inferentie.
  • De beloning is een gewogen som: $R_t = \sum (At)_i (\alpha R_{ppl, i} - \beta R_{mem, i})$, waarbij hyperparameters $\alpha$ en $\beta$ de afweging tussen nauwkeurigheid en geheugen afstemmen.

3. Uitvoeringsflow

1. De RL-agent observeert de realtime omgeving (invoerlengte, batch, geheugentijdbeschikbaarheid).
2. Het selecteert een pruning-actie (welke blokken te verwijderen) op basis van de huidige toestand.
3. Het basismodel voert inferentie uit op de geprepareerde architectuur.
4. De agent ontvangt een beloning op basis van het resulterende geheugengebruik en de perplexiteit, en werkt zijn beleid bij om toekomstige beslissingen te optimaliseren.

Belangrijkste Bijdragen

1. Runtime-Adaptief Framework: RAP is het eerste framework dat modelgewichten en KV cache-geheugen gezamenlijk "on the fly" overweegt, en dynamisch de sparsiteitsniveaus aanpast op basis van de onmiddellijke werklast en toestands van het apparaat.
2. RL-gedreven Beleid: De auteurs formuleren pruning als een RL-probleem, waarbij ze een agent trainen om een optimaal beleid te leren dat geheugenbesparing afweegt tegen generatietrouw, en zo statische heuristische baselines overtreft.
3. GSI voor Structurele Pruning: De introductie van Greedy Sequential Importance adresseert het probleem van inter-lagen afhankelijkheid in transformer-pruning, en biedt een robuustere methode voor het bepalen van blokbelangrijkheid onder agressieve compressie.
4. Empirische Validatie: Uitgebreide experimenten tonen aan dat RAP consequent state-of-the-art statische pruning-baselines (bijv. LLM-Pruner, SliceGPT, ShortGPT) overtreft over diverse modellen (Llama2, Llama3, Qwen) en geheugenbeperkingen.

Experimentele Resultaten

Experimenten werden uitgevoerd op Llama2-7B, Llama3-8B, Qwen1.5-7B en Qwen2.5-7B onder 80% en 60% verenigde geheugenbudgetten (dekkend zowel parameters als KV cache).

• Generatievermogen: RAP vertoonde de kleinste perplexiteitsdrift. Op Llama2-7B met een 80% budget verhoogde RAP de perplexiteit met slechts +6,33 (WikiText2) en +21,47 (PTB), wat significant beter was dan de op één na beste methode met marges van respectievelijk 22,95 en 146,93.
• Downstream Taken: Bij zero-shot gezond verstand redeneren behield RAP 86% van de nauwkeurigheid van het dichte model onder een 80% budget, en overtrof leidende baselines met +8,2% (Llama2-7B) en +11,5% (Llama3-8B).
• Weerbaarheid: Onder zware 60% geheugenbeperkingen was RAP de enige methode die meer dan 60% van de prestaties van het dichte model behield, terwijl andere methoden catastrofale degradatie leden.
• Ablatiestudies:
  • Het verwijderen van GSI (met behulp van statische eenmalige scoring) veroorzaakte een daling van 13,17% in de gemiddelde nauwkeurigheid van gezond verstand.
  • Het verwijderen van de RL-agent (met willekeurige blokdropping) resulteerde in een enorme toename van de perplexiteit (+301,71 op WikiText2) vergeleken met het geleerde beleid, wat de noodzaak van adaptieve besluitvorming benadrukt.
• Overhead: De RL-controller voegt verwaarloosbare runtime-overhead toe (<1% toename in latentie) vanwege zijn compacte architectuur (18K parameters), terwijl de offline GSI-berekening een eenmalige kostenpost is die over de implementatie wordt afgeschreven.

Betekenis en Claims

Het paper claimt dat RAP de kloof overbrugt tussen statische compressietechnieken en dynamische real-world implementatiescenario's. Door weg te bewegen van "one-size-fits-all" heuristieken, stelt RAP het volgende in staat:

• Schaalbare Edge Implementatie: Het stelt LLMs in staat om op het apparaat te passen onder strakke, fluctuerende geheugenbeperkingen door substantiële gewichten te pruningen wanneer nodig, terwijl het volledige capaciteit behoudt wanneer middelen dit toelaten.
• Efficiëntie zonder Kwaliteitsverlies: Het bereikt superieure geheugenbesparingen en snellere inferentie met minimale impact op de outputkwaliteit, en navigeert effectief door de afwegingen tussen parameter-zware en KV-cache-zware knelpunten.
• Generaliseerbaarheid: Het framework blijkt architectuur-agnostisch te zijn en presteert goed over verschillende modelfamilies (Llama, Qwen) en generaties.

De auteurs positioneren RAP als een flexibel, contextbewust compressiemechanisme dat LLM-inferentie haalbaar maakt in heterogene omgevingen waar vaste strategieën falen.