用数学推导优化线性注意力的更新步长,提升长文本建模精度和速度。
Kaczmarz Linear Attention
- 基于Kaczmarz投影法推导出归一化步长,自动调节更新幅度。
- 在10亿词预算下,困惑度降至8.09,优于GDN的8.50。
- 适合需要高效长序列建模的场景,如超长文档理解与生成。
长上下文语言建模仍是现代序列建模的核心挑战,但Transformer注意力的二次复杂度使扩展计算成本过高。线性递归模型通过将上下文压缩为固定大小的状态来缓解此瓶颈,状态的遗忘、写入与编辑机制成为设计关键。为解决状态维护问题,门控增量网络(GDN)结合门控状态衰减与增量规则残差写入,使用可学习系数平衡遗忘与更新强度。然而该系数是经验性学习而非由底层目标推导,可能导致次优更新幅度。本文重新审视GDN背后的在线回归目标,受Kaczmarz投影法启发,推导出关键范数归一化的动态步长 $β_t = η_t / ( |k_t |_2^2 + ε)$ 用于残差更新。提出Kaczmarz线性注意力(KLA),仅对GDN进行单标量修改,保持状态形状、门控结构、线性递归及分块并行算法。在0.4B参数量、10亿词预算下,KLA在所有线性时间基线中取得最低验证困惑度(8.09,对比GDN的8.50),且在65K上下文长度下仍稳定。在控制任务中,KLA实现100%单针检索准确率,多查询关联召回率较GDN提升7.03点,32K上下文时解码吞吐量提高2.1倍。结果表明,关键范数归一化的Kaczmarz系数是增量规则序列模型的一阶设计维度:无需改变递归状态或硬件内核即可同时提升精度、外推能力与解码效率。
原文摘要 · Abstract (English)
Long-context language modeling remains central to modern sequence modeling, but the quadratic cost of Transformer attention makes scaling computationally prohibitive. Linear recurrent models address this bottleneck by compressing the context into a fixed-size state, making the rule that forgets, writes, and edits information a central design problem. To address state maintenance, Gated DeltaNet (GDN) combines gated state decay with delta-rule residual writes, using a learnable coefficient to balance forgetting and update magnitude. However, this coefficient is learned empirically rather than derived from the underlying objective, which can lead to suboptimal update magnitudes. We revisit the online-regression objective underlying GDN and, inspired by the Kaczmarz projection method, derive the key-norm-normalized dynamic step size $β_t = η_t / (\|k_t\|_2^2 + ε)$ for residual updates. We propose Kaczmarz Linear Attention (KLA), a one-scalar modification of GDN that preserves the state shape, gates, linear recurrence, and chunkwise parallel algorithm. At the 0.4B scale with a 1B-token budget, KLA achieves the lowest validation perplexity among evaluated linear-time baselines, 8.09 versus 8.50 for GDN, and remains stable up to 65K tokens. On controlled tasks, KLA reaches 100% on single-needle-in-a-haystack retrieval, improves 8x multi-query associative recall by 7.03 points over GDN, and delivers 2.1x higher decode throughput at 32K context. These results suggest that the key-norm-normalized Kaczmarz coefficient is a first-order design axis for delta-rule sequence models: it improves accuracy, extrapolation, and decoding efficiency without changing the recurrent state or hardware kernel.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。