🤖 machine learning

Preconditioned DeltaNet: Curvature-aware Sequence Modeling for Linear Recurrences

This paper introduces Preconditioned DeltaNet, a curvature-aware framework that enhances subquadratic recurrent models like DeltaNet and Mamba-2 by deriving preconditioned variants from online least squares theory, thereby achieving consistent performance improvements in long-context language modeling and synthetic recall tasks.

Original authors: Neehal Tumma, Noel Loo, Daniela Rus

Published 2026-04-24

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Neehal Tumma, Noel Loo, Daniela Rus

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 teach a robot to read a very long book and remember specific details from the beginning to answer questions at the end. This is what Large Language Models (LLMs) do.

For a long time, the best way to do this was like a photocopier. Every time the robot read a new word, it had to look back at every single word it had read before to see how they connected. This is called "Attention." It's incredibly smart, but it's slow and expensive. If the book gets twice as long, the robot has to do four times as much work.

To fix this, scientists invented Linear Recurrences (like DeltaNet). Think of these as a notebook. Instead of looking back at the whole book, the robot just updates its current note. It's super fast, but it's a bit "dumb." It treats every new piece of information the same way, often forgetting important details or getting confused when the book gets too long.

The Problem: The "Flat" Notebook

The paper argues that these "notebook" models are like trying to walk through a hilly landscape using a flat map. They take steps based on the slope right under their feet (the immediate error), but they don't understand the shape of the hill (the curvature).

The Analogy: Imagine you are trying to find the bottom of a valley.
- Standard models (DeltaNet) are like a blindfolded hiker who just takes a step downhill based on the slope right under their foot. If the valley is steep on one side and flat on the other, they might zigzag wildly or get stuck.
- The Goal: They want to know the shape of the valley so they can take a giant, direct step straight to the bottom.

The Solution: Preconditioned DeltaNet (The "Curvature-Aware" Hiker)

The authors introduce Preconditioning. In math, this is like giving the hiker a pair of special glasses that reshape the landscape so the hills look flat and the valleys look like straight paths.

The "Preconditioner" (The Glasses): The model learns to adjust its "notebook" based on how often it has seen certain patterns before. If a specific type of information is very common, the model knows to be careful with it. If it's rare, it knows to pay extra attention.
The "Write Key" vs. "Read Key": In the old models, the robot used the same finger to read a note and write a new one. The new model uses two different fingers.
- Read Key: "What am I looking at?"
- Write Key: "How should I update my memory based on what I'm looking at?"
- By separating these, the model can decide how to update its memory without messing up what it just read.

The Magic Trick: The "Diagonal" Shortcut

Calculating the perfect "glasses" (the exact shape of the valley) is computationally expensive, like trying to map every single rock in the valley.

The authors found a clever shortcut. Instead of mapping the whole 3D shape of the valley, they just measured the height at each point (a diagonal approximation).

The Analogy: Instead of drawing a full topographic map, they just wrote down "High here, Low there" on a simple list.
Why it works: It turns out that for language, knowing the "height" (importance) of each dimension is 90% of the battle. This keeps the model fast (like the old notebook) but smart (like the photocopier).

The Results: Faster and Smarter

The paper tested this new "Preconditioned DeltaNet" (and its cousins like PGDN and PKDA) on:

Synthetic Games: Like finding a needle in a haystack (remembering a specific fact from a long list).
Real Language: Writing essays and answering questions.

The Verdict: The new models were consistently better at remembering long contexts and understanding language, without slowing down the computer. They managed to get the "best of both worlds": the speed of a notebook and the memory of a photocopier.

Summary in One Sentence

The authors gave fast, simple language models a pair of "curvature-aware glasses" and a two-fingered writing system, allowing them to remember long stories much better without needing a supercomputer to do it.

1. Problem Statement

Large Language Models (LLMs) face significant computational bottlenecks when scaling to long context lengths due to the quadratic complexity of standard softmax attention. To address this, researchers have developed sub-quadratic recurrent operators, such as DeltaNet, Gated DeltaNet (GDN), and Kimi Delta Attention (KDA). These models are often interpreted through the Test-Time Regression (TTR) framework, where the recurrent state learns a linear map from keys to values by minimizing a least-squares objective.

However, existing delta-rule recurrences (e.g., DeltaNet) are essentially first-order approximations (similar to Stochastic Gradient Descent) of this least-squares problem. They ignore the curvature of the loss landscape (governed by the key Gram matrix), leading to suboptimal convergence and potentially poor long-context retrieval capabilities. While exact second-order methods exist (e.g., MesaNet), they suffer from either a lack of parallelizability (sequential updates) or numerical instability (conjugate gradient approximations).

The Core Challenge: How to incorporate second-order curvature information (preconditioning) into linear recurrences to improve performance and stability while maintaining the efficiency of chunkwise parallel training.

2. Methodology

The authors propose Preconditioned DeltaNet (PDN) and its variants, which introduce a preconditioner to the recurrence updates. This approach is grounded in the theory of online least squares.

A. Theoretical Foundation: Equivalence under Exact Preconditioning

The paper first establishes a theoretical equivalence between Linear Attention and DeltaNet when an exact preconditioner (the inverse of the key Gram matrix, $G_t^{-1}$ ) is applied.

Linear Attention (Offline Solve): Interpreted as a single full-batch gradient descent step. By applying the inverse Gram to the query ( $\tilde{q}_t = G_t^{-1} q_t$ ), it recovers the exact least-squares solution.
DeltaNet (Online Solve): Interpreted as an online SGD step. By applying the inverse Gram to the write key ( $\tilde{k}_t = G_t^{-1} k_t$ ), the update rule becomes a preconditioned SGD step that also recovers the exact least-squares solution.
Key Insight: In the exact case, Linear Attention with query-preconditioning (ATQ) and DeltaNet with key-preconditioning (ATK) are mathematically equivalent.

B. Practical Realization: Diagonal Approximation

Computing the exact inverse Gram matrix is computationally expensive and sequential, preventing efficient chunkwise parallel training. To solve this, the authors propose a diagonal approximation of the key Gram matrix:
$A_t = A_{t-1} + k_t \odot k_t$
where $A_t$ stores the coordinate-wise second moments of the keys. The preconditioner is then $P_t = \text{Diag}(A_t)^{-1}$ .

Efficiency: The diagonal update is an additive prefix sum, which admits an efficient chunkwise parallel form.
Stability: To prevent numerical instability (heavy tails in the distribution of $A_t$ ), the authors introduce a learnable, non-linear transformation. They map the log-scaled $A_t$ through a "squash" function to bound the preconditioner values within a stable interval $[1/x, x]$ (where $x \le 2$ ). This avoids explicit matrix inversion and ensures recurrent eigenvalues remain within the unit disk.

C. The DGPS Taxonomy

The paper introduces a new design taxonomy for linear recurrences called DGPS, which parameterizes the space along four orthogonal axes:

Decay ( $\alpha_t$ ): Controls state forgetting (e.g., Mamba-2, GDN).
Gain ( $\beta_t$ ): Controls the magnitude of the update.
Preconditioner ( $P_t$ ): The curvature-aware transform (ATQ for queries, ATK for keys).
Solve: Offline (Linear Attention) vs. Online (DeltaNet).

This framework unifies existing models and allows for the systematic construction of new, efficient variants like Preconditioned GDN (PGDN) and Preconditioned KDA (PKDA).

3. Key Contributions

Theoretical Unification: Proved that Linear Attention and DeltaNet are equivalent under exact preconditioning, revealing them as two realizations of the same operator.
Efficient Algorithms: Derived chunkwise parallel forms for preconditioned recurrences using a diagonal approximation. This enables training on modern GPUs with high throughput, unlike sequential exact solvers.
Stable Parameterization: Proposed a novel, learnable parameterization for the diagonal preconditioner that ensures numerical stability and avoids the training instabilities observed in previous exact methods (like MesaNet).
New Taxonomy (DGPS): Defined a modular design space for linear recurrences that decouples decay, gain, preconditioning, and solve strategies, facilitating future model design.
Implementation: Open-sourced custom Triton kernels for these models, demonstrating that the overhead of the preconditioner is minimal (~10%) compared to the base models.

4. Experimental Results

The authors evaluated their models (PDN, PGDN, PKDA) at 340M and 1B parameter scales on the SlimPajama dataset.

Synthetic Benchmarks (MQAR): Preconditioned models consistently matched or outperformed their unpreconditioned counterparts in associative recall tasks, demonstrating improved ability to track key-value pairs over long sequences.
Long-Context Retrieval (S-NIAH): On the "Needle In A Haystack" benchmark, preconditioned models showed significant improvements, particularly in the GDN and KDA variants. This suggests the preconditioner helps the model better distinguish and retrieve specific information from long contexts.
Language Modeling:
- Commonsense Reasoning: PGDN achieved the best performance among 340M models on commonsense benchmarks.
- Perplexity: Preconditioned variants generally achieved lower perplexity (better language modeling) than their baselines.
- Comparison: While adding negative eigenvalues to baseline models improved them slightly, the preconditioned models still outperformed them, proving that the gains come from curvature awareness rather than just spectral properties.
Efficiency: The preconditioned models incurred only a ~10% overhead in tokens-per-second (TPS) compared to their non-preconditioned baselines, while being >20% faster than MesaNet.

5. Significance

This work bridges the gap between theoretical optimality (exact least-squares regression) and practical efficiency in sequence modeling. By introducing curvature-aware updates via a diagonal preconditioner, the authors demonstrate that linear recurrences can be made significantly more expressive and effective for long-context tasks without sacrificing the sub-quadratic training speed that makes them attractive.

The DGPS taxonomy provides a structured way to innovate in the space of linear recurrences, suggesting that "preconditioning" is a powerful, underutilized lever for improving state-tracking and memory in efficient LLMs. The results indicate that future high-performance linear models should likely incorporate second-order information (preconditioning) to compete with the recall capabilities of full attention mechanisms.