DART：用于高效长上下文序列建模的循环状态解码注意力

中文摘要 AI 辅助

现代语言模型主要由Transformer、循环模型及其混合架构构成。Transformer依赖token级注意力记忆，而状态空间模型（SSM）和线性注意力等循环模型则维护紧凑的循环状态。这些架构通常单独实例化或在层间交错，仍存在一个待解决的问题：共享记忆表示能否同时支持循环压缩和注意力式检索。我们通过Mamba-2的状态空间对偶性（SSD）视角研究该问题，其中SSM状态可被解释为压缩的关联键值（KV）缓存。我们观察到，Mamba-2从该状态解码token条件化的值，但未解码token条件化的键。基于此观察，我们提出DART（Decoded Attention over Recurrent sTates，循环状态解码注意力），其保留Mamba-2分块扫描生成的分块状态贡献作为分块状态记忆，从这些记忆中解码token条件化的键和值，并对生成的KV对执行状态-记忆注意力（SMA）。检索到的输出随后通过门控残差连接与原生Mamba-2输出结合。DART通过复用Mamba-2分块扫描并将SMA实现为FlashAttention式计算，支持实用训练。我们的分析与实验表明，与匹配的注意力基线相比，DART大幅减少了依赖长度的推理缓存（例如，当分块大小S=256、状态大小N=128时，缓存节省75%）。与Mamba-2相比，DART在保持通用语言建模质量的同时，大幅提升了关联召回率和检索性能。

英文摘要

Modern language models are built primarily from Transformers, recurrent models, and their hybrid architectures. Transformers rely on token-level attention memories, while recurrent models such as state space models (SSMs) and linear attention maintain compact recurrent states. These architectures are typically instantiated separately or interleaved at the layer level, leaving open whether a shared memory representation can support both recurrent compression and attention-style retrieval. We study this question through the state space duality (SSD) view of Mamba-2, where the SSM state can be interpreted as a compressed associative key--value (KV) cache. We observe that Mamba-2 decodes token-conditioned values from this state but does not decode token-conditioned keys. Based on this observation, we propose DART (Decoded Attention over Recurrent sTates), which retains the chunk state contributions produced by the Mamba-2 chunked scan as chunk state memories, decodes token-conditioned keys and values from these memories, and performs state-memory attention (SMA) over the resulting KV pairs. The retrieved output is then combined with the native Mamba-2 output through a gated residual connection. DART supports practical training by reusing the Mamba-2 chunked scan and implementing SMA as a FlashAttention-style computation. Our analysis and experiments show that DART substantially reduces the length-dependent inference cache compared with a matched attention baseline (e.g., $75\%$ savings when the chunk size is $S=256$ and the state size is $N=128$). Compared with Mamba-2, DART substantially improves associative recall and retrieval while preserving general language-modeling quality.

发表机构

• College of Control Science and Engineering, Zhejiang University（浙江大学控制科学与工程学院）
• National University of Singapore（新加坡国立大学）
• Hong Kong Polytechnic University（香港理工大学）
• School of Science, Huzhou Normal University（湖州师范学院理学院）

机构由 AI 辅助整理，请以论文原文为准。