Technische Samenvatting: DistillCache

Probleemstelling

Transformer-gebaseerde Large Language Models (LLMs) vertrouwen op Key-Value (KV) caches om de herberekening van token-representaties tijdens autoregressieve decodering te vermijden. Echter, de omvang van de KV-cache groeit lineair met de sequentielengte, wat een ernstige geheugenbottleneck creëert die het contextvenster en de doorvoer van long-context inferentie beperkt.

Bestaande oplossingen vertrouwen primair op heuristische eliminatiestrategieën (bijv. H2O, SnapKV) die tokens behouden op basis van statische aandachtstatistieken of positionele signalen. De auteur stelt dat deze heuristieken vaak falen in het vastleggen van de toekomstige voorspellende invloed van een token, met name in scenario's met langetermijnafhankelijkheden of gestructureerd redeneren waar de belangrijkheid van een token afhankelijk is van de gevolgen voor de downstream generatie in plaats van op het verleden van de aandachtspatronen.

Methodologie: DistillCache

Het artikel stelt DistillCache voor, een reinforcement learning (RL) framework dat KV-cache eliminatie formuleert als een sequentieel beslissingsprobleem onder strikte geheugenbudgetten. In tegen tegenstelling tot heuristische methoden, leert DistillCache een beleid (policy) om een subset van tokens te behouden die de outputdistributie van een full-cache model behoudt.

Kerncomponenten

1. Toestandsruimte (Feature Representatie):
   Voor elk token in de cache observeert het beleid een 5-dimensionale feature vector afgeleid van interne modelsignalen die beschikbaar zijn tijdens de decodering:

   • $\alpha_i$: De gemiddelde aandachtsmassa ontvangen door het token over alle heads en lagen.
   • $\|v_i\|_2$: De $\ell_2$-norm van de waardevector van het token.
   • $\sigma_i$: De variantie van de ontvangen aandachtsmassa over de heads.
   • $H_i$: De aandacht-gewogen gemiddelde entropie van recente queries die naar het token kijken (vangt context-gevoelige belangrijkheid op).
   • $pos_i$: Genormaliseerde sequentiepositie.

2. Actieruimte:
   Een lichtgewicht Multi-Layer Perceptron (MLP) policy netwerk wijst een scalaire score toe aan elk gecachte token. Tijdens de training wordt een stochastische rangschikking geïnduceerd door Gumbel-ruis toe te voegen aan deze scores, waardoor het beleid een top-$B$ subset van tokens kan samplen (waarbij $B$ het budget is) uit de Plackett–Luce distributie. Tijdens de inferentie worden de top-$B$ tokens deterministisch geselecteerd.

3. Trainingsarchitectuur (Teacher-Student setup):
   Om het beleid te trainen zonder grondwaarheid-labels, hanteert de auteur een teacher-student distillatie-aanpak:

   • Teacher Cache: Houdt de volledige, niet-gepruned geschiedenis van tokens bij om de "gouden" next-token distributie ($P_t$) te berekenen.
   • Student Cache: Houdt de geprunede geschiedenis bij volgens de beslissingen van het beleid om de geprunede distributie ($\hat{P}_t$) te berekenen.
   • Beloningssignaal (Reward Signal): De beloning bij elke stap wordt gedefinieerd als de negatieve Kullback-Leibler (KL) divergentie tussen de teacher en student distributies: $R_t = -D_{KL}(P_t \parallel \hat{P}_t)$. Deze doelstelling optimaliseert direct voor distributionele getrouwheid bij elke decimeringstap, in plaats van te vertrouwen op proxy-signalen zoals aandachtsscores.

4. Optimalisatie:
   Het beleid wordt geoptimaliseerd met behulp van het REINFORCE algoritme met entropie-regularisatie en een variantiereductie-baseline (Leave-One-Out). Het trainingsproces wordt offline uitgevoerd, waarbij het geheugengebruik tijdelijk wordt verdubbeld, maar inferentie vereist slechts één forward pass met een verwaarloosbare overhead voor feature extractie en policy scoring.

Belangrijkste Bijdragen

• RL Framework voor Cache Eliminatie: Introductie van DistillCache, een algemeen bruikbaar RL-framework dat tokenretentie optimaliseert voor downstream distributionele getrouwheid onder geheugenbeperkingen, waarmee het verder gaat dan statische heuristische scoring.
• Distributie-behoudende Doelstelling: Demonstratie dat het optimaliseren van de per-stap KL-divergentie meetbare winst oplevert ten opzien van heuristische en gesuperviseerde alternatieven door expliciet lange-horizon trade-offs te leren.
• Uitgebreide Evaluatie: Uitgebreide evaluatie op Mistral-7B-Instruct-v0.3 en zero-shot transfer naar Llama-3-8B-Instruct over LongBench, redeneer-benchmarks (GSM8K, BBH) en throughput metrieken.

Experimentele Resultaten

Geëvalueerd op Mistral-7B-Instruct-v0.3 met een 25% KV-cache budget:

• LongBench Prestaties: DistillCache behoudt 94,2% van de full-cache accuratesse, waarmee het sterke heuristische baselines (H2O, SnapKV) met wel 2,7 absolute punten verslaat en gelijktijdige RL-gebaseerde methoden (ForesightKV, RLKV) met wel 1,4 punten.
• Redeneertaken: Op redeneer-benchmarks is DistillCache competitief met de beste gelijktijdige methoden en presteert het beter onder agressieve compressie (12,5% budget).
• Cross-Model Transfer: Het beleid getraind op Mistral-7B transfeert zero-shot naar Llama-3-8B-Instruct en blijft competitief met heuristische baselines zonder dat per-model hertraining vereist is (een vereiste voor de gelijktijdige RL-baselines).
• Efficiëntie: De methode levert tot wel 2,1× full-cache throughput. Hoewel het een kleine overhead introduceert (~1,4 ms/token) voor feature extractie en policy evaluatie, blijft het vergelijkbaar met andere geleerde methoden en is het aanzienlijk sneller dan full-cache inferentie.
• Foutanalyse: DistillCache vertoont de laagste percentages repetitieve fouten en onjuiste antwoorden vergeleken met heuristieken en andere RL-methoden, wat duidt op een superieure preservatie van sequentie-coherentie en feitelijke accuratesse.

Betekenis en Claims

Het artikel stelt dat DistillCache de effectiviteit demonstreert van geleerde, distributie-bewuste beleidsregels voor geheugenefficiënte long-context LLM inferentie. Door de beslissingen voor tokenretentie te optimaliseren tegen toekomstige distributionele getrouwheid in plaats van retrospectieve aandachtstatistieken, leert het framework budget-beperkte beleidsregels die rekening houden met downstream generatiegevolgen.

De auteur benadrukt dat DistellCache niet louter een geleerde proxy is voor heuristische metrieken, maar een onderscheidende aanpak die direct optimaliseert voor het behoud van de outputdistributie van het model. Hoewel beperkingen zoals trainingskosten (ca. 130 GPU-uren) en een kleine prestatiekloof met RLKV op redeneertaken bij gematigde budgetten worden erkend, concludeert het artikel dat distributie-behoudende adaptieve eliminatie een veelbelovende richting is voor het mogelijk maken van long-context inferentie op hardware met beperkte middelen. Het werk suggereert dat het combineren van token-niveau distributionele getrouwheid met head-level gating strategieën (zoals verkend in hybride experimenten) een complementair pad biedt voor toekomstige optimalisatie.