Preconditioned DeltaNet: Curvature-aware Sequence Modeling for Linear Recurrences

TL;DR

This paper introduces preconditioning to delta-rule recurrences, enhancing their performance in sequence modeling tasks by incorporating curvature information from online least squares theory.

Contribution

It develops a theoretical framework connecting linear attention and delta rule with preconditioning, and implements practical preconditioned variants of DeltaNet, GDN, and KDA.

Findings

Preconditioned delta-rule recurrences improve performance on synthetic recall benchmarks.

They also enhance language modeling at 340M and 1B scale.

Abstract

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.