提出GatedFWA，让窗口注意力更稳定高效。
GatedFWA: Linear Flash Windowed Attention with Gated Associative Memory
- 用可学习门控机制控制记忆更新，避免梯度消失
- 在语言建模中实现接近FlashAttention的吞吐量
- 适合需要长序列建模与稳定训练的自回归任务
现代自回归模型依赖注意力机制，但Transformer中的Softmax全注意力随序列长度呈平方增长。滑动窗口注意力（SWA）通过限制注意力模式实现线性时间编码/解码，但在关联记忆视角下，其差分式更新导致训练目标实质上无界。相反，Softmax注意力会归一化更新，引发记忆萎缩和梯度消失。本文提出GatedFWA：一种带记忆门控的闪光窗口注意力机制，在保持SWA效率的同时稳定记忆更新并可控梯度流动。本质上，GatedFWA将每个标记/头的门控值累积为衰减偏置加入注意力logits，作为记忆递推中的可学习收缩项。我们实现了融合的一次性门控预处理与兼容FlashAttention的核函数，在滑动掩码下注入门控，确保输入输出效率与数值稳定性。在语言建模基准上，GatedFWA以可忽略开销达到竞争性吞吐量，并更好利用全局上下文，且可无缝集成令牌压缩/选择方法如NSA，泛化至多种自回归领域。
原文摘要 · Abstract (English)
Modern autoregressive models rely on attention, yet the Softmax full attention in Transformers scales quadratically with sequence length. Sliding Window Attention (SWA) achieves linear-time encoding/decoding by constraining the attention pattern, but under an \textit{Associative Memory} interpretation, its difference-style update renders the training objective effectively \emph{unbounded}. In contrast, Softmax attention normalizes updates, leading to \emph{memory shrinkage and gradient vanishing}. We propose GatedFWA: a Memory-\underline{Gated} (\underline{F}lash) \underline{W}indowed \underline{A}ttention mechanism that preserves SWAs efficiency while stabilizing memory updates and making gradient flow controllable. In essence, GatedFWA accumulate a per-token/head gate into a decay bias added to the attention logits, acting as a learnable contraction in the memory recurrence. We implement a fused one-pass gate preprocessing and a FlashAttention-compatible kernel that injects the gate under a sliding mask, ensuring I/O efficiency and numerical stability. On language modelling benchmarks, GatedFWA delivers competitive throughput with negligible overhead and better use of global context, and it integrates cleanly with token compression/selection methods such as NSA and generalizes to various autoregressive domains.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。