通过引入曲率感知预处理,提升序列建模在长上下文中的性能。
Preconditioned DeltaNet: Curvature-aware Sequence Modeling for Linear Recurrences
- 基于在线最小二乘理论,为递归运算引入对角预处理机制。
- 在340M和1B参数规模语言模型上实现稳定性能提升。
- 适用于长序列建模,尤其适合需高精度记忆的场景。
为应对softmax注意力在长上下文计算上的瓶颈,已有若干次二次方递归算子被提出,如Mamba-2、DeltaNet、Gated DeltaNet(GDN)和Kimi Delta Attention(KDA)。随着递归形式增多,相关研究开始对其进行分类。其中一种有说服力的观点是测试时回归(TTR)框架,将递归视为在线最小二乘更新,学习从键到值的线性映射。现有delta规则递归可视为该目标的一阶近似,但忽略了最小二乘损失的曲率信息。本文通过引入预处理机制解决此问题。基于在线最小二乘理论,我们推导出精确预处理下线性注意力与delta规则的等价性。进一步,提出对角近似实现方案,并构建了DeltaNet、GDN和KDA的预处理变体,同时设计高效分块并行算法。实验表明,预处理后的delta规则递归在合成回忆基准及340M与1B规模语言建模任务中均实现一致性能提升。
原文摘要 · Abstract (English)
To address the increasing long-context compute limitations of softmax attention, several subquadratic recurrent operators have been developed. This work includes models such as Mamba-2, DeltaNet, Gated DeltaNet (GDN), and Kimi Delta Attention (KDA). As the space of recurrences grows, a parallel line of work has arisen to taxonomize them. One compelling view is the test-time regression (TTR) framework, which interprets recurrences as performing online least squares updates that learn a linear map from the keys to values. Existing delta-rule recurrences can be seen as first-order approximations to this objective, but notably ignore the curvature of the least-squares loss during optimization. In this work, we address this by introducing preconditioning to these recurrences. Starting from the theory of online least squares, we derive equivalences between linear attention and the delta rule in the exactly preconditioned case. Next, we realize this theory in practice by proposing a diagonal approximation: this enables us to introduce preconditioned variants of DeltaNet, GDN, and KDA alongside efficient chunkwise parallel algorithms for computing them. Empirically, we find that our preconditioned delta-rule recurrences yield consistent performance improvements across synthetic recall benchmarks and language modeling at the 340M and 1B scale.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。