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

AI summary

Overview

Research area: Efficient sequence modeling with linear recurrent neural networks (linear RNNs), specifically the theory and practice of delta-rule architectures for large language models, combined with formal-language and state-tracking expressivity analysis.

Technical level: Advanced. The paper combines matrix-algebra proofs about transition spectra, group-theoretic state-tracking results, and kernel-level implementation engineering.

Scope (one sentence): The paper shows that Kimi Delta Attention's channel-wise diagonal gate can be turned into a rotation mechanism by extending its parameter ranges, yielding "Complex KDA," and characterizes exactly what this buys in expressivity, state tracking, and language modeling at 1.3B parameters.

What This Paper Is About

Linear RNNs built on the delta rule update their recurrent state with a diagonal matrix plus a rank-one correction. That structure is cheap but limited: the transition matrices have a restricted geometry, and prior work needed to compose two delta-rule updates per token (DeltaProduct 2) to obtain a 2D rotation, which raises the rank and the cost of each update. This paper asks whether Kimi Delta Attention (KDA), which already uses a separate gate value per channel, can reach the same rotational expressivity with a single delta-rule transformation. The authors answer yes: by allowing gate entries in [-1,1] and the delta-rule coefficient β in [0,2], a signed channel-wise gate supplies a coordinate reflection that composes with the Householder reflection into a planar rotation, giving a model they call Complex KDA (CKDA).

Key Contributions

1. Spectral characterization and completeness result. The authors characterize the spectrum of CKDA transitions and prove (Theorem 1) that every orthogonal diagonal-plus-rank-one (DPR1) matrix is exactly a CKDA transition matrix, i.e. a signed-Householder form (I − 2kkᵀ)S with ‖k‖₂ = 1 and S = Diag(sᵢ), sᵢ ∈ {−1,+1}. They also study products of CKDA matrices, showing CKDA needs at most max{1, 2n−2} factors in dimension n (and max{1, n−1} for orthogonal matrices, which is sharp).

2. Single-layer and multi-layer state-tracking bounds. One CKDA layer tracks every finite group isomorphic to a subgroup of SO(3), including S₃, S₄ and A₅ (Theorem 3). Three layers solve arbitrary finite-group word problems, and with β > 2 they simulate weighted finite automata in polynomial precision (Theorem 5). These bounds match DeltaProduct 2 and use one layer fewer than (Gated) DeltaNet.

3. An impossibility result. The paper rules out one-layer S₅ tracking for CKDA and for DeltaProduct k with k ≤ 3 when transitions are non-expansive and the set of RNN states is finite (Theorem 4).

4. Empirical and systems validation. Combining both range extensions gives the strongest length extrapolation among the tested KDA range settings on S₃, S₄ and periodic waveform continuation. At 1.3B parameters and 100B tokens of language modeling, CKDA performs on par with a KDA baseline while outperforming Transformers and other linear RNN architectures, retaining 96–97% of standard KDA throughput with a minor modification of KDA recurrence kernels from FLA.

Main Findings

• Scalar gates cannot rotate, diagonal signed gates can. For Gated DeltaNet the gate is a scalar αᵢ = α, so the transition α(I − βkkᵀ) is symmetric with a real spectrum and the scalar factors out of any product. For a channel-wise gate, the commutation condition (H_k D − D H_k)ᵢⱼ = β kᵢ kⱼ (αᵢ − αⱼ) = 0 fails whenever the key has nonzero components on two coordinates with different gate values. Positive gates still give a real spectrum, because H_k D is similar to the symmetric matrix D^(1/2) H_k D^(1/2); signed gates are needed to escape it.

• A 2D rotation appears in closed form. For β = 2, the transition A_{α,θ} = (−α cos 2θ, −sin 2θ; −α sin 2θ, cos 2θ) has discriminant Δ = (1 − α)² cos²(2θ) + 4α < 0 exactly when α < 0. Complex eigenvalues have modulus |λ| = √(−α), moving toward the unit circle as α → −1. At α = −1, Diag(−1,1) = H_{e₁}, and the transition becomes a composition of two reflections whose mirror lines differ by θ, i.e. a rotation by 2θ.

• One rotation plane per transition. Proposition 2 gives a sign-magnitude decomposition A = à Diag(|α|) with à = (I − βkkᵀ)S; the restriction of à to span{k₋, k₊} is a 2×2 block that is a rotation by 2θ when β = 2, with non-real eigenvalues exactly when β² cos²(2θ) < 4(β − 1). Complex eigenvalues require at least two gate coordinates of opposite sign, a key vector spanning coordinates with opposite signs, and β > 1. Theorem 8 shows any CKDA matrix has at most one non-real conjugate eigenvalue pair, requiring both β > 1 and a negative gate entry; Theorem 9 extends this to every non-expansive DPR1 matrix.

• Single-layer group tracking with one fewer layer. Theorem 3: one CKDA layer (one head) tracks every finite group isomorphic to a subgroup of SO(3), realizing every finite cyclic (ℤₙ) and dihedral (Dₙ) group in d = 2, A₄ and S₄ in d = 3, and tracking A₅ in d = 4 (with a many-to-one decoder). S₄ works by reorienting the cube; no orientation works for the icosahedral rotations of A₅, so four dimensions are used instead.

• An obstruction for S₅ in one layer. Theorem 4: if every head transition is non-expansive (‖A‖₂ ≤ 1), has at most one non-real conjugate eigenvalue pair on the unit circle counted with algebraic multiplicity, and reaches only finitely many states, then a single layer with any finite number of independent heads cannot track S₅. The proof builds products R₁, R₂ whose commutator C satisfies C¹⁰ = I but decodes to a three-cycle, a contradiction.

• Depth savings for regular languages and automata. Theorem 5: three CKDA layers solve every finite group-word problem with a fixed exact datatype; allowing β > 2 lets three layers recognize every regular language and compute every WFA over ℚ in polynomial precision. The clock-buffer-accumulator construction uses CKDA's planar rotation as a clock in one layer, versus 4 layers for DeltaNet/GDN.

• Empirical state tracking matches the theory. On S₃, S₄ and A₅ word problems, CKDA (α ∈ [−1,1]ⁿ, β ∈ [0,2]) extrapolates well on S₃ and S₄ while other KDA range settings fail. Extending only the gate or only β gives long-length S₃ scaled accuracy near 0.2, matching parity-only discrimination of the two A₃ cosets. Training lengths reach 32, and accuracies are reported as best of three seeds, scaled from chance (0) to perfect (1).

• Learned models recover the predicted mechanism. On S₃, head 11 approaches the reflection limit β = 2, its channel-wise gate becomes nearly sign-valued with coordinates close to ±1, the first three principal components of the keys explain approximately 95% of the variance, and the resulting transition spectra show complex eigenvalues.

• A₅ needs initialization, not just range extension. Standard training fails to learn A₅; a smaller, fully trainable model initialized near the paper's quaternion construction achieves length extrapolation. The authors note training retains ordinary next-state cross-entropy and a standard MLP readout, and that the architecture and training schedule also differ, so this comparison does not isolate the effect of initialization.

• Throughput cost is small. Forward–backward kernel throughput was measured on an H100 in BF16 with 16 heads, d_k = d_v = 128, and 32k tokens per step. The kernel implementations retain approximately 96–97% of KDA throughput. DeltaProduct 2 achieves competitive throughput with two delta-rule updates per token, helped by its identity-plus-rank-one factors that avoid CKDA's channel-wise gating computations; the authors do not claim an inherent efficiency advantage over it.

• Language modeling. At 1.3B parameters and 100B tokens, CKDA performs on par with a KDA baseline and outperforms Transformers and other linear RNNs, with promising scaling behavior.

Methodology in Plain English

The authors start from the observation that a KDA update is the product of two symmetric matrices: a Householder reflection (I − βkkᵀ) and a diagonal gate matrix. Two symmetric matrices that commute produce something symmetric with only real eigenvalues; two that do not commute can produce rotation-like behavior. They show that with a scalar gate the factors always commute, but with a per-channel gate they generally do not — and once the gate entries are allowed to go negative, the product can have complex eigenvalues sitting on or near the unit circle.

They then prove that this construction is not just capable of rotations but exhaustive for the orthogonal rank-one family, and use that to build explicit group constructions (cyclic, dihedral, S₄, A₅) and depth bounds. For multi-layer results they reuse the standard clock–buffer–accumulator recipe, where a clock counts position modulo some period, a buffer stores block transitions, and an accumulator applies them; CKDA can implement the clock directly with a planar rotation, saving a layer.

On the empirical side they train one-layer models on group word problems, train audio-continuation models on a synthetic periodic groove after a half-bar cue followed by zero inputs, and train 1.3B-parameter language models on 100B tokens. Because KDA stores gate magnitudes in log-space, signed gates required separating sign from magnitude; they implement this three ways — compiled PyTorch operations that absorb cumulative signs into keys and queries without changing kernels, Triton kernels that fuse the signs into normalization and its backward pass, and a TileLang-backed hybrid. They also introduce a "spread" gate initialization with roughly equal proportions of positive and negative entries, and remove the SiLU activation on key and query projections for the state-tracking experiments.

Why This Matters

Impact on research. The paper converts a large open question — what does a channel-wise gate add over a scalar gate in delta-rule models? — into a precise spectral statement, and shows that two previously separate "range extensions" (signed gates and β up to 2) only produce their full effect when combined. It also supplies matching upper and lower expressivity bounds, including a rare impossibility result (no one-layer S₅ tracking under stated assumptions) that constrains what the whole non-expansive DPR1 family can do, not just CKDA. The layer savings for group tracking and regular languages unify several prior results across the literature.

Real-world applications (potential, based on the properties studied here):

• Long-context language models and assistants, where linear RNNs offer linear scaling in sequence length and a fixed-size recurrent state, and better length extrapolation means more stable behavior past the training window.
• Audio and music generation, where the paper's periodic waveform continuation task tests whether complex eigenvalues help preserve phase over long horizons.
• Reasoning and algorithmic tasks that require composing input-dependent operations over time, such as tracking parity, modular arithmetic or permutation composition.
• Efficient inference and training pipelines, since the CKDA kernel retains 96–97% of KDA throughput and is a minor modification of existing FLA recurrence kernels.

Industry relevance. KDA and GDN variants are already deployed in recent large language models cited in the paper (Qwen Team; Kimi Team and others; Upstage Solar Team; Z.ai; inclusionAI). Because CKDA requires only a change to parameter ranges and a sign-aware kernel, it is a low-friction upgrade path for anyone already running KDA, and it offers an alternative to DeltaProduct 2 for obtaining rotational expressivity without a second delta-rule update per token.

Future Directions

1. Closing the A₅ optimization gap. Standard training fails to learn A₅ even though the CKDA construction exists; the paper shows near-construction initialization works. Understanding whether this is an optimization obstacle, a readout limitation, or something else — and how to make it learnable from scratch — remains open.

2. Better long-horizon robustness. The paper notes a gap between the expressive capacity of linear RNNs and their robustness in state tracking, with perturbations that can accumulate and compromise state representations. CKDA's non-expansive transitions and near-unit-circle eigenvalues are relevant to this, but S₄ accuracy still declines at long lengths.

3. Scaling beyond 1.3B parameters. The language-modeling results are at 1.3B parameters and 100B tokens, with "promising scaling behavior" reported. Whether the parity with KDA holds at larger scale, and whether the added expressivity translates into measurable downstream gains rather than parity, is unresolved.

4. Extending the theory past the finite-reachability and non-expansive assumptions. Theorem 4 relies on finite reachability and non-expansiveness; the paper notes that finite reachability holds in most group-tracking constructions in the literature, but whether the impossibility can be circumvented outside those assumptions — and how the one-rotation-plane restriction interacts with more general products or with β > 2 regimes — is a natural next question.

Target Audience

This paper is for researchers and engineers working on efficient sequence-model architectures, especially linear RNNs, state-space models and delta-rule variants such as Mamba, GLA, DeltaNet, Gated DeltaNet, KDA and DeltaProduct. It suits readers with a background in linear algebra, group theory or formal-language expressivity who want to understand what a channel-wise gate buys; practitioners who want a drop-in modification of KDA with a small throughput cost; and theorists interested in expressivity lower bounds for non-expansive diagonal-plus-rank-one recurrences.

Authors’ abstract

Linear RNNs based on the delta-rule enable efficient sequence modeling, but their linear updates with a low-rank correction constrain their expressivity. Prior work has shown that composing two delta-rule transitions in a single recurrent update can model a 2D rotation, but this increases the rank and the cost of the updates compared to a single transition. We show that Kimi Delta Attention (KDA) can realize 2D rotations by combining a single delta-rule transformation with a second reflection supplied by its channel-wise gate. This requires extending the parameter ranges of KDA by combining two existing range extensions: allowing gates in $[-1,1]$ and the delta-rule coefficient $β$ in $[0,2]$. We call the resulting model Complex KDA (CKDA). It preserves KDA's stability and efficiency, with transitions that remain diagonal-plus-rank-one and non-expansive, while reaching the state-tracking expressivity of DeltaProduct$_2$. We characterize the expressivity of CKDA and prove that every orthogonal diagonal-plus-rank-one matrix is exactly a CKDA transition matrix. A single CKDA layer can track every finite group isomorphic to a subgroup of $\mathrm{SO}(3)$, and many state-tracking results use one fewer layer for CKDA compared to other diagonal-plus-rank-one Linear RNNs. Empirically, combining both extensions yields the strongest length extrapolation among tested KDA range settings on $S_3$, $S_4$, and periodic audio continuation. In language modeling, CKDA outperforms Transformers and other linear RNNs, obtains similar results to a KDA baseline, and shows promising scaling behavior. Our code is open source at https://github.com/OpenEuroLLM/ComplexKDA and our models are available at https://huggingface.co/collections/openeurollm/complexkda.