Technische Samenvatting: CoMem (Comprehension Memory)

Probleemstelling

Long-context taalmodellering staat voor twee primaire knelpunten: de kwadratische prefill-kosten van self-attention en de lineaire geheugengroei van de Key-Value (KV) cache. Bestaande efficiënte methoden pakken dit doorgaans aan door KV-toestanden te verwijderen of te comprimeren langs de token-as (een subset van tokens behouden) of door externe tekstunits op te halen. Deze benaderingen worstelen echter vaak met het balanceren van geheugenefficiëntie met de getrouwheid (fidelity) die vereist is voor complexe redenering over lange contexten. Bovendien vertrouwen veel methoden op het trainen van gespecialiseerde geheugennetwerken of het comprimeren van de volledige context, wat aanzienlijke overhead of verlies van getrouwheid kan introduceren.

De auteurs stellen dat Transformer-lagen een taakverdeling vertonen: semantische informatie wordt toegankelijk in de lagere en middelste lagen, terwijl de hogere lagen steeds meer specialiseren in representaties voor de huidige query en voorspelling. Huidige methoden maken geen volledig gebruik van deze diepte-gebaseerde structuur om geheugen te organiseren.

Methodologie: CoMem

CoMem (Comprehension Memory) reorganiseert long-context geheugen langs de diepte-as in plaats van uitsluitend de token-as. Het combineert gebonden tokenselectie met diepte-gepartitioneerde computatie-hergebruik.

Kernmechanisme

1. Schrijffase (Diepte-partitionering):

   • De input-context wordt gestreamd in disjuncte chunks (bijv. 512 tokens).
   • Elke chunk wordt alleen verwerkt door een specifieke intermediaire laag $j$ (de "split depth").
   • In plaats van volledige diepte KV-toestanden op te slaan, cachet CoMem de residuele tensor $h_j$ (de hidden state) op laag $j$ voor elke token in de chunk.
   • Deze tensoren worden opgeslagen in een externe index met de chunk-tekst als sleutel.
   • Opslagefficiëntie: Voor Qwen3-8B vereist het opslaan van $h_j$ slechts 1/18de van de bytes van een volledige bf16 KV-paar per token.

2. Leesfase (Gebonden Retrieval & Herberekening):

   • Bij ontvangst van een query selecteert een selector (met behulp van BM25) een gebonden set van top-$k$ relevante chunks.
   • De gecachte residuele tensoren $h_j$ voor deze chunks worden samengepakt met een "sink" token en de huidige query.
   • Het model herrekent alleen de hogere lagen $[j:L]$ over deze vaste grootte pack met behulp van volledige cross-chunk attention.
   • Complexiteit: Voor een vastgesteld retrieval-budget ($k$) en chunk-grootte zijn de model-zijde online read-compute en het geheugen werkset onafhankelijk van de totale lengte van de opgeslagen context.

Training en Adaptatie

• Backbone: Het basismodel (Qwen3-8B) blijft bevroren.
• Zelf-distillatie: Om het verlies van getrouwheid te beperken dat wordt veroorzaakt door het overslaan van lagere lagen tijdens de leesfase, wordt een lichtgewicht LoRA (Rank-32) getraind via zelf-distillatie. Het student-model (bij split $j=12$) leert de logits van een bevroren teacher-model (bij $j=0$) te matchen op de top-64 logit support.
• Trainingsdata: De LoRA wordt getraind op plain PG-19 tekst zonder taakspecifieke supervisie of retrieval-data, waardoor het model in staat is de "diepte split"-capaciteit te generaliseren.

Belangrijkste Bijdragen

1. Diepte als Hergebruik-as: Het paper stelt vast dat semantische toegankelijkheid piekt nabij de middelste diepte, terwijl de "zero-shot readable boundary" (waar gecachte toestanden direct bruikbaar zijn zonder herberekening) afhankelijk is van de model-schaal. CoMem behandelt deze diepte split ($j$) als een instelbare knop tussen vooraf berekende arbeid en query-tijd kosten.
2. Gebonden Diepte-gepartitioneerd Geheugen: CoMem cacheert één residuele staat per token en hervat exact de bovenliggende transformer. Het combineert gebonden retrieval met volledige cross-chunk attention, waarbij een opslagvoetafdruk van 1/18de van de volledige KV-caches per token wordt bereikt op Qwen3-8B.
3. Gecontroleerde Kwaliteit-Efficiëntie Afweging: De auteurs leveren bewijs dat diepere caching de per-query computation vermindert maar een verlies aan getrouwheid veroorzaakt, wat door zelf-distillatie substantieel wordt hersteld. De methode isoleert de voordelen van retrieval van de kosten van diepte-hergebruik.

Experimentele Resultaten

Geëvalueerd op een uniform, chat-template-vrij protocol met behulp van Qwen3-8B:

• Benchmark Prestaties:
  • RULER: CoMem behaalt 97.05, waarmee het de full-context KV-Direct baseline (78.80) verslaat.
  • LoCoMo (Dialooggeheugen): CoMem scoort 38.27 versus 34.59 voor KV-Direct. Dit voordeel blijft bestaan onder conversatie-cluster resampling en onafhankelijke judge audits.
  • LongEval & LongBench: CoMem evenaart of overtreft de baselines (69.0 op LongEval; 12.15 macro F1 op LongBench).
• Efficiëntie (NVIDIA H20 bij 128k context):
  • Geheugen: CoMem gebruikt 18.26 GB piekgeheugen vergeleken met 89.36 GB voor de full-context KV-Direct.
  • Snelheid: CoMem behaalt een 7.83× versnelling in de full-write-inclusive prefill tijd.
  • Decode: De resumed-band decode blijft binnen 11% van dense inference tot 64k en is 1.07× sneller bij 128k.
• Ablaties:
  • Retrieval vs. Diepte: Gebonden selectie alleen (zonder diepte-hergebruik) verbetert het dialooggeheugen, maar diepte-hergebruik is cruciaal voor efficiëntie.
  • Distillatie: Zelf-distillatie herstelt ~13.75 LoCoMo punten die verloren gingen door de bevroren split bij $j=12$.
  • Cross-Chunk Attention: Volledige cross-chunk attention is essentieel voor multi-fact disambiguatie; block-diagonal hergebruik veroorzaakt significante prestatiedalingen.

Betekenis en Claims

Het paper claimt dat long-context geheugen effectief georganiseerd kan worden langs de laag-as, niet alleen de token-as. CoMem demonstreert dat:

• Semantische informatie toegankelijk is op intermediaire dieptes, wat het mogelijk maakt om query-blinde toestanden te cachen die later gespecialiseerd worden door de bovenliggende lagen te herberekenen.
• Een gebonden retrieval-strategie in combinatie met diepte-gepartitioneerde computation de online inference kosten kan ontkoppelen van de opgeslagen contextlengte.
• Lichtgewicht zelf-distillatie de inherente verlies aan getrouwheid door het overslaan van lagere lagen kan herstellen, wat hoge nauwkeurigheid mogelijk maakt zonder de backbone te hertrainen.

De auteurs benadrukken dat CoMem een efficiëntiemethode is die een kwaliteit-geheugen-compute afweging biedt, in plaats van een state-of-the-art nauwkeurigheidsoplossing voor elke in-window taak. Ze stellen expliciet dat hun resultaten op contexten die de native window (40.960 tokens) overschrijden "onverlengde stress-referenties" zijn en geen superioriteit claimen ten opzichte van modellen die gebruikmaken van position-extended technieken zoals YaRN. De methode wordt gepresenteerd als een complementaire benadering van bestaande token-as compressie en retrieval methoden.