OSDN: Improving Delta Rule with Provable Online Preconditioning in Linear Attention

📄 arXiv: 2605.13473v1 📥 PDF

作者: Chenyu Zhou, Hongpei Li, Yuerou Liu, Jianghao Lin, Dongdong Ge, Yinyu Ye

分类: cs.LG, cs.CL

发布日期: 2026-05-13

💡 一句话要点

OSDN：通过可证明的在线预处理改进线性注意力中的Delta规则

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 线性注意力 Delta规则 在线预处理 上下文学习 超梯度反馈

📋 核心要点

  1. 线性注意力模型在上下文学习中存在关联回忆能力不足的问题，限制了其应用。
  2. OSDN通过在线更新的对角预处理器来调整Delta规则的步长，实现特征级别的自适应缩放。
  3. 实验结果表明，OSDN在不同参数规模下均能显著提升上下文回忆能力，并在下游任务中保持竞争力。

📝 摘要（中文）

线性注意力和状态空间模型为softmax注意力提供了恒定内存的替代方案，但通常在上下文关联回忆方面表现不佳。Delta规则通过在线梯度下降一步写入每个token来缓解这个问题。然而，它的步长依赖于一个标量门，忽略了内部目标的特征曲率。我们提出了在线缩放DeltaNet（OSDN），它通过超梯度反馈在线更新的对角预处理器来增强标量门。关键是，这种右预处理在代数上等价于write-side key的每个特征缩放。这种等价性允许OSDN严格保留DeltaNet的硬件友好型分块并行流水线，而不会产生高维状态开销。理论上，通过利用内部回归损失的精确二次结构，我们建立了针对右牛顿比较器的超几何收敛，并证明了算法对齐的token局部残差收缩界限。为了处理非平稳上下文，我们进一步引入了自适应预处理器遗忘（APF）来动态刷新过时的校准。实验表明，OSDN在不同规模上表现出强大的性能。在3.4亿参数规模下，OSDN比DeltaNet在JRT风格的上下文回忆中提高了32%。扩展到13亿参数，它实现了39%的回忆残差率降低，同时保持了一般下游任务（例如，困惑度和LongBench）的均等性——表明我们的在线预处理机制有效地转移并在十亿参数规模上放大。

🔬 方法详解

问题定义：线性注意力模型，如DeltaNet，在处理上下文关联回忆任务时，由于步长调整依赖于单一标量门，无法有效捕捉特征级别的曲率信息，导致性能受限。现有方法难以在不引入额外计算开销的前提下，提升模型的上下文学习能力。

核心思路：OSDN的核心思路是通过在线预处理（Online Preconditioning）来改进Delta规则。具体来说，它引入一个对角预处理器，该预处理器通过超梯度反馈进行在线更新，从而实现对write-side key的特征级别缩放。这种方法旨在更精确地调整每个特征的步长，从而提高模型的学习效率和泛化能力。

技术框架：OSDN的整体框架与DeltaNet类似，仍然采用chunkwise并行流水线，以保证硬件友好性。主要区别在于，OSDN在Delta规则中引入了一个在线更新的对角预处理器。该预处理器根据内部回归损失的梯度信息进行更新，并用于缩放write-side key。为了处理非平稳上下文，OSDN还引入了自适应预处理器遗忘（APF）机制，用于动态刷新过时的校准。

关键创新：OSDN的关键创新在于将在线预处理技术应用于线性注意力模型中的Delta规则。通过引入对角预处理器，OSDN能够自适应地调整每个特征的步长，从而更有效地学习上下文信息。此外，OSDN还证明了这种右预处理在代数上等价于write-side key的特征级别缩放，从而保证了模型的高效性和可扩展性。

关键设计：OSDN的关键设计包括：1) 对角预处理器的在线更新机制，该机制基于超梯度反馈，能够自适应地调整每个特征的步长；2) 自适应预处理器遗忘（APF）机制，用于处理非平稳上下文，动态刷新过时的校准；3) 理论分析，证明了OSDN的超几何收敛性和token局部残差收缩界限。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

OSDN在3.4亿参数规模下，相比DeltaNet在JRT风格的上下文回忆任务中提升了32%。在扩展到13亿参数时，OSDN实现了39%的回忆残差率降低，同时在困惑度和LongBench等下游任务中保持了与DeltaNet相当的性能。这些结果表明，OSDN的在线预处理机制能够有效地提升模型的上下文学习能力，并且可以扩展到更大的参数规模。

🎯 应用场景

OSDN的潜在应用领域包括自然语言处理、语音识别、计算机视觉等。它可以用于提升各种序列模型的上下文学习能力，例如，在长文本建模、对话生成、图像描述等任务中，OSDN可以帮助模型更好地理解上下文信息，从而生成更准确、更连贯的结果。此外，OSDN的硬件友好型设计使其易于部署在各种平台上，具有广泛的应用前景。

📄 摘要（原文）

Linear attention and state-space models offer constant-memory alternatives to softmax attention, but often struggle with in-context associative recall. The Delta Rule mitigates this by writing each token via one step of online gradient descent. However, its step size relies on a single scalar gate that ignores the feature-wise curvature of the inner objective. We propose Online Scaled DeltaNet (OSDN), which augments the scalar gate with a diagonal preconditioner updated online via hypergradient feedback. Crucially, this right-preconditioning is algebraically equivalent to a per-feature scaling of the write-side key. This equivalence allows OSDN to strictly preserve the hardware-friendly chunkwise parallel pipeline of DeltaNet without incurring high-dimensional state overhead. Theoretically, by exploiting the exact-quadratic structure of the inner regression loss, we establish super-geometric convergence against a right-Newton comparator and prove an algorithm-aligned token-local residual contraction bound. To handle non-stationary contexts, we further introduce Adaptive Preconditioner Forgetting (APF) to dynamically refresh stale calibration. Empirically, OSDN demonstrates strong performance across scales. At the 340M-parameter scale, OSDN improves JRT-style in-context recall by 32% over DeltaNet. Scaling to 1.3B parameters, it achieves a 39% reduction in the recall residual ratio while maintaining parity on general downstream tasks (e.g., perplexity and LongBench) -- demonstrating that our online-preconditioning mechanism effectively transfers and amplifies at the billion-parameter scale.