Laat Elke Token Tellen: Naar Het Verbeteren Van Long-Context Prestaties Met KV-Cache Verwijdering

Samenvatting

De sleutel-waarde (KV) cache is een belangrijk knelpunt bij inferentie in lange contexten, waarbij geheugen en rekenkracht toenemen met de sequentielengte. Bestaande KV-verwijderingsmethoden verminderen deze kosten, maar leiden doorgaans tot prestatievermindering ten opzichte van inferentie met volledige cache. Ons belangrijkste inzicht is dat aandacht met volledige cache niet altijd optimaal is: in lange contexten kunnen irrelevante tokens de aandacht wegleiden van nuttig bewijs, dus selectieve, leerbare verwijdering kan de generatie verbeteren in plaats van alleen de volledige cache te benaderen. We introduceren een globale retentie-gebaseerde KV-verwijderingsmethode die het toekomstige nut van elk token leert onder een uniform geheugenbudget. Lichtgewicht retentiepoorten kennen nutsscores toe aan opgeslagen KV-items, en een gedeelde finale scoreprojectie kalibreert deze scores over alle lagen en koppen. Dit maakt één enkel globaal verwijderingsbeleid mogelijk waarin tokens uit verschillende lagen, koppen en modaliteiten rechtstreeks concurreren om cachecapaciteit. We bieden verder theoretische analyse die aantoont dat het preferentieel behouden van nuttige tokens de aandachtverdunning vermindert, en we rechtvaardigen geometrische retentie als een query-agnostische proxy voor toekomstig nut. In diverse benchmarks voor taal en visueel-taalkundig redeneren over lange contexten en voor meerstapsdialogen vermindert onze methode het KV-geheugen aanzienlijk terwijl ze inferentie met volledige cache evenaart of overtreft. Deze resultaten suggereren dat geleerde, globaal gekalibreerde KV-verwijdering niet alleen een compressietechniek is, maar ook een mechanisme om redeneren in lange contexten te verbeteren.

English

The key-value (KV) cache is a major bottleneck in long-context inference, where memory and computation grow with sequence length. Existing KV eviction methods reduce this cost but typically degrade performance relative to full-cache inference. Our key insight is that full-cache attention is not always optimal: in long contexts, irrelevant tokens can dilute attention away from useful evidence, so selective, learnable eviction can improve generation rather than merely approximate the full cache. We introduce a global retention-based KV eviction method that learns each token's future utility under a unified memory budget. Lightweight retention gates assign utility scores to cached KV entries, and a shared final scoring projection calibrates these scores across all layers and heads. This enables a single global eviction policy in which tokens from different layers, heads, and modalities compete directly for cache capacity. We further provide theoretical analysis showing that preferentially retaining useful tokens reduces attention dilution, and we justify geometric retention as a query-agnostic proxy for future utility. Across diverse long-context language and vision-language reasoning, and multi-turn dialogue benchmarks, our method substantially reduces KV memory while matching or surpassing full-cache inference. These results suggest that learned, globally calibrated KV eviction is not only a compression technique, but also a mechanism for improving long-context reasoning.