Adaptive KV Memory for Continual Learning and Long Reasoning

Two workstreams on adaptive KV cache memory: gating compressed memories for continual learning, and online KV cache compaction for long reasoning.

Mentors:
Alham Fikri Aji Farid Adilazuarda M. Reza Qorib
Mentees (5):
Faeyza Rishad ArdiIzaaz InharPhudish PrateepamornkulQuang Minh NguyenTri Q. Vo

Project Proposal

Long-context inference and long reasoning are expensive because the key-value (KV) cache grows with sequence length. This project studies how to keep KV memory small while it stays useful, in two workstreams. Both papers were accepted at the COLM 2026 Workshop on Context Beyond the Window.

Workstream 1: Gating compressed KV memory

Compressed memories such as Cartridges distill a long context into a small, trainable KV cache. Real deployments, however, see new corpora arrive continually, and a memory with a fixed budget must absorb each one while retaining what it already holds.

This workstream studies gating mechanisms that govern what enters a compressed memory and where it is stored: a statistical TF-IDF gate that updates only the slots a new corpus attends to, and a learned mixture-of-experts gate that routes each query to a small subset of expert sub-caches. Gated memories improve both forgetting and acquisition over a dense baseline of equal capacity, and routing helps mainly by limiting which parts of the cache can change during an update.

Team: Quang Minh Nguyen, Phudish Prateepamornkul, Tri Q. Vo.

Workstream 2: Adaptive online KV cache compaction

Long Chain-of-Thought reasoning makes KV cache growth a severe memory and latency bottleneck, and existing learnable cache compressions are typically offline with rigid, uniform compression ratios.

This workstream proposes an online compaction method that combines token merging with learnable token dropping. It extends the STILL compactor with Gumbel-Sigmoid gating to prune unimportant latent tokens from a dense, merged initialization, which allows heterogeneous cache allocation across layers and heads. On Qwen3-1.7B it outperforms vanilla STILL and DMS on AIME 26 and HMMT-Feb 26, reduces peak KV payload by about 83%, and improves throughput by more than 22% over vanilla STILL.

Team: Izaaz Inhar, Faeyza Rishad Ardi.

Relevant publications:

  • Empirical Study on Updating Key-Value Memories in Transformer Feed-Forward Layers
  • Memory Layers at Scale
  • Continual Learning via Sparse Memory Finetuning

Publications

  1. Beyond Merging and Dropping: Adaptive Online KV Cache Compaction
    Izaaz Inhar, Faeyza Rishad Ardi, Muhammad Reza Qorib, Alham Fikri Aji, Farid Adilazuarda
    Context Beyond the Window Workshop @ COLM • 2026
    Abstract

    As Large Language Models (LLMs) increasingly rely on long Chain-of-Thought reasoning, linear Key-Value (KV) cache growth creates severe memory and latency bottlenecks. Existing learnable cache compressions are typically offline and enforce rigid, uniform compression ratios. To address this, we propose an adaptive online compaction method that combines token merging with learnable token dropping. By extending the STILL compactor with Gumbel-Sigmoid gating, our approach dynamically prunes unimportant latent token representations from a dense, merged initialization. This hybrid mechanism enables optimal, heterogeneous cache allocation across layers and heads while retaining reasoning capabilities. Trained sequentially on Open-R1, our compactor reaches stable 8x compression while keeping teacher–student divergence bounded across repeated sequential compaction. On Qwen3-1.7B, our method achieves 27.50% Avg@4 on AIME 26 and 19.69% on HMMT-Feb 26, outperforming vanilla STILL by 2.50 and 0.76 percentage points and DMS by 1.67 and 1.51 points, respectively. It reduces peak KV payload by approximately 83% relative to uncompressed inference and improves throughput by more than 22% over vanilla STILL.

  2. Gating Mechanisms for Continually Adapting Compressed KV Memory
    Quang Minh Nguyen, Phudish Prateepamornkul, Tri Q. Vo, Muhammad Reza Qorib, Alham Fikri Aji, Farid Adilazuarda
    Context Beyond the Window Workshop @ COLM • 2026
    Abstract

    Serving large language models over long contexts is expensive because cost scales with the key–value (KV) cache, and accuracy degrades when relevant evidence is buried among distracting tokens. Compressed memories address this problem by distilling a long context into a small, trainable KV cache. Cartridges learn such a cache through a synthetic question-answering objective called self-study. Real deployments, however, are not static: new corpora arrive continually, and a memory with a fixed budget must absorb each one while retaining what it already holds. We study this problem through gating mechanisms that govern what enters a compressed memory and where it is stored. We introduce two mechanisms for the dense Cartridge: a statistical TF–IDF gate that updates only the slots attended to by a new corpus but not by an earlier one, and a learned mixture-of-experts gate that routes each query to a small subset of expert sub-caches. Across two instruction-tuned models (Llama-3.2-3B and Qwen3-4B) in a continual setting, the best routed configurations improve both forgetting and acquisition relative to a dense baseline of equal capacity. However, router analysis shows little corpus-specific specialization, suggesting that routing helps mainly by limiting which parts of the cache can change during an update.

    PDF