CoMeT: Collaborative Memory Transformer for Efficient Long Context Modeling

TL;DR

CoMeT introduces a plug-in Collaborative Memory Transformer that achieves constant memory usage and linear time for long-context modeling by coupling a global memory with a gated update and a temporary FIFO memory for recent details. The architecture enables processing contexts up to 1M tokens, with strong extrapolation demonstrated via passkey retrieval and competitive performance on Scrolls and real-world tasks compared to full attention. A layer-level pipeline parallelism strategy further enables efficient fine-tuning on very long sequences, achieving substantial speedups. Overall, CoMeT makes practical long-context processing feasible for LLMs, with demonstrated efficiency, effectiveness, and flexibility for future memory extensions.

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://anonymous.4open.science/r/comet-B00B/

Figures (11)

• Figure 1: CoMeT is trained on the passkey task munkhdalai2024leave (i.e., the Needle-in-a-Haystack test) with a 32k context, yet it can retrieve a passkey from any position within a 1M-token context. Moreover, its inference time scales linearly with the context length, while GPU memory usage remains constant.
• Figure 2: Overview of CoMeT. At layer $i$, the global memory $\mathbf{G}^i_{\tau}$ and temporary memory $\mathbf{T}^i_{\tau}$ are prepended to the current chunk's hidden states $\mathbf{H}^i_{\tau}$. Compression tokens $\mathbf{C}^i_{\tau}$ are interleaved within the hidden states for fine-grained information capture, while readout tokens $\mathbf{R}^i_{\tau}$ are appended at the end to distill key information for updating the global state. All tokens interact through causal self-attention, enabling the model to retrieve relevant historical information while processing the current chunk.
• Figure 3: Architecture of the global memory mechanism. At each layer $i$ and chunk $\tau$, the global state $\mathbf{S}^i_{\tau}$ is transformed by a $\mathrm{RLA}$ to produce the global memory $\mathbf{G}^i_{\tau}$. The state is then updated for the next chunk via a gating mechanism that selectively integrates information from the normalized readout tokens $\mathbf{R}^{i+1}_{\tau}$.
• Figure 4: The architecture of the temporary memory mechanism. CoMeT employs a fixed-capacity FIFO queue to manage compressed representations of recent chunks. As new information from the current chunk is enqueued, the oldest memory entry is discarded. This rolling window of memory provides the model with a high-resolution view of the most recent context while maintaining a constant memory footprint.
• Figure 5: Naive context parallelism. Workers process chunks sequentially. Worker $j+1$ must wait for worker $j$ to complete its entire computation before starting, creating a large pipeline bubble and leading to significant resource under-utilization.