Signed KDA gates yield 2D rotations in one update, keeping 96% of KDA throughput

Complex KDA: Understanding and Enhancing the Expressivity of Kimi Delta Attention

Julien Siems, Riccardo Grazzi, Korbinian Pöppel, Jaisidh Singh, Arber Zela, Timur Carstensen, Jenia Jitsev, Frank Hutter, Volkan Cevher, Antonio Orvieto, Aaron Klein

cs.LG

2026-09-21

CKDA extends KDA gates to [-1,1] and β to [0,2], forming 2D rotations in one update. One layer tracks SO(3) subgroups; at 1.3B it matches KDA (54.06%) at 96–97% throughput.

What problem this solves

Linear RNNs that use the delta rule scale linearly in sequence length and keep a fixed-size state. The typical update is a diagonal gate times a rank-one correction: Kimi Delta Attention (KDA) writes this as (I − βkk⊤)Diag(α). That form is cheap, and it is also expressively tight. The spectrum stays real, and mixing across channels is only rank one.

DeltaProduct₂ already showed that stacking two delta-rule factors in one step can realize a planar rotation, at the cost of extra rank and extra work per token. KDA's channel-wise gate looks like a small upgrade over Gated DeltaNet's scalar gate. This paper's claim is sharper: extend the gate from [0,1] to [-1,1] and β to [0,2], and the gate itself supplies a coordinate reflection. Compose that with a Householder reflection, and a 2D rotation appears inside a single diagonal-plus-rank-one update.

Method

The resulting model is Complex KDA (CKDA). In code, β = 2σ(b) and each gate is r = 2σ(a) − 1. The kernels stay the KDA recurrence from FLA, with a sign-handling patch.

Geometry is straightforward. At β = 2 the Householder is a reflection. A diagonal entry of −1 is another reflection, this time axis-aligned. Two reflections whose mirrors differ by θ compose into a rotation by 2θ. A scalar gate cannot do this: it always commutes with the Householder, so Gated DeltaNet stays symmetric with a real spectrum. Extending a scalar gate into the negatives only flips a global sign.

A few theorems pin the construction down:

Products of CKDA maps cover every square non-expansive matrix with at most 2n−2 factors in dimension n, and every orthogonal matrix with n−1.

Results

On one-layer S₃ and S₄ word problems, the two range extensions together give the strongest length extrapolation among the KDA settings tested. Extending only the gate or only β leaves long S₃ scaled accuracy near 0.2, which is parity of the two A₃ cosets. A successful head drives β toward 2, gates toward ±1, and a complex pair toward the unit circle, matching the theory. A₅ does not learn from random init; a quaternion-near initialization does extrapolate.

On periodic waveform continuation, the model sees a half-bar cue, then zeros, and must keep phase without output feedback. Trained up to length 136, CKDA holds 38.1 dB SNR at length 264; a causal Transformer falls to 2.8 dB. A GRU still wins this task.

Language modeling: at 340M parameters on 15B Nemotron-CC tokens, CKDA's average downstream accuracy is 52.30% against 51.32% for standard KDA. At 1.3B on 100B FineWeb-Edu tokens, recurrent CKDA scores 54.06% average, their bounded-gate KDA 54.09%, Gated DeltaNet-2 53.11%, and an SWA-hybrid Transformer 50.86%. Wiki perplexity is 15.78 for CKDA and 15.73 for KDA. A 3:1 recurrent-to-attention hybrid reaches 54.37% for CKDA and 53.92% for KDA. Kernel throughput is 96–97% of standard KDA.

Why it matters

Planar rotations no longer require two rank-one updates per token. The change is a parameter-range extension plus sign bookkeeping, and the kernel barely slows down. For state tracking, periodic signals, and non-commutative composition, this is a cheaper path than DeltaProduct₂. On language modeling the gap versus KDA is small; the honest reading is "expressivity filled in, quality held."

Code is at OpenEuroLLM/ComplexKDA; weights are public.

Limitations

The rank-one DPLR form allows only one persistent complex pair per step, so multi-plane rotations still need products over time. A₅ is representable and not learned from the standard random start, so expressivity is not learnability. The weighted-automaton construction needs β > 2 and drops the non-expansive guarantee. At 1.3B, CKDA and KDA are essentially tied, and Wiki perplexity slightly favors KDA; the 340M edge depends on spread initialization. GRU remains better on the audio task. Comparisons with published KDA numbers also mix different gate implementations and hybrid ratios.

Terms

Source

What people are saying

Related papers

All paper explainers