ゲート付き連想メモリ：効率的なシーケンスモデリングのための並列O(N)アーキテクチャ

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

August 30, 2025

著者: Rishiraj Acharya

cs.AI

要旨

Transformerアーキテクチャは、自己注意機構を基盤としており、シーケンスモデリングタスクにおける事実上の標準となっています。しかし、その中核となる計算プリミティブはシーケンス長に対して二次的にスケールする（O(N^2)）ため、長いコンテキストを処理する際に重大なボトルネックを生み出します。本論文では、シーケンス長に対して線形の計算量（O(N)）を実現する、新規の完全並列アーキテクチャであるGated Associative Memory（GAM）ネットワークを提案します。GAMブロックは、自己注意層を2つの並列パスに置き換えます。1つは局所的で位置依存のコンテキストを効率的に捕捉する因果畳み込み、もう1つはグローバルで内容ベースのパターンをモデル化する並列連想メモリ検索機構です。これらのパスはゲーティング機構によって動的に融合され、各トークンに対して局所情報とグローバル情報を柔軟に組み合わせることが可能になります。我々はGAMをゼロから実装し、WikiText-2ベンチマークにおいて標準Transformerモデルと最新の線形時間ベースライン（Mamba）に対して、またTinyStoriesデータセットにおいてTransformerに対して厳密な比較分析を行いました。実験の結果、GAMは一貫して高速であり、トレーニング速度において両ベースラインを上回り、全てのデータセットにおいて優れたまたは競争力のある最終検証パープレキシティを達成し、シーケンスモデリングにおける有望で効率的な代替手段として確立されました。

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.