CoMeT: Collaborative Memory Transformer for Efficient Long Context Modeling
This addresses a major bottleneck for deploying large language models in real-world applications requiring long-context understanding, such as agent tasks and user behavior QA, though it builds incrementally on existing memory-augmented approaches.
The paper tackles the problem of quadratic complexity and growing memory requirements in Transformers for long-context processing by introducing CoMeT, a plug-in module that enables handling of arbitrarily long sequences with constant memory and linear time. It achieves accurate retrieval from 1M token sequences and competitive performance on benchmarks like SCROLLS.
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/