POP: Prefill-Only Pruning for Efficient Large Model Inference

Probleemstelling

Grote Taalmodellen (LLMs) en Vision-Language Modellen (VLMs) vertonen indrukwekkende prestaties, maar hun implementatie wordt gehinderd door hoge rekenkosten. Bestaande methoden voor gestructureerde pruning (het verwijderen van hele lagen of kanalen) zijn weliswaar hardware-efficiënt, maar leiden vaak tot aanzienlijke nauwkeurigheidsverlies, vooral bij open-ended generatieve taken.

De auteurs stellen dat dit falen voortkomt uit een stadium-agnostische ("one-size-fits-all") aanpak. Traditionele pruning-methoden behandelen de twee fasen van autoregressieve inferentie identiek:

1. Prefill-fase: Het parallelle verwerken van de input-prompt om de Key-Value (KV) cache te genereren voor context.
2. Decode-fase: Het sequentiële genereren van tokens, waarbij de huidige token wordt gecodeerd én de kansverdeling voor de volgende token wordt gemodelleerd.

De auteurs betogen dat deze fasen functioneel asymmetrisch zijn en verschillende gevoeligheid voor pruning vereisen. Bestaande methoden negeren deze asymmetrie, wat leidt tot suboptimale prestaties.

Methodologie: Prefill-Only Pruning (POP)

Het paper introduceert POP, een stadia-bewuste inferentiestrategie die diepe lagen uitsluitend verwijdert tijdens de computatie-intensieve prefill-fase, terwijl het volledige model behouden blijft voor de gevoelige decode-fase.

De methode bestaat uit drie kerncomponenten:

1. Virtuele Gates voor Lagenbelang (Virtual Gate Mechanism)
Om te bepalen welke lagen redundant zijn, introduceren de auteurs een "virtuele gate" ($g_l$) voor elke laag. Door de residual-branches (Attention en FFN) te vermenigvuldigen met deze parameter, kunnen ze het belang van een laag schatten zonder de laag fysiek te verwijderen.

• Ze benaderen de toename in verlies wanneer $g_l$ van 1 naar 0 gaat.
• Door gebruik te maken van een tweede-orde Taylor-expansie en eigenschappen van de Fisher-informatie (gebaseerd op sampling van de modelverdeling in plaats van ground-truth), kunnen ze een efficiënte schatter voor het belang van de laag berekenen met slechts één forward-backward pass.

2. Stadia-bewuste Belanganalyse
De analyse onthult een opvallend verschil tussen de fasen:

• Decode-fase: Diepe lagen zijn kritiek voor het voorspellen van de volgende token.
• Prefill-fase: Diepe lagen zijn grotendeels redundant voor het coderen van contextinformatie; de eerste lagen zijn hier het belangrijkst.
  Op basis hiervan wordt een statische strategie gekozen: de laatste 1/3 van de lagen wordt verwijderd tijdens de prefill-fase.

3. Implementatie en Overgangshandtering
Om de overgang tussen de "geprunte" prefill-fase en de "volledige" decode-fase naadloos te maken, worden twee mechanismen ingevoerd:

• Onafhankelijke KV-projecties: Hoewel de zware Attention- en FFN-berekeningen worden overgeslagen tijdens de prefill, worden de Key- en Value-projecties ($W^K, W^V$) toch uitgevoerd op de input. Dit zorgt ervoor dat de decode-fase, die het volledige model gebruikt, een complete KV-cache heeft voor alle lagen. De rekenkosten hiervan zijn verwaarloosbaar (< 5%).
• Boundary Handling: De laatste input-token ($x_N$) wordt niet verwerkt door het geprunte model, maar door het volledige model. Dit is cruciaal omdat de eerste gegenereerde token ($x_{N+1}$) sterk afhankelijk is van de diepe lagen. Door $x_N$ als de eerste decode-stap te behandelen, wordt de nauwkeurigheid van de eerste gegenereerde token behouden.

Belangrijkste Resultaten

De auteurs hebben POP geëvalueerd op diverse modellen (Llama-3.1, Qwen3-VL, Gemma-3) en taken (redeneren, code generatie, lange context, multimodaal).

• Nauwkeurigheid: POP behoudt bijna volledige nauwkeurigheid vergeleken met het volledige model. In tegenstelling tot bestaande gestructureerde pruning-methoden (zoals SliceGPT en ShortGPT) die catastrofaal falen bij open-ended generatie (bijv. daling van 79% naar <1% op GSM8K), presteert POP consistent goed.
  • Voorbeeld: Op Llama-3.1 behoudt POP 97% van de prestaties op GSM8K en 95% op HumanEval.
• Snelheid (TTFT): POP bereikt een Time-to-First-Token (TTFT) versnelling van maximaal 1,37x voor lange contexten (2048 tokens).
  • De versnelling is het grootst bij lange sequenties, omdat de prefill-fase dan rekenintensief wordt (compute-bound).
  • Bij korte sequenties is de versnelling beperkter omdat de prefill-fase dan geheugen-gebonden is (memory-bound) en de laatste token toch het volledige model nodig heeft.
• Vergelijking: POP overtreft gestructureerde pruning-methoden in zowel snelheid als nauwkeurigheid. Het presteert vergelijkbaar met ongestructureerde pruning (zoals Wanda) qua nauwkeurigheid, maar biedt wel hardware-compatibiliteit omdat het structurele lagen verwijdert (geen speciale kernels nodig).

Bijdragen en Significantie

De belangrijkste bijdragen van dit werk zijn:

1. Inzicht in Asymmetrie: Het paper demonstreert dat diepe lagen essentieel zijn voor generatie (decode) maar redundant zijn voor contextcodering (prefill).
2. POP-strategie: Een nieuwe, stadia-bewuste pruning-methode die de efficiëntie van de prefill-fase verhoogt zonder de generatieve capaciteit te schaden.
3. Praktische Toepasbaarheid: POP vereist geen hertraining van het model en werkt met standaard hardware, wat het een "plug-and-play" oplossing maakt voor lange contexten en multimodale taken.

Beperkingen:
POP vermindert niet het piekgeheugengebruik (VRAM), omdat het volledige model gewicht voor de decode-fase moet worden geladen. Het is dus het meest effectief in scenario's die beperkt zijn door rekencapaciteit (compute-bound) en niet door geheugen (capacity-bound).

Conclusie:
POP biedt een nieuwe richting voor efficiënte inferentie van grote modellen door de functionaliteit van inferentiefasen te ontleden. Het overwint de klassieke afweging tussen snelheid en nauwkeurigheid die bestaande gestructureerde pruning-methoden beperkt, en opent de deur voor verdere optimalisaties in kwantisatie en architectuurontwerp die rekening houden met de specifieke rol van inferentiefasen.