GatedFWA：带门控关联记忆的线性闪存窗口注意力
机器学习
2026-01-08 v2
摘要
现代自回归模型依赖注意力机制，但Softmax注意力在Transformer中的全注意力随序列长度呈二次方比例扩大。滑动窗口注意力 (SWA) 通过约束注意力模式实现线性时间编码/解码，但在关联记忆解释下，其差分风格更新导致训练目标实际上是无界的。相比之下，Softmax注意力对更新进行归一化，导致记忆缩小和梯度消失。我们提出GatedFWA：一种记忆-门控 (Flash) 窗口 注意力机制，保持SWA的效率，同时稳定记忆更新并可控梯度流。本质上，GatedFWA 将每个标记/头的门累积到衰减偏置中，作为记忆 recurrence 中的可学习收缩因子。我们实现了一种单次门处理和兼容FlashAttention的内核，将门在滑动掩码下注入，确保I/O效率和数值稳定性。在语言建模基准测试中，GatedFWA在竞争性吞吐量方面表现良好，开销可忽略，更好地利用全局上下文，并且能够与NSA等标记压缩/选择方法无缝集成，并可推广到各种自回归领域。
引用
@article{arxiv.2512.07782,
title = {GatedFWA: Linear Flash Windowed Attention with Gated Associative Memory},
author = {Jiaxu Liu and Yuhe Bai and Xiangyu Yin and Christos-Savvas Bouganis},
journal= {arXiv preprint arXiv:2512.07782},
year = {2026}
}