Triadic Linear Attention: Three-Dimensional Recurrent States for Long-Context Sequence Modeling
Oliver Sieberling, Bharat Runwal, David Jin, Ryan Chin, Rameswar Panda, Yoon Kim
cs.LG, cs.CL
2026-09-29
A key × second-key × value outer product fills a 3D tensor state, buying E-fold capacity for two projections; E=8 adds 1.2% params and wins on long-context LM and recall.
Linear attention swaps softmax for an outer-product read and write, compressing the whole history into a fixed-size matrix state so inference cost stops growing with context. The catch is capacity: a d×d state separates at most d orthogonal keys, and prior work shows a fixed-state model cannot even copy sequences past a certain length. Gated DeltaNet and Mamba-2 manage that state with forget gates and the delta rule, but the substrate is still one d² matrix, and on recall-heavy tasks they clearly trail Transformers.
Enlarging it the obvious way is expensive. More heads or wider values inflate every projection of the sequence mixer; in the 4× state comparison, two baselines needed so many extra parameters that the MLP width dropped from 2816 to 640 to stay on budget. Hence the question: can the state grow without paying in parameters?
The starting point is an observation: the matrix state of linear attention is the vector state of a classic RNN lifted one order by the outer product of key and value, and the construction generalizes.
Two moves keep training fast. The Kronecker structure of the joint key, (q⊗q′)ᵀ(k⊗k′) = (qᵀk)(q′ᵀk′), holds the intra-chunk cost at C²(d+E) instead of C²·dE. And the state is tiled along the value axis into 32-column blocks that stay in registers, so no streaming multiprocessor ever holds the full state; at E=8 one head's state is 512 KiB in FP32, twice the register file of a Hopper SM.
An ablation pins down one detail: the second key and query need a non-negative activation (softplus or sigmoid). Signed entries let a token read and write different slices with opposite signs, and the contributions cancel at aggregation.
| 400M model | WikiText ↓ | PG19 16k–64k ↓ | Recall ↑ |
| Transformer | 11.11 | 13.92 | 41.6 |
| GDN (base) | 11.25 | 14.15 | 26.2 |
| 4× state: larger heads (d=512) | 11.25 | 14.24 | 28.5 |
| 4× state: grouped values | 11.35 | 14.36 | 27.3 |
| 4× state: Triadic (E=4) | 10.93 | 13.79 | 31.1 |
At matched parameter count and state size, triadic wins every PG19 range and the recall suite; at 4× state the three conventional routes all do worse than vanilla GDN. On synthetic MQAR, E=16 stores roughly 16× the key-value associations for 1.08× the non-embedding parameters, and E=8 costs 1.2% more parameters at 400M. On PG19 books, triadic GDN with E=8 predicts the next token better than the Transformer across the full 64k context, while vanilla GDN is overtaken near 10k; at 1.3B, recall rises from 36.1 to 44.4. A pretrained GDN can be upcycled by expanding E from 1 to 8 during long-context extension, recovering half to three quarters of the from-scratch gain (42.0 vs 44.4 at 1.3B, single seed). The learned gates develop half-lives from 0.20 to 1080 tokens across the eight slices, where GDN's single gate sits at a median of 5.1, so the new axis really holds multiple timescales. Speed at 1.3B on one H100: E=8 is 28–30% slower than vanilla GDN, 3% slower than the Transformer at 4k, and 5.1× faster at 64k.
For anyone building on linear attention or hybrid stacks (Kimi Linear, Qwen3-style layering of GDN and attention), this is a state-size knob that costs almost nothing in parameters, with open kernels. The hybrid experiment answers a practical question directly: at equal budget, quadrupling the GDN state beats doubling the attention KV heads on perplexity, zero-shot and needle-in-a-haystack (56.3 vs 53.0), while using nearly half the memory at 64k. The upcycling result suggests a recipe: pretrain with a small state and expand during long-context extension, mirroring how a KV cache grows with context. The authors also float a broader reading: part of the gains attributed to test-time training variants may simply be state-capacity effects. This is architecture research capped at 1.3B on 65B tokens, but the direction is clear and the code is public.
The authors name three: E=8 slows training by about 30%; some recall tasks still trail the Transformer (41.6 vs 33.1 at 400M, a real gap); only GDN and sGLA were tested. Two gaps go unaddressed. Decode throughput is never reported: the state grows from 6.3 MB to 50.3 MB at 400M, per-step state traffic grows proportionally, and the paper does not quantify that cost at inference. And language-modeling gains diminish as the state grows, while the 1.3B upcycling number comes from a single seed, so whether the recall gains carry to larger scale remains open.