💬 NLP

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

The paper introduces Online Scaled DeltaNet (OSDN), a linear attention model that enhances the Delta Rule by incorporating an online-updated diagonal preconditioner for feature-wise scaling, thereby achieving provable super-geometric convergence and significantly improving in-context associative recall while maintaining hardware-efficient parallelism.

Original authors: Chenyu Zhou, Hongpei Li, Yuerou Liu, Jianghao Lin, Dongdong Ge, Yinyu Ye

Published 2026-05-14

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Chenyu Zhou, Hongpei Li, Yuerou Liu, Jianghao Lin, Dongdong Ge, Yinyu Ye

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to memorize a very long story so you can answer questions about it later. In the world of AI, this is called "in-context learning." The AI reads the story (the context) and updates its internal memory to remember important details.

For a long time, the best way to do this was like a giant library with a perfect index (called "Softmax Attention"). But this library is slow and takes up a massive amount of space. Newer, faster methods (like DeltaNet) act more like a person taking notes in a small notebook. They are fast and efficient, but they sometimes struggle to remember specific details if the story gets too long or if the same words appear many times. They tend to "blur" the notes together.

This paper introduces a new method called OSDN (Online Scaled DeltaNet) to fix that blurring. Here is how it works, using simple analogies:

1. The Problem: The "One-Size-Fits-All" Pen

Imagine the AI is writing in its notebook. Every time it sees a new word (a "token"), it writes a note.

Old Method (DeltaNet): The AI uses a single pen with a fixed ink flow. If it needs to write a tiny, delicate detail, the pen is too thick and smudges it. If it needs to write a huge, bold headline, the pen is too thin and the ink runs out. It treats every piece of information with the exact same "step size" or intensity.
The Issue: Some facts in the story are rare and need a strong, clear note. Others are common and need a light touch. Using the same "pen pressure" for everything leads to mistakes in recall.

2. The Solution: The "Smart, Adjustable Pen" (OSDN)

OSDN gives the AI a smart, adjustable pen. Instead of one fixed setting, the pen has a dial for every single letter of the alphabet (or every feature of the data).

How it learns: As the AI reads the story, it constantly checks: "Did I write that note correctly?" If the note is too faint, it turns the dial up for that specific letter next time. If it's too dark, it turns it down.
The Magic Trick: The paper proves that this "dial" doesn't need a complicated, slow computer to calculate. It can be figured out instantly, just by looking at the word being written. This allows the AI to keep its speed while getting much smarter about how it writes.

3. The "Forgetting" Mechanism (APF)

Sometimes, the story changes topics. A fact that was important at the beginning might be irrelevant by the end.

The Problem: If the AI keeps adjusting its pen based on old topics, it might get confused when the story switches to a new subject.
The Fix (APF): OSDN includes a feature called Adaptive Preconditioner Forgetting. Think of this as a "fresh page" button. If the AI notices the topic has shifted, it gently resets its pen dials for the new topic, so it doesn't get stuck using settings from the previous chapter.

4. Why This Matters (The Results)

The authors tested this on AI models of different sizes (from small to very large).

The "Memory Test": They gave the AI a story and then asked it to recall specific details, especially when those details appeared twice (like a character being introduced, then mentioned again later).
The Outcome:
- At a medium size, OSDN improved the ability to remember repeated details by 32% compared to the old method.
- At a massive size (1.3 billion parameters), it improved memory recall by 39% while keeping the AI just as good at general tasks (like writing sentences or answering general questions).
- Crucially, it did this without slowing down the AI significantly. It's like upgrading a car's engine to be more precise without making the car heavier or slower.

Summary

Think of OSDN as teaching an AI to be a better note-taker. Instead of scribbling everything with the same pressure, it learns to press harder on important, rare details and lighter on common ones. It also knows when to wipe the slate clean for a new topic. This allows the AI to remember complex stories much better without losing its speed or efficiency.

What the paper does NOT claim:

It does not claim this makes AI "conscious" or able to feel emotions.
It does not claim this solves all AI problems (like reasoning or math); it specifically targets the ability to remember and retrieve information from a long text.
It does not suggest this is ready for medical diagnosis or critical real-world deployment yet; it is a research breakthrough in how AI models process sequences of text.

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

1. Problem Statement

Linear attention and state-space models (SSMs) offer $O(N)$ inference complexity and constant-memory alternatives to softmax attention, making them suitable for long-sequence processing. However, they often struggle with in-context associative recall, particularly when compared to softmax-based transformers.

The Delta Rule (as implemented in DeltaNet) attempts to bridge this gap by interpreting the recurrent state update as a single step of online gradient descent (OGD) on a per-token regression loss: $f_t(S) = \frac{1}{2}\|S k_t - v_t\|_F^2$ . The update rule is $S_t = S_{t-1} + \beta_t (v_t - S_{t-1}k_t)k_t^\top$ , where $\beta_t$ is a scalar gate.

The Core Limitation: The scalar gate $\beta_t$ acts as a uniform learning rate applied to all key dimensions. This ignores the feature-wise curvature of the inner objective. In associative recall tasks, different keys (e.g., frequent vs. rare identifiers, stable vs. high-entropy values) exhibit different empirical curvatures. A single scalar step size forces a compromise: increasing it to help overwrite stale memory for one direction may over-correct another direction that is already well-calibrated. Standard adaptive optimizers (like Adam) use diagonal preconditioning to address this, but applying such mechanisms to recurrent layers typically requires dense second-order states or breaks the hardware-friendly chunkwise parallelization pipelines (like the WY transform) that make DeltaNet efficient.

2. Methodology: Online Scaled DeltaNet (OSDN)

The authors propose OSDN, which augments the scalar gate with a diagonal preconditioner $D_t = \text{diag}(d_t)$ , learned online via hypergradient feedback.

2.1. Algebraic Equivalence and Hardware Preservation

A key insight of OSDN is that right-preconditioning the gradient step is algebraically equivalent to a per-feature scaling of the write-side key.
The standard update:
$S_t = S_{t-1} - \beta_t \nabla f_t(S_{t-1}) D_t = S_{t-1} + \beta_t (v_t - S_{t-1}k_t)(d_t \odot k_t)^\top$
By defining a preconditioned key $\tilde{k}_t = d_t \odot k_t$ , the update becomes:
$S_t = S_{t-1} + \beta_t (v_t - S_{t-1}k_t)\tilde{k}_t^\top$
This transformation allows OSDN to strictly preserve the chunkwise WY parallel pipeline of DeltaNet. The implementation only requires substituting the storage key $K$ with $\tilde{K}$ in the chunk Gram matrix and intra-chunk scores, incurring no high-dimensional state overhead beyond the $O(K)$ preconditioner vector.

2.2. Decoupled Hypergradient Feedback

The preconditioner $d_t$ is updated via Online Gradient Descent (OGD) on a hypergradient surrogate derived from the Online Scaled Gradient Method (OSGM) framework [18].
The surrogate loss for the preconditioner is:
$h_t(d_t) = \frac{f_t(S_t) - f_t(S_{t-1})}{\|\nabla f_t(S_{t-1})\|_F^2}$
Crucially, due to the exact-quadratic structure of the inner regression loss $f_t$ , this surrogate admits a closed-form expression that depends only on the key $k_t$ and scalar gate $\beta_t$ , decoupling it from the high-dimensional state $S_t$ , value $v_t$ , or residual $u_t$ .
$h_t(d_t) = \frac{(1 - \beta_t \langle d_t, k_t^2 \rangle)^2 - 1}{2 \|k_t\|_2^2}$
This decoupling enables a two-phase implementation:

Phase 1: A lightweight $O(K)$ scan over the key stream to compute the sequence of preconditioned keys $\{\tilde{k}_t\}$ and update $d_t$ .
Phase 2: The standard chunkwise DeltaNet pass using $\tilde{K}$ instead of $K$ .

2.3. Adaptive Preconditioner Forgetting (APF)

To handle non-stationary language contexts where key distributions shift, the authors introduce APF. This adds a learned, token-wise, head-wise retention gate $r_{t,h}$ applied only to the preconditioner state $d_t$ (not the memory state $S_t$ ):
$\bar{d}_{t+1} = r_{t,h} d_t + \eta \nabla_d h_t(d_t)$
This maintains the affine recurrence structure required for the two-phase scan while allowing the meta-optimizer to "forget" stale calibration.

3. Theoretical Guarantees

The paper establishes two convergence guarantees based on the exact-quadratic nature of the inner loss:

Population-Limit Super-Geometric Convergence (Theorem 5.1): Against an ideal right-Newton comparator ( $D^\star = \Sigma_k^\dagger$ ), if the online learner achieves sublinear regret, the suboptimality of the global objective contracts at a super-geometric rate of $O((C/T)^T)$ . This is significantly faster than standard geometric (linear) convergence.
Algorithm-Local Residual Contraction (Theorem 5.2): For the implemented diagonal update, the product of token-local residual ratios $\prod_{t=1}^T \frac{f_t(S_t)}{f_t(S_{t-1})}$ is bounded. If a diagonal comparator satisfies the gated Newton condition ( $\beta_t \langle d^\star, k_t^2 \rangle = 1$ ), the contraction rate is also super-geometric, specifically $O((1/\sqrt{T})^T)$ under standard OGD regret bounds.

4. Experimental Results

Experiments were conducted on DeltaNet, Gated DeltaNet (GDN), and Kimi Delta Attention (KDA) backbones at 340M and 1.3B parameter scales.

4.1. In-Context Recall (340M Scale)

Performance: Vanilla OSDN improved JRT-style in-context recall by 32% over the baseline DeltaNet (0.150 $\to$ 0.198). The gain was concentrated on repeated-context tasks (e.g., "JRT-twice" variants), where keys recur, allowing the preconditioner to calibrate.
APF Impact: OSDN-APF retained a 17% gain (0.176) and provided the best WikiText/LAMBADA perplexity among DeltaNet variants (30.99 vs. 32.00).
Generalization: On broader benchmarks (Commonsense, LongBench), OSDN variants remained at parity with their baselines, confirming the mechanism is targeted rather than a universal performance booster.

4.2. Scaling to 1.3B Parameters

Residual Contraction: The mechanism-level signal amplified at scale. The geometric mean residual ratio ( $q_{geo}$ ) dropped from 0.432 to 0.265 (a 39% reduction), nearly doubling the improvement seen at 340M.
Downstream Parity: Despite the massive improvement in the internal residual contraction, downstream metrics (WikiText/LAMBADA perplexity, Commonsense, LongBench) remained at parity with the 1.3B DeltaNet baseline, demonstrating that the mechanism transfers effectively without degrading general capabilities.

4.3. Efficiency

Throughput: OSDN variants incurred minimal inference overhead. At 340M, throughput changes were within $\pm 2.2\%$ . At 1.3B, OSDN-APF was 6.8% slower than the baseline, attributed to the additional sequential scan for the preconditioner.
State Overhead: The persistent recurrent state increased by only $\le 0.05\%$ due to the $O(K)$ preconditioner vector.

5. Significance and Claims

The paper positions OSDN as a minimal, targeted extension to the Delta Rule that bridges the gap between first-order scalar updates and adaptive optimization.

Mechanism-Level Improvement: The primary contribution is a provable, mechanism-level improvement in associative recall via online preconditioning, validated by a direct diagnostic of the residual contraction ratio.
Hardware Compatibility: Unlike other second-order or dense preconditioning methods, OSDN preserves the chunkwise parallel pipeline and $O(K)$ state complexity, making it practical for large-scale deployment.
Scalability: The results demonstrate that the benefits of online preconditioning transfer and amplify at the billion-parameter scale, suggesting that the "curvature mismatch" in linear attention is a fundamental bottleneck that can be addressed without architectural overhaul.
Limitations: The authors explicitly note that OSDN is not a universal benchmark improver; gains are concentrated on retrieval tasks involving repeated keys. The theoretical guarantees rely on the exact-quadratic assumption of the inner loss and monotone descent, which are idealized conditions. The method does not claim to solve all in-context learning limitations, but specifically addresses the write-step geometry of the Delta Rule.