CoMeT: 효율적인 장문 컨텍스트 모델링을 위한 협업 메모리 트랜스포머

CoMeT: Collaborative Memory Transformer for Efficient Long Context Modeling

표준 트랜스포머의 2차 복잡도와 무한정으로 증가하는 키-값(KV) 캐시는 장문 컨텍스트 처리에 큰 걸림돌이 됩니다. 이를 극복하기 위해, 본 논문에서는 협업 메모리 트랜스포머(CoMeT)라는 새로운 아키텍처를 소개합니다. CoMeT는 LLM이 일정한 메모리 사용량과 선형 시간 복잡도로 임의의 길이의 시퀀스를 처리할 수 있도록 합니다. 효율적인 플러그인 모듈로 설계된 CoMeT는 최소한의 파인튜닝만으로 사전 훈련된 모델에 통합될 수 있습니다. CoMeT는 시퀀스 데이터 청크를 처리하며, 듀얼 메모리 시스템을 사용하여 컨텍스트를 관리합니다. 여기에는 최근 이벤트에 대한 FIFO 큐를 사용하는 임시 메모리와 장거리 의존성에 대한 게이트 업데이트 규칙을 사용하는 글로벌 메모리가 포함됩니다. 이러한 메모리는 다음 청크에 대한 동적 소프트 프롬프트 역할을 합니다. 매우 긴 컨텍스트에 대한 효율적인 파인튜닝을 위해, 새로운 레이어 레벨 파이프라인 병렬 처리 전략을 도입했습니다. 제안하는 방식의 효과는 매우 뛰어납니다. CoMeT를 탑재하고 32k 컨텍스트로 파인튜닝된 모델은 1M 토큰 시퀀스 내의 임의의 위치에서 정확하게 패스키를 검색할 수 있습니다. SCROLLS 벤치마크에서 CoMeT는 다른 효율적인 방법들을 능가하며, 요약 작업에서 전체 어텐션 기반 모델과 비교 가능한 성능을 달성합니다. 또한, 실제 에이전트 및 사용자 행동 질의응답 작업에서 실용적인 효과가 검증되었습니다. 코드는 다음 주소에서 확인할 수 있습니다: https://anonymous.4open.science/r/comet-B00B/

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/