让长序列模型用更少内存实现高效注意力,提升记忆检索能力。
DART: Decoded Attention over Recurrent States for Efficient Long-Context Sequence Modeling
- 从Mamba-2的递归状态中解码出键值对,实现状态内存注意力
- 在块大小256、状态128时,推理缓存减少75%
- 适合需要长上下文记忆和低资源推理的场景
现代语言模型主要基于Transformer、循环模型及其混合架构。Transformer依赖于逐标记的注意力记忆,而像状态空间模型(SSMs)和线性注意力这样的循环模型则维护紧凑的递归状态。这些架构通常单独或分层交错使用,尚未明确是否存在共享的记忆表示可同时支持递归压缩与注意力式检索。我们通过Mamba-2的状态空间对偶(SSD)视角研究该问题,发现其状态可被解释为压缩的键值(KV)缓存。观察到Mamba-2能从状态中解码出依赖标记的值,但无法解码依赖标记的键。基于此,我们提出DART(Decoded Attention over Recurrent sTates),保留Mamba-2分块扫描生成的块状态贡献作为状态记忆,从中解码出依赖标记的键和值,并对所得键值对执行状态内存注意力(SMA)。最终输出通过门控残差连接与原Mamba-2输出融合。DART通过重用Mamba-2分块扫描并以FlashAttention风格实现SMA,支持实际训练。分析与实验表明,相比匹配的注意力基线,DART在块大小S=256、状态大小N=128时,显著降低长度相关推理缓存(减少75%);相较于Mamba-2,DART大幅改善关联召回与检索性能,同时保持语言建模质量。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。