CoMeT: Collaborative Memory Transformer for Efficient Long Context Modeling

AI summary

Overview

Research area: Efficient long-context modeling for Large Language Models (LLMs) — specifically, reducing the memory and compute cost of processing very long sequences.

Technical level: Intermediate (requires familiarity with Transformer attention, KV caches, and basic memory architectures, but the core ideas are explained accessibly).

Scope: The paper presents CoMeT, a plug-in dual-memory module that lets pre-trained LLMs process arbitrarily long contexts with constant memory and linear time, validated on academic benchmarks, real-world agent tasks, and 1M-token passkey retrieval.

What This Paper Is About

Standard Transformers scale poorly with context length: attention costs grow quadratically, and the key-value (KV) cache grows linearly, making million-token inputs prohibitively expensive. Existing fixes either compress context (which still grows with input length) or use fixed-size memory (which forgets important details and lacks gating). CoMeT's goal is to solve both problems at once — retaining both long-range facts and recent fine-grained details — while remaining a drop-in module for existing pre-trained LLMs.

Key Contributions

1. A dual-memory architecture. CoMeT combines a fixed-size global memory with a gated update rule (for long-range dependencies) and a temporary memory managed by a FIFO queue (for high-fidelity recent context). These memories act as a dynamic soft prompt prepended to each incoming chunk.

2. A lightweight state-to-memory transform. A Residual Low-Rank Adapter (RLA) with rank 8 converts the persistent global state into memory tokens, adding only 3.95M parameters (0.098%) to Qwen3-4B-Instruct-2507 — small enough to keep performance stable.

3. Layer-level pipeline parallelism for long-context training. Instead of having GPU workers wait for a full chunk computation, the method transmits memory states layer-by-layer, interleaving communication with computation and yielding a 2.7× speedup over naive context parallelism. This makes fine-tuning at 128k context feasible on 16×80GB GPUs.

4. Strong empirical results across three settings. Academic benchmarks (Scrolls), real-world tasks (user behavior QA and a terminal agent task), and extreme-length passkey retrieval, all with competitive or superior performance relative to full attention at a fraction of the cost.

Main Findings

• Extreme length extrapolation from short training: Trained only on 32k-length sequences, CoMeT retrieves a passkey from any position in a 1M-token context, with a 21× inference speedup and 10× smaller memory footprint than full attention at that length.

• Competitive with full attention on Scrolls: CoMeT achieved the highest average score among efficient methods (40.10), and on summarization tasks (GovReport, SummScreenFD) matched the fine-tuned full-attention baseline using only ~2.5k tokens of memory.

• Broad plug-and-play generalization: Evaluated on Qwen3-4B (95.0% of full attention), Qwen3-14B (97.6%), and Llama-3.1-8B-Instruct, where CoMeT actually surpassed the fine-tuned full-attention baseline by 4.6%. Gains scale with model size.

• Real-world utility: On a real e-commerce user behavior QA task, CoMeT beat the xRAG baseline by 2.7 accuracy points and a 4k-truncation baseline by 27.4 points. On Terminal-Bench (128k-token agent trajectories), it was competitive with full attention while training far more efficiently.

• Constant memory, linear time: Inference uses a flat ~10GB of GPU memory regardless of context length, with per-token decode latency stable at ~22ms, versus full attention's linear growth and out-of-memory failure at 128k.

• Memory roles are distinct: Temporary memory drives performance on in-domain lengths (saturating at 2,048 tokens), while global memory is decisive for extrapolation. Removing the gate on global memory causes complete failure on 1M-token retrieval.

• Gating visualization: In deep layers (e.g., layer 33), gate values drop to zero exactly when the passkey is encountered, then stay at 1 to shield that information from being overwritten. Other layers and states show varied forgetting rates, producing multi-scale memory retention.

Methodology in Plain English

The input is processed in chunks rather than all at once. At each Transformer layer and for each chunk, the model prepends two sets of memory tokens to the chunk's hidden states: a global memory (a persistent state vector transformed by a small low-rank adapter) and a temporary memory (a rolling FIFO queue of compressed representations of the most recent chunks). It also interleaves "compression tokens" inside the chunk to capture fine-grained local information, and appends "readout tokens" at the end to summarize what was just seen.

All tokens interact through standard causal self-attention, so the current chunk can read from both memories. After processing, the readout tokens are normalized and mixed into the global state via a sigmoid gate — the gate decides how much old information to keep versus overwrite. The compression tokens are transformed and pushed into the FIFO queue, discarding the oldest entry. This keeps memory bounded no matter how long the input is.

To train this on very long contexts, the authors split chunks across GPUs and send memory states between workers layer by layer rather than chunk by chunk, so downstream GPUs can start work as soon as they receive the first layer's state instead of idling until the entire previous chunk finishes.

Why This Matters

Impact on research: CoMeT shows that a lightweight, plug-in memory module can achieve asymptotic efficiency gains (O(N) time, O(1) space) without retraining a model from scratch — an important practical path forward, since recurrent alternatives like Mamba and RetNet cannot be grafted onto existing LLMs. It also demonstrates that gating and stratified memory (long-term vs. recent) matter more than raw compression capacity.

Real-world applications:

• Long-document analysis — summarizing or querying books, legal contracts, and technical manuals without truncation.
• Personalized assistants — reasoning over months or years of a user's interaction history.
• Codebase and agent workflows — comprehending large repositories or long tool-use trajectories within a single context window.
• Enterprise knowledge work — question answering over aggregated internal corpora on modest hardware.

Industry relevance: Because the module adds under 0.1% parameters and is compatible with existing pre-trained backbones, it lowers the barrier for deploying long-context capability on limited GPU budgets. The 128k-context training recipe on 16 GPUs is directly actionable for teams that cannot afford massive full-attention long-context training runs.

Future Directions

1. Episodic memory integration. The authors note their framework does not yet incorporate test-time training or episodic memory, which could allow the model to update weights in response to new information.

2. External memory interfaces. Coupling CoMeT with notebooks, retrieval-augmented knowledge bases, or other external stores could extend its capacity beyond the fixed global state.

3. Longer contexts and more scales. Extending beyond 1M tokens and validating on larger backbones (their largest test was Qwen3-14B) would clarify how the architecture behaves at frontier scale.

4. Privacy safeguards. The authors explicitly raise the risk that efficient processing of long user behavior logs could enable surveillance, and call for regulation-aligned deployment practices — an open problem for practitioners.

Target Audience

Researchers and engineers working on efficient LLM inference, long-context modeling, or memory architectures will get the most value. It is also useful for practitioners deploying long-context applications — agent frameworks, document QA, personalization — who need a plug-in solution without retraining a model from scratch. Readers should be comfortable with Transformer attention, KV caching, and chunked sequence processing, though the paper's core mechanisms are explained clearly enough for a motivated newcomer at the intermediate level.

Authors’ abstract

The quadratic complexity and indefinitely growing key-value (KV) cache of standard Transformers pose a major barrier to long-context processing. To overcome this, we introduce the Collaborative Memory Transformer (CoMeT), a novel architecture that enables LLMs to handle arbitrarily long sequences with constant memory usage and linear time complexity. Designed as an efficient, plug-in module, CoMeT can be integrated into pre-trained models with only minimal fine-tuning. It operates on sequential data chunks, using a dual-memory system to manage context: a temporary memory on a FIFO queue for recent events, and a global memory with a gated update rule for long-range dependencies. These memories then act as a dynamic soft prompt for the next chunk. To enable efficient fine-tuning on extremely long contexts, we introduce a novel layer-level pipeline parallelism strategy. The effectiveness of our approach is remarkable: a model equipped with CoMeT and fine-tuned on 32k contexts can accurately retrieve a passkey from any position within a 1M token sequence. On the SCROLLS benchmark, CoMeT surpasses other efficient methods and achieves performance comparable to a full-attention baseline on summarization tasks. Its practical effectiveness is further validated on real-world agent and user behavior QA tasks. The code is available at: https://github.com/LivingFutureLab/Comet