Adaptive Layer Selection for Layer-Wise Token Pruning in LLM Inference

Probleemstelling

Grote Taalmodellen (LLMs) worden steeds vaker ingezet voor taken met lange contexten (zoals documentanalyse en uitgebreide conversaties). Een kritieke beperking bij de inferentie van deze modellen is de Key-Value (KV) cache, die de key- en value-vectoren van eerdere tokens opslaat om efficiënte generatie mogelijk te maken. Bij zeer lange contexten wordt het geheugengebruik van deze cache een bottleneck.

Om dit op te lossen, zijn er technieken ontwikkeld om tokens te "prunen" (verwijderen) uit de KV-cache. Een populaire aanpak is layer-wise token pruning, waarbij op specifieke lagen in het Transformer-netwerk een subset van belangrijke tokens wordt geselecteerd en de rest wordt verwijderd.

De kernbeperking van bestaande methoden:
Huidige methoden (zoals FastKV en GemFilter) gebruiken vast gedefinieerde, statische lagen voor token-selectie. Dit ontwerp is inflexibel omdat de optimale laag voor selectie sterk varieert afhankelijk van de taak:

• Bij eenvoudige taken (bijv. vragen beantwoorden over een tekst) kunnen relevante tokens vaak al in eerdere lagen worden geïdentificeerd.
• Bij moeilijke taken (zoals KV-retrieval of "Needle in a Haystack") is de semantische gelijkenis tussen vraag en context hoog, waardoor de modelattention pas in diepere lagen convergeert naar de juiste tokens.
  Als een vaste laag wordt gebruikt, leidt dit tot een aanzienlijke prestatiedaling bij moeilijke taken, of vereist het een groter KV-budget (meer geheugen), waardoor het voordeel van compressie verloren gaat.

Methodologie: Adaptive Selection Layer (ASL)

De auteurs stellen ASL voor, een trainingsvrije methode die de selectielag adaptief bepaalt op basis van de waargenomen attention-patronen tijdens de inferentie (prefilling-fase).

Kernidee:
De methode monitort de variantie van de rangschikking (ranks) van tokens, gebaseerd op hun attention-scores, over opeenvolgende lagen.

1. Observatie: In vroege lagen zijn attention-scores vaak uniform verdeeld. In diepere lagen focussen de scores consistent op een stabiele subset van tokens.
2. Meting: ASL berekent de variantie van de rangschikking van de top-$k$ tokens over een venster van recente lagen ($L_{obs}$).
3. Beslissing: Wanneer de relatieve variantie (variantie gedeeld door de initiële variantie) daalt onder een door de gebruiker gespecificeerde drempelwaarde ($\tau$), wordt aangenomen dat de attention stabiel is geconvergeerd. De huidige laag wordt dan gekozen als de selectielag.
4. Implementatie:
   • Prefilling: ASL berekent en slaat gepoolde attention-scores op. Zodra de drempel wordt bereikt, worden tokens eenmalig geselecteerd ("one-shot") en doorgegeven aan alle daaropvolgende lagen.
   • Decoding: De geselecteerde tokens worden gebruikt voor verdere generatie.
   • Integratie: ASL kan worden gecombineerd met bestaande methoden zoals SnapKV (voor compressie vóór de selectielag) en GemFilter (voor een two-pass strategie), zodat het voldoet aan specifieke KV-budgetvereisten.

Algoritme:

• Het algoritme start vanaf een minimale laag $L_{min}$ (bijv. 1/3 van de totale lagen).
• Het berekent de variantie van de rangschikking van tokens over de laatste $L_{obs}$ lagen.
• Als $relative\_variance < \tau$, wordt de selectie uitgevoerd.

Belangrijkste Bijdragen

1. Identificatie van inflexibiliteit: De auteurs tonen aan dat vaste selectielagen leiden tot grote prestatievariaties tussen taken en dat de optimaliteit sterk afhankelijk is van de taakmoeilijkheid.
2. ASL-methode: Een nieuwe, adaptieve methode die de selectielag bepaalt door de stabiliteit van attention-rangschikkingen te monitoren, zonder extra training te vereisen.
3. Flexibele integratie: ASL kan naadloos worden geïntegreerd met bestaande KV-cache-reductiemethoden om zowel de prefilling- als decoding-fase te optimaliseren binnen een vastgesteld geheugenbudget.
4. Uitgebreide evaluatie: De methode is getest op drie benchmarks (InfiniteBench, RULER, NIAH) met contextlengten tot 256k, waarbij ASL superieure resultaten behaalt ten opzichte van state-of-the-art methoden.

Resultaten

De experimenten zijn uitgevoerd op modellen zoals Llama-3.1-8B-UltraLong en Qwen2.5-7B.

• Nauwkeurigheid (Accuracy):

  • ASL presteert aanzienlijk beter dan FastKV, GemFilter en PyramidInfer, vooral bij moeilijke taken zoals KV-retrieval en multi-key NIAH.
  • Bij InfiniteBench behaalt ASL (in een two-pass configuratie) de hoogste gemiddelde score.
  • Bij de NIAH-benchmark (Needle in a Haystack) behaalt ASL bijna 100% nauwkeurigheid over alle contextlengten, terwijl FastKV en SnapKV bij zeer lange contexten (bijv. 148k) significant falen.
  • Wanneer ASL wordt gecombineerd met volledige KV-cache vóór selectie, nadert de prestatie die van volledige KV-cache, terwijl FastKV sterk achterblijft.

• Efficiëntie:

  • TTFT (Time to First Token): ASL is iets langzamer dan FastKV omdat de selectielag vaak dieper ligt (afhankelijk van de taak), maar blijft aanzienlijk sneller dan volledige attention. De TTFT is een acceptabele trade-off voor de winst in nauwkeurigheid.
  • TPOT (Time Per Output Token): De decoding-snelheid is vergelijkbaar met andere pruning-methoden omdat na de selectie slechts een kleine subset van tokens wordt verwerkt.
  • Geheugengebruik: ASL voegt verwaarloosbare overhead toe aan het geheugengebruik (alleen voor het opslaan van gepoolde scores en rangschikkingen) en voldoet strikt aan het ingestelde KV-budget.

• Trade-off:

  • De auteurs tonen aan dat ASL een betere afweging maakt tussen snelheid en nauwkeurigheid. Door de drempelwaarde $\tau$ aan te passen, kunnen gebruikers de balans sturen: een lagere $\tau$ leidt tot diepere selectie (hoger nauwkeurigheid, langzamere TTFT), terwijl een hogere $\tau$ leidt tot vroegere selectie (sneller, maar mogelijk minder nauwkeurig).

Significantie

Dit paper is significant omdat het een fundamenteel probleem in de huidige KV-cache-compressie aanpakt: de eenheid van maatstaf (one-size-fits-all) voor taakselectie.

• Adaptiviteit: In plaats van handmatig te proberen de beste laag te vinden voor elke nieuwe toepassing, stelt ASL het systeem in staat om dynamisch te reageren op de complexiteit van de input.
• Toepasbaarheid: Omdat de methode trainingsvrij is en werkt tijdens de inferentie, kan deze direct worden toegepast op bestaande LLM's zonder kostbare fine-tuning.
• Toekomstperspectief: Het werk opent de deur voor nog geavanceerdere, contextbewuste inferentiestrategieën die de prestaties van LLM's bij extreem lange contexten (honderdduizenden tokens) aanzienlijk kunnen verbeteren zonder in te leveren op de kwaliteit van het antwoord.

Kortom, ASL biedt een robuuste oplossing om de "memory wall" van LLM-inferentie te doorbreken, waarbij de nauwkeurigheid bij complexe taken behouden blijft terwijl de geheugenefficiëntie wordt gemaximaliseerd.