arXiv:2509.00605cs.CLcs.LG2025-09

提出线性复杂度的GAM模型,解决Transformer长序列计算慢问题

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

  • 用并行路径替代注意力:局部卷积+全局记忆检索
  • 训练速度比Transformer快,验证困惑度更优或相当
  • 适合长序列建模,尤其对效率敏感的应用

Transformer凭借自注意力机制成为序列建模的标准架构,但其核心计算复杂度随序列长度呈平方增长(O(N²)),在处理长上下文时存在显著瓶颈。本文提出门控关联记忆(Gated Associative Memory, GAM)网络,一种全新的全并行序列建模架构,具有与序列长度线性相关的复杂度(O(N))。GAM模块用两条并行路径替代自注意力层:因果卷积用于高效捕捉位置相关的局部上下文,关联记忆检索机制用于建模内容相关的全局模式。两条路径通过门控机制动态融合,实现每个词元对局部与全局信息的灵活组合。我们从零实现GAM,并在WikiText-2基准和TinyStories数据集上与标准Transformer及现代线性时间基线Mamba进行严格对比。实验表明,GAM在训练速度上持续领先,且在所有数据集上的最终验证困惑度均优于或相当,证实其作为高效序列建模替代方案的潜力。

原文摘要 · Abstract (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.

序列建模线性复杂度Transformer改进高效架构

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。