论文
DART:解码循环状态的注意力以实现高效的长上下文序列建模
DART: Decoded Attention over Recurrent States for Efficient Long-Context Sequence Modeling
摘要
现代语言模型主要由 Transformer、循环模型及其混合架构构建。 Transformer 依赖于 词元级 注意力记忆,而状态空间模型(SSM)和线性注意力等循环模型保持紧凑的循环状态。这些架构通常在层级别单独实例化或交错,从而使共享内存表示是否可以支持循环压缩和注意力式检索保持开放。我们通过 Mamba-2 的状态空间对偶 (SSD) 视图来研究这个问题,其中 SSM 状态可以解释为压缩的关联键值 (KV) 缓存。我们观察到 Mamba-2 从该状态解码词元条件值,但不解码词元条件密钥。基于这一观察,我们提出了 DART(Decoded Attention over Recurrent sTates),它将 Mamba-2 分块扫描产生的块状态贡献保留为块状态记忆,从这些记忆中解码词元条件键和值,并对生成的 KV 对执行状态记忆注意力(SMA)。然后,检索到的输出通过门控残差连接与本机 Mamba-2 输出组合。 DART 通过重用 Mamba-2 分块扫描并将 SMA 实现为 FlashAttention 式计算来支持实际训练。我们的分析和实验表明,与匹配的注意力基线相比,DART 大大减少了长度相关的推理缓存(例如,当块大小为 $S=256$ 且状态大小为 $N=128$ 时,节省 $75\%$)。与 Mamba-2 相比,DART 显着提高了联想回忆和检索,同时保留了一般语言建模质量。