Triadic Linear Attention: Three-Dimensional Recurrent States for Long-Context Sequence Modeling

AI summary

Overview

Research area: Efficient sequence modeling / recurrent neural networks — specifically linear attention and long-context language modeling.

Technical level: Advanced. The paper assumes familiarity with linear attention, associative memory / fast-weight programming, tensor products, and chunkwise-parallel training kernels.

Scope: The paper generalizes the matrix-valued recurrent state of linear attention to a third-order (3D) tensor state built from a triadic outer product of a key, a second key, and a value, and evaluates that construction on language modeling and recall benchmarks at 400M and 1.3B parameters.

What This Paper Is About

Linear attention models compress the past into a fixed-size state, and that state's size determines how much of the context the model can recall. The standard state is a matrix built from a key–value outer product, so enlarging it usually means enlarging all the model's projections. This paper asks whether the state can be enlarged in a parameter-efficient way, and answers by binding each value to a second key, turning the matrix state into a third-order tensor of size d × E × d that is read out by contracting two queries against two key axes.

Key Contributions

1. Triadic linear attention. A generalization of linear attention in which the state is a third-order tensor S_t of size d × E × d, updated by the outer product k_t ⊗ k'_t ⊗ v_t and read by contracting the two key axes with q_t and q'_t. Ordinary linear attention is the special case E = 1 with k'_t = q'_t = 1. An E-dimensional second key yields an E-fold increase in state size (d²E entries instead of d²) while adding only two projections.

2. Compatibility with modern linear attention machinery. The construction extends to data-dependent forgetting (each slice S_t[:, e, :] along the second-key axis gets its own decay gate α_{t,e}), the delta rule (erasing the value currently stored under both keys jointly before writing the new one), and chunkwise-parallel training. Flattening the two keys into a joint key κ_t = k_t ⊗ k'_t shows the delta-rule form corresponds to Gated DeltaNet with key dimension d·E read by a joint query q_t ⊗ q'_t.

3. Efficient kernels. The Kronecker structure of the joint key factorizes the masked attention of a chunk into an ordinary causal matrix masked by tril(Q'K'^T), costing C²(d + E) instead of C²(d·E) operations. The state is tiled along the value axis into blocks of 32 columns, one per thread block, so the full third-order state of a head is never held in one place — at E = 8 the state of one head occupies 512 KiB in FP32 (twice the register file of a Hopper streaming multiprocessor) while one block occupies 128 KiB.

4. Application to two base mixers, plus upcycling and hybrid evaluation. Triadic linear attention is applied to Gated DeltaNet (GDN) and scalar-gated linear attention (sGLA); a pretrained model can be converted post hoc from E = 1 to E = 8 during long-context extension; and a 3:1 GDN/GQA-8 hybrid is tested with either a larger softmax key-value cache or a larger linear-attention state.

Main Findings

• State capacity scales as expected in isolation. On multi-query associative recall (MQAR) with two layers of four heads, d = 16, and no convolutions, forgetting, non-linearities, or delta rule, every doubling of E shifts the accuracy curve to the right by roughly a doubling of N (tested for N from 32 to 4096). At E = 16 the model stores roughly 16 times as many key-value associations as ordinary linear attention, for only a 1.08-fold increase in non-embedding parameters.

• Long-context language modeling improves. GDN predicts early tokens better than the Transformer but falls behind at around 10k tokens. With E = 2, triadic GDN pushes the crossover considerably further out; with E = 8 it predicts the next token better than the Transformer even at 64k context. Gains grow with state size but show diminishing returns.

• Triadic variants beat other ways of enlarging the state at matched state size. At both 2× and 4× state, triadic GDN achieves lower WikiText perplexity and lower PG19 perplexity in every context range and the highest recall average (Triadic E = 2: recall 28.4 vs. GDN base 26.2; Triadic E = 4: recall 31.1 vs. 26.2). Alternatives (larger heads, wider values, grouped values, more heads) barely improve on vanilla GDN at 2× and at 4× they degrade on every PG19 range. At 4×, "wider values" and "grouped values" force the MLP width down from 2816 to 640 to stay within the parameter budget, whereas triadic linear attention adds only two small projections and preserves the full MLP. The same pattern holds for sGLA, where triadic sGLA has the lowest perplexity at both state sizes.

• Recall-intensive tasks gain sharply. Triadic GDN substantially improves recall-intensive tasks at both scales, with particularly large gains on FDA and SWDE, which require copying information from a long document. On RULER needle-in-a-haystack, a larger state keeps retrieval accurate up to longer contexts. Recall averages at 400M: GDN base 26.2, Triadic E=2 28.4, Triadic E=4 31.1, Triadic E=8 from scratch 33.1; at 1.3B: GDN base 36.1, Triadic E=8 from scratch 44.4.

• Upcycling works but recovers only part of the benefit. Expanding a pretrained model from E = 1 to E = 8 (copying forget gates to every slice, initializing the second key/query projections and convolutions from scratch, then long-context extension) consistently improves over the base model at both mixers and both scales. Upcycled models recover about half to three quarters of the gain of a triadic model pretrained from scratch (400M GDN: upcycled recall 29.7 vs. from-scratch 33.1 vs. base 26.2; 1.3B GDN: upcycled 42.0 vs. from-scratch 44.4 vs. base 36.1).

• In hybrids, enlarging the linear-attention state beats enlarging the key-value cache. In a 3:1 GDN/GQA-8 hybrid at 400M, the triadic variant (E = 4) has the lowest perplexity at every context range and the best zero-shot (53.8) and NIAH averages (56.3), and trails the GQA-4 variant only on recall (44.5 vs. 45.3, described as within seed spread). The GQA-4 hybrid requires more memory than the triadic hybrid beyond about 4.6k tokens and almost twice as much at 64k tokens (407 MB vs. 220 MB).

• Training overhead is moderate. GDN matches the Transformer at 2k tokens; triadic GDN overtakes it at 4k for E = 2 and E = 4 and at 8k for E = 8. At E = 8 it is 3% slower than the Transformer at 4k and faster at every longer context, reaching 5.1 times faster at 64k. Overhead over vanilla GDN is 28%–30% for E = 8, 14%–15% for E = 4, and 9%–11% for E = 2.

• Key-dimension trade-off is mild, activation choice matters more. At a fixed joint key dimension d_k · E = 1024, moving from d_k = 128, E = 8 to d_k = 32, E = 32 slightly degrades perplexity (WikiText 10.80 to 10.90) and lowers recall (33.4 to 32.5). Non-negative activations on the second key and query (softplus, sigmoid) consistently outperform activations that can take negative values (SiLU, no activation) in perplexity — a possible explanation being that signed entries let a token read and write to different slices with opposite signs, causing cancellation when the second query aggregates slices.

• Costs and limitations are acknowledged. The larger state slows training by around 15% for E = 4 and around 30% for E = 8 despite the optimized kernels, and performance on some recall-intensive tasks still trails Transformers, though only with a far larger state size at long context.

Methodology in Plain English

The authors start from the observation that linear attention is an associative memory: each key–value pair is stored as an outer product and retrieved by matching the query against the key, with capacity limited by the state's shape (a d × d matrix can separate at most d mutually orthogonal keys). Because a binding is itself a vector, it can be bound to a further role, which raises the order of the memory. The paper binds each value to a second key as well, giving a d × E × d tensor that can in principle store d² associations, and unbinds with two queries.

They first validate the capacity claim in a deliberately minimal MQAR setup — two layers, four heads, d = 16, no forgetting, no delta rule, no non-linearities — sweeping E over {2, 4, 8, 16} and N from 32 to 4096. They then extend the construction to scalar data-dependent forgetting (one decay gate per slice along the second-key axis) and the delta rule (erase the joint two-key readout before writing), and derive chunkwise-parallel training forms.

For efficiency, they exploit the Kronecker structure of the joint key so that within-chunk attention factorizes into an ordinary linear attention term masked by an E-dimensional correction matrix R', and they tile the third-order state along the value axis into 32-column blocks, one per thread block, so no single unit holds a full head's state.

The empirical evaluation trains GDN and sGLA variants at 400M (24 layers, d_model = 1024, 8 heads) and 1.3B (24 layers, d_model = 2048, 16 heads) parameters with head dimension d = 128, for E ∈ {1, 2, 4, 8}. Training uses 50 tokens per parameter (2.5× Chinchilla-optimal, i.e., 20B tokens at 400M and 65B at 1.3B) on Fineweb-Edu at 4k context, then long-context extension to 64k on 5 tokens per parameter using Fineweb-Edu, PG19 books, and scientific PDFs, with AdamW (peak learning rate 3×10⁻⁴ pretraining, 10⁻⁴ for extension, weight decay 0.1, linear warmup then cosine decay to 10% of peak, effective batch about 0.5M tokens at 400M and 1M at 1.3B). Evaluation uses PG19 perplexity by context position, WikiText perplexity, the recall benchmark suite of Arora et al. (2024), ten zero-shot tasks, and RULER needle-in-a-haystack. Baselines include Transformers with RoPE, QK-norm and grouped-query attention (eight query heads per key-value head), and state-matched alternatives at 2× and 4× state.

Why This Matters

The paper identifies state size as a distinct, controllable axis for improving recurrent sequence models, and shows that it can be enlarged with negligible parameter overhead rather than by widening every projection. This reframes a design question that has usually been answered by scaling heads or value dimensions.

Real-world applications suggested by the properties studied:

• Long-document question answering and analysis, where models must retrieve facts scattered across tens of thousands of tokens (the paper's own recall suite includes tasks requiring copying from a long document).
• Retrieval-augmented generation, where the fixed-size state replaces a growing key-value cache and therefore bounds memory per sequence.
• Streaming or long-session inference, since constant-time inference with a fixed state avoids an unbounded cache; the paper's hybrids show the triadic state using 220 MB at 64k tokens versus 407 MB for the enlarged softmax cache.
• Hybrid deployment stacks, where linear-attention layers are interleaved with softmax attention blocks — the paper's finding that enlarging the linear-attention state beats doubling key-value heads is directly relevant to how such stacks are budgeted.

Industry relevance follows from the training-efficiency measurements: triadic GDN is 3% slower than a Transformer at 4k at E = 8 and 5.1 times faster at 64k, and the Triadic GDN/GQA-8 hybrid at 400M scores 56.3 on NIAH average and 53.8 zero-shot while using less long-context memory than the GQA-4 alternative.

Future Directions

• Revisiting other matrix-state innovations for three-dimensional states. The authors note they applied the construction only to GDN and sGLA, a small subset of the linear attention literature, and that the additional axis opens up possibilities for entirely new variants.
• Making the state-size increase cheaper. Training still slows by around 15% for E = 4 and 30% for E = 8 even with the custom kernels, so kernel and system work remains open.
• Closing the remaining recall gap. Performance on some recall-intensive tasks still trails Transformers, though only when the Transformer uses a far larger state at long context.
• State growth as a training-stage schedule. Because short-context data requires little state but makes up most of pretraining, the authors suggest a model could be pretrained with a small state that is enlarged in later training stages with higher sequence lengths, which the upcycling experiment partially supports.
• Factorizations with smaller projections, which the key-dimension ablation suggests are slightly worse in perplexity but may be relevant for mixture-of-experts architectures.

Target Audience

Researchers and engineers working on efficient attention, linear RNNs, and long-context language modeling, particularly those implementing chunkwise-parallel kernels or designing hybrid architectures that mix linear-attention and softmax-attention layers. It is also relevant to readers interested in associative memory theory, since the paper connects higher-order correlation-matrix memory capacity to practical sequence-mixer design. Readers without background in linear attention, fast-weight programming, or tensor contractions will find the derivations difficult.

Authors’ abstract

Recurrent neural networks (RNNs) compress the historical context into a memory state of fixed size, thus allowing for constant-time inference. The memory state size is a crucial factor in their performance, as exemplified by the strong performance and resurgence of linear attention, which extends the vector-valued hidden states of ordinary RNNs to matrix-valued hidden states. Crucially, linear attention does so in a parameter-efficient way, in particular by using an outer product of the key and value vectors to write to the matrix-valued hidden state. We generalize this construction and propose triadic linear attention, which writes the triadic outer product of a key, a second key, and a value, into a third-order (i.e., 3D) tensor state, and reads from it by contracting both key axes with two queries. An $E$-dimensional second key thus yields an $E$-fold increase in state size while adding only two projections. Triadic linear attention is compatible with data-dependent forgetting, the delta rule, and chunkwise-parallel training. Applied to Gated DeltaNet and scalar-gated linear attention, triadic linear attention substantially improves long-context language modeling and recall, outperforming alternatives that enlarge the state.