🤖 machine learning

LOCKS: Page-Local Compact Key Summaries for Efficient Long-Context Decoding

LOCKS is een drop-in plugin voor vLLM die decoding met een lange context versnelt door elke geheugenpagina een compacte, laag-rangige spectrale samenvatting toe te wijzen om de aandachtsmassa efficiënt te schatten en alleen de meest relevante pagina's te selecteren, waardoor de latentie en het geheugengebruik aanzienlijk worden verminderd terwijl een nauwkeurigheid die bijna gelijk is aan full-attention behouden blijft.

Oorspronkelijke auteurs: Junsung Hwang

Gepubliceerd 2026-07-28

📖 4 min leestijd☕ Koffiepauze-leesvoer

CC BY 4.0

Oorspronkelijke auteurs: Junsung Hwang

Oorspronkelijk artikel gelicentieerd onder CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dit is een AI-gegenereerde uitleg van het onderstaande artikel. Het is niet geschreven of goedgekeurd door de auteurs. Raadpleeg het oorspronkelijke artikel voor technische nauwkeurigheid. Lees de volledige disclaimer

Stel je voor dat je een enorme bibliotheek vol boeken probeert te lezen om één enkele vraag te beantwoorden. In de wereld van kunstmatige intelligentie zijn Large Language Models (LLM's) als briljante studenten die het hele internet hebben gelezen, maar wanneer ze proberen een vraag te beantwoorden op basis van een zeer lang document, krijgen ze te maken met een lastig probleem. Om na te kunnen denken, moeten ze een "geheugen" bijhouden van alles wat ze tot nu toe hebben gelezen. Dit geheugen wordt de KV-cache genoemd (Key-Value cache). Beschouw dit als een gigantisch whiteboard waarop het model elk woord dat het heeft verwerkt, opschrijft.

Het probleem is dat naarmate het verhaal langer wordt, dit whiteboard enorm groot wordt. Elke keer dat het model het volgende woord wil schrijven, moet het het gehele whiteboard scannen om te beslissen welke eerdere woorden belangrijk zijn. Als het verhaal 100.000 woorden lang is, moet het model bij elke nieuwe letter die het typt, naar 100.000 woorden kijken. Dit is traag en verbruikt een enorme hoeveelheid computergeheugen, alsof je probeert een specifieke naald in een hooiberg te vinden door de hele hooiberg te verplaatsen bij elke knipoog. Wetenschappers hebben geprobeerd uit te vogelen hoe ze het model de saaie delen van het verhaal kunnen laten negeren en alleen naar de spannende delen kunnen laten kijken, maar ze hadden moeite om dit te doen zonder het vermogen te verliezen om het juiste antwoord te vinden.

Hier komt een nieuwe methode genaamd LOCKS om de hoek kijken. De onderzoekers achter dit artikel ontdekten een slimme truc om zaken te versnellen zonder de draad kwijt te raken. Ze realiseerden zich dat hoewel het hele verhaal complex is, kleine stukken ervan (genaamd "pagina's") hun eigen eenvoudige, unieke patronen hebben. In plaats van te proberen het hele verhaal te samenvatten met één gigantische, rommelige kaart, geeft LOCKS elke individuele pagina zijn eigen kleine, hoogwaardige "spectrale samenvatting".

Denk aan het volgende: Stel je voor dat je een detective bent die een mysterie oplost in een roman van 1.000 pagina's. In plaats van elk woord van elke pagina te lezen om de moordenaar te vinden, maak je een piepkleine "spiekbrief" van slechts 10% van de grootte voor elke pagina. Deze spiekbrief bevat niet alleen de woorden; het legt de vibe en de belangrijkste richtingen van de inhoud van die specifieke pagina vast. Wanneer de detective (de AI) moet weten waar hij het eerst moet kijken, leest hij niet de volledige pagina's. Hij werpt slechts een blik op deze kleine spiekbriefjes om te zien welke pagina's de meeste "aanwijzingen" (attention mass) bevatten.

Het artikel laat zien dat deze methode ongelooflijk effectief is. Door gebruik te maken van deze pagina-specifieke spiekbriefjes kan het model ongeveer 98% van de tekst overslaan in een context van 100.000 tokens, en toch vindt het het juiste antwoord bijna even goed als wanneer het alles had gelezen. Sterker nog, bij moeilijke wiskunde- en redeneertests falen andere methoden die proberen te raden welke pagina's belangrijk zijn vaak volledig, maar LOCKS houdt de "drager"-pagina's — de pagina's die daadwerkelijk het antwoord bevatten — veilig en wel.

De onderzoekers bewezen dat het proberen te gebruiken van één enkele kaart voor het hele boek (een "gedeelde" samenvatting) niet werkt, omdat verschillende pagina's verschillende geheimen hebben die verloren gaan in de mix. Ze toonden ook aan dat hun methode "training-free" is, wat betekent dat het werkt met bestaande AI-modellen zonder dat er iets opnieuw geleerd hoeft te worden. Toen ze het testten op echte hardware, ontdekten ze dat het de tijd die nodig is om elk woord te genereren voor zeer lange documenten met de helft verminderde. Het is alsof je een trage, slepende wandeling door een bibliotheek verandert in een hogesnelheidsteleportatiesysteem dat alleen stopt bij de planken die er echt toe doen.

Kortom, LOCKS lost de "long-context" bottleneck op door te beseffen dat elke pagina van een verhaal zijn eigen unieke vingerafdruk heeft. Door een compacte, pagina-specifieke samenvatting voor elke pagina te maken, kan de AI onmiddellijk weten welke pagina's ze moet lezen en welke ze moet negeren, waardoor het mogelijk wordt om met modellen te chatten over boeken die honderdduizenden woorden lang zijn zonder dat de computer overbelast raakt.

Technische Samenvatting: LOCKS – Pagina-lokale Compacte Sleutel-samenvattingen voor Efficiënte Long-Context Decoding

1. Probleemstelling

Het serveren van Large Language Models (LLM's) bij lange contextlengtes wordt steeds meer beperkt door de Key-Value (KV) cache, specif으로 door de KV-leesbandbreedte in plaats van geheugencapaciteit. In productieregimes met hoge concurrentie worden de modelgewichten over een batch geamortiseerd, maar elke aanvraag moet zijn eigen unieke KV-cache lezen. Voor modellen zoals Llama-3.1-8B vereist een context van 128K ongeveer 16 GiB aan KV-cache per aanvraag, wat de omvang van de modelgewichten overstijgt. Gevolgelijk is de cache, en niet de gewichten, de limiterende factor voor het aantal aanvragen dat een accelerator kan bedienen.

Hoewel aandacht globaal hoog-rangig is, is het lokaal laag-rangig: enkele pagina's (blokken tokens) dragen op elk decode-stap bijna de volledige aandachtsmassa van een query. Sparse decoding streeft ernaar alleen aandacht te besteden aan deze kritieke pagina's. Bestaande selectiemethoden kampen echter met een trade-off:

Exacte selectie (elke sleutel lezen) heft het doel van sparsity op door de volledige bandbreedtekosten te veroorzaken.
Residentie-samenvattingen (compacte representaties van pagina's) zijn vereist om het lezen van volledige sleutels te vermijden, maar bestaande benaderingen falen vaak in het behouden van de specifieke geometrische structuur die nodig is om "carrier"-pagina's (die het antwoord bevatten) te identificeren onder agressieve budgetbeperkingen. Gedeelde projecties over pagina's heen zijn bewezen blind voor pagina-specifieke inhoud, wat leidt tot taalklap (task collapse) in retrieval- en redeneerscenario's.

2. Methodologie: LOCKS

LOCKS (Page-Local Compact Key Summaries) pakt dit aan door de volledige KV-cache resident te houden, maar het lezen van de meeste sleutels tijdens het decoderen te vermijden. Het werkt met een pagina-lokale spectrale samenvatting die training-vrij is en fungeert als een drop-in plugin voor vLLM.

Kernmechanisme

Pagina-lokale Spectrale Decompositie:
- In plaats van een gedeelde globale of sequentie-niveau basis te gebruiken, construeert LOCKS een unieke laag-rangige basis voor elke pagina van de KV-cache.
- Wanneer een pagina wordt afgerond (tijdens prefill of een page-crossing refresh), ondergaat deze een kleine eigendecompositie van zijn sleutel-deviaties (gecentreerde sleutels).
- De samenvatting slaat de top- $r$ eigenvectoren (basis) en de bijbehorende coëfficiënten voor elke token in de pagina op.
- Kwantisatie: De basis wordt opgeslagen in int4 en coëfficiënten/centroids in int8. Bij een rang $r=8$ en paginagrootte $B=16$ , is de samenvatting ongeveer 10% van de oorspronkelijke KV-bytes van de pagina.
Reconstructie en Selectie:
- Bij elke decode-stap reconstrueert LOCKS de binnen-pagina logits voor de query met behulp van de compacte samenvatting:
  $\hat{s}_j(q) = \log \sum_{i \in P_j} \exp\left( \frac{q^\top \mu_j + (V_j^\top q)^\top c_i}{\sqrt{d}} \right)$
  waarbij $\mu_j$ het pagina-centroid is, $V_j$ de pagina-lokale basis, en $c_i$ de coëfficiënten.
- Het berekent de Log-Sum-Exp (LSE) van deze gereconstrueerde logits om de aandachtsmassa van de pagina te schatten.
- Selectie vindt plaats door pagina's te rangschikken op basis van deze geschatte massa's. Cruciaal is dat geen enkele kandidaat-sleutel of -waarde wordt gelezen tijdens deze selectiefase.
Group-Shared Selection (GSA):
- Voor Grouped-Query Attention (GQA) middelt LOCKS de geschatte massadeel over de query-heads in een groep om één gedeelde selectietabel te vormen. Dit zorgt ervoor dat alle heads in een groep naar dezelfde set pagina's kijken, wat verkeersfragmentatie voorkomt.
Implementatie:
- Geïmplementeerd als een plugin voor onveranderde vLLM.
- De gehele decode-stap (scoring, selectie en sparse attention) draait binnen volledige CUDA-graphs, wat zorgt voor nul overhead door Python-dispatch en efficiënte kernel-fusie mogelijk maakt.
- De constructie van de samenvatting is ontkoppeld van het kritieke pad van het decoderen en vindt alleen plaats wanneer pagina's worden afgerond.

3. Belangrijkste Bijdragen

Theoretische Bevindingen

Localiteit van Structuur: Het paper toont aan dat de geometrische structuur die relevant is voor pagina-selectie geconcentreerd is op de pagina-scope, niet op de sequentie- of cache-scope. De eigenbasis van een pagina zelf (bij rang-8) vangt de meerderheid van de energie van de sleutel op, terwijl gedeelde bases (zelfs wanneer ze een vergelijkbare staat-omvang hebben) er niet in slagen de pagina-specifieke richtingen te vangen.
Onmogelijkheid van Gedeelde Projecties: De auteurs bewijzen (Propositie 1) dat elke vaste, query-onafhankelijke gedeelde laag-rangige lineaire projectie inhoud-blind is voor een $(d-\rho)$ -dimensionale familie van binnen-pagina verschillen. Dit betekent dat gedeelde bases niet in staat zijn om onderscheid te maken tussen pagina's die alleen verschillen in specifieke inhoudsrichtingen, wat leidt tot selectiefouten.
Retentiegaranties: Het paper biedt een samengestelde theoretische grens (Theorem 1) die aantoont dat de fout in de aandacht-output direct wordt gecontroleerd door de spectrale staartenergie van de pagina-lokale reconstructie. De pagina-lokale basis wordt bewezen als residu-optimaal onder orthogonale rang- $r$ reconstructies.

Empirische Resultaten

Retrieval en QA (LongBench-v1, RULER): LOCKS volgt de "lees-elke-sleutel" oracle (exacte LSE-selectie) tot aan de kleinste budgetten. Op retrieval-intensieve taken (RULER) presteert het aanzienlijk beter dan state-of-the-art selectors (Quest, KVzip, ShadowKV, RocketKV) op het gebied van carrier retention (het behouden van de pagina's die het antwoord bevatten), wat de primaire drijfveer is voor taaksucces.
Long-Form Reasoning (AIME26, MATH-500): In redeneertaken waarbij het model duizenden tokens genereert, bezwijken baseline selectors vaak bij strikte budgetten. LOCKS behoudt een hoge nauwkeurigheid en evenaart de FullKV-prestaties bij 100K+ contexten, terwijl het slechts naar ~2% van de tokens kijkt.
Efficiëntie:
- Bij een contextlengte van 1M bereikt LOCKS een 2.0× versnelling in per-token latentie vergeleken met dense attention (FlashAttention-3/FlashInfer).
- Het vermindert het aantal gelezen KV-bytes per decode-stap met een factor 9.8× bij 1M context.
- De throughput-versnelling schaalt met de batchgrootte en bereikt 1.80× bij 256K context met grote batches.

4. Betekenis en Claims

Het paper claimt dat pagina-lokale spectrale concentratie de ontbrekende schakel is voor efficiënte sparse decoding. Door de scope van representatie te verschuiven van globaal/sequentie naar lokaal (per pagina), behoudt LOCKS de noodzakelijke "carrier"-pagina's voor taakvoltooiing zonder dat training of aanpassing van de modelarchitectuur vereist is.

Drop-in Integratie: In tegenstelling tot veel eerdere werken die model-forks of training vereisen, werkt LOCKS met onveranderde vLLM en standaard GQA-modellen.
Theoretisch-Praktische Afstemming: De methode is een van de weinige die een theoretische grens op reconstructiefout direct verbindt met end-to-end taakkwaliteit en het overleven van de 'carrier', wat valideert dat de "keten" van spectrale concentratie naar taakprestaties standhoudt.
Schaalbaarheid: De aanpak schaalt naar extreme contexten (1M+ tokens) waar geheugenbandbreedte de primaire bottleneck is, wat een praktisch pad biedt voor het serveren van long-context modellen zonder offloading naar CPU of getierde geheugensystemen.

De auteurs blijven bescheiden over beperkingen: de retentie van de toegepaste gekwantiseerde samenvatting wordt empirisch gemeten in plaats van strikt per stap begrensd, en de methode gaat momenteel uit van een vaste paginagrootte. Desalniettemin vestigen de resultaten dat een training-vrije, pagina-lokale spectrale samenvatting voldoende is om de oracle-selectie te volgen tot aan agressieve sparsity-budgetten, wat praktische versnelling van long-context mogelijk maakt.