게이트 연관 메모리: 효율적인 시퀀스 모델링을 위한 병렬 O(N) 아키텍처

Gated Associative Memory: A Parallel O(N) Architecture for Efficient Sequence Modeling

August 30, 2025

저자: Rishiraj Acharya

cs.AI

초록

자기 주의 메커니즘(self-attention mechanism)을 기반으로 하는 트랜스포머(Transformer) 아키텍처는 시퀀스 모델링 작업에서 사실상의 표준이 되었습니다. 그러나 이 아키텍처의 핵심 계산 원리는 시퀀스 길이에 대해 제곱적으로 증가하는(O(N^2)) 복잡도를 가지며, 이는 긴 문맥을 처리하는 데 있어 상당한 병목 현상을 일으킵니다. 본 논문에서는 시퀀스 모델링을 위한 완전히 병렬적인 새로운 아키텍처인 Gated Associative Memory(GAM) 네트워크를 제안합니다. 이 아키텍처는 시퀀스 길이에 대해 선형적인 복잡도(O(N))를 보입니다. GAM 블록은 자기 주의 레이어를 두 개의 병렬 경로로 대체합니다: 하나는 지역적이고 위치에 의존적인 문맥을 효율적으로 포착하기 위한 인과적 컨볼루션(causal convolution)이며, 다른 하나는 전역적이고 내용 기반 패턴을 모델링하기 위한 병렬 연관 메모리 검색 메커니즘입니다. 이 두 경로는 게이팅 메커니즘을 통해 동적으로 융합되어, 모델이 각 토큰에 대해 지역적 및 전역적 정보를 유연하게 결합할 수 있도록 합니다. 우리는 GAM을 처음부터 구현하고, WikiText-2 벤치마크에서 표준 트랜스포머 모델과 현대적인 선형 시간 기반선(Mamba)에 대해 엄격한 비교 분석을 수행하였으며, TinyStories 데이터셋에서도 트랜스포머와 비교 실험을 진행했습니다. 실험 결과, GAM은 일관적으로 더 빠른 속도를 보이며, 훈련 속도에서 두 기반선을 모두 능가했고, 모든 데이터셋에서 우수하거나 경쟁력 있는 최종 검증 혼란도(perplexity)를 달성하여 시퀀스 모델링을 위한 유망하고 효율적인 대안으로 자리매김했습니다.

English

The Transformer architecture, underpinned by the self-attention mechanism, has become the de facto standard for sequence modeling tasks. However, its core computational primitive scales quadratically with sequence length (O(N^2)), creating a significant bottleneck for processing long contexts. In this paper, we propose the Gated Associative Memory (GAM) network, a novel, fully parallel architecture for sequence modeling that exhibits linear complexity (O(N)) with respect to sequence length. The GAM block replaces the self-attention layer with two parallel pathways: a causal convolution to efficiently capture local, position-dependent context, and a parallel associative memory retrieval mechanism to model global, content-based patterns. These pathways are dynamically fused using a gating mechanism, allowing the model to flexibly combine local and global information for each token. We implement GAM from scratch and conduct a rigorous comparative analysis against a standard Transformer model and a modern linear-time baseline (Mamba) on the WikiText-2 benchmark, as well as against the Transformer on the TinyStories dataset. Our experiments demonstrate that GAM is consistently faster, outperforming both baselines on training speed, and achieves a superior or competitive final validation perplexity across all datasets, establishing it as a promising and efficient alternative for sequence modeling.