Source-linked AI summary

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

Oliver Sieberling, Bharat Runwal, David Jin, Ryan Chin, Rameswar Panda, Yoon Kim

arXiv:2609.36529v1cs.LGcs.CL

TL;DR

Linear attention’s fixed-size state limits recall, while naive state enlargement can significantly increase parameter count. Triadic linear attention writes a key, a second key, and a value into a tensor state; at 400M and 1.3B scales, it improves perplexity and recall while outperforming alternative enlargement methods.

  • Problem

    Linear attention models still struggle on recall-intensive and long-context tasks, while naively increasing state size can substantially increase parameter count.

    Source: §1 · p. 1
  • Method

    Triadic linear attention writes the outer product of a key, a second key, and a value into a third-order state, reading it with two queries and adding only two projections.

    Source: §1 · p. 1
  • Results

    Within 1.3 times vanilla Gated DeltaNet training time at E = 8, triadic Gated DeltaNet consistently lowers perplexity and improves recall at 400M and 1.3B scales.

    Source: §4 · p. 9, §6 · p. 10
  • Takeaways & Limitations

    State size is an important axis for improving linear RNNs, particularly on recall-intensive tasks and at long context.

    Source: §4 · p. 9

    Larger states slow training by around 15% at E = 4 and 30% at E = 8, and the construction has only been applied to GDN and sGLA.

    Source: §4 · p. 9

Abstract

from arXiv · show

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.

1 INTRODUCTION

Triadic linear attention expands linear attention’s matrix state into a third-order tensor using only two additional projections. Applied to Gated DeltaNet and scalar-gated linear attention, it improves long-context language modeling and recall while outperforming other state-enlargement approaches.

  • Linear attention models can struggle with recall-intensive and long-context tasks, while naively enlarging their state can substantially increase parameter count.Adding heads enlarges mixer projections, and increasing value dimensions or adding value heads grows value and output projections.Source: §1 · p. 1
  • An E-dimensional second key expands the state E-fold while requiring only two additional projections.The model writes the outer product of a key, second key, and value into a third-order state, then reads by contracting both key axes with two queries.Source: §1 · p. 1
  • Triadic linear attention supports data-dependent forgetting, the delta rule, and chunkwise-parallel training.It uses separate forget gates for second-key slices, removes the jointly stored value before delta-rule writes, and tiles state along the value axis for chunkwise-parallel execution.Source: §1 · p. 2
  • Triadic linear attention improves long-context language modeling and recall in Gated DeltaNet and scalar-gated linear attention, outperforming alternatives that enlarge the state.In a GDN/Transformer hybrid, triadic GDN lowers long-book perplexity more than doubling key-value heads while using significantly less memory at long contexts.Source: §1 · p. 2

2 TRIADIC LINEAR ATTENTION

Triadic linear attention binds each value to two keys in a third-order state, read by contracting both key axes with queries. The added state capacity is paired with small parameter overhead and compatibility with forgetting, the delta rule, and chunkwise-parallel training.

2.1 LINEAR ATTENTION AS AN ASSOCIATIVE MEMORY

  • Linear attention stores associations in a d²-entry matrix state, which can separate at most d mutually orthogonal keys before interference is inevitable.The tensor-product view treats values as fillers and keys as roles; retrieval contracts the state with a query, and non-orthogonal queries produce interference.Source: §2.1 · p. 2

2.2 BINDING TO A SECOND KEY

  • With orthonormal roles, an n-th-order tensor product representation stores d^(n−1) associations exactly, so the second key raises capacity from d to d² associations.The second key need not correspond to anything in the data; the model can choose how to use both keys.Source: §2.2 · p. 3
  • Triadic linear attention writes the outer product of a key, second key, and value into a d×E×d state and reads it with two queries.The second key and query are produced through projections; the state has d²E entries rather than d², while parameter overhead comes from two additional projections and is negligible for small E such as 8.Source: §2.2 · p. 3

2.3 CAPACITY IN ISOLATION

  • 16-fold state-capacity increase comes with only a 1.08-fold increase in non-embedding parameters for E=16 on MQAR.Figure 1 compares plain linear attention with state 16×16 against triadic states 16×E×16; each doubling of E shifts the accuracy curve right by roughly a doubling of N.Source: Figure 1 · p. 3, §2.3 · p. 3

2.4 FORGETTING AND THE DELTA RULE

  • Forgetting assigns each second-key slice its own decay gate, while the delta rule removes the association stored under both keys before writing the new one.The delta update can use a data-dependent write strength and can be combined with decay before erasure.Source: §2.4 · p. 4

2.5 EFFICIENT IMPLEMENTATION

  • Chunkwise-parallel training computes within-chunk interactions in parallel and carries recurrent state between chunks.Factorizing joint-key inner products makes the masked attention work with C×C matrices; forming them costs C²(d+E) operations rather than C²(d·E), with an additional C²E cost for R′.Source: §2.5 · p. 4
  • Tiling confines the E-fold cost to the state, whose value-axis blocks keep all E slices in registers without holding the full state in one place.At E=8, one head’s state is 512 KiB in FP32, while a 32-column block occupies 128 KiB.Source: §2.5 · pp. 4–5

3 EMPIRICAL STUDY

The study evaluates triadic linear attention in long-context language modeling, recall, hybrid models, and training efficiency. Across these settings, larger triadic states improve performance, while the method remains competitive in parameter use and long-sequence speed.

3.1 EXPERIMENTAL SETUP

  • The experiments compare GDN and sGLA triadic variants at 400M parameters and triadic GDN at 1.3B, alongside state-matched alternatives and long-context evaluations.E = 8 adds 1.2% more parameters; evaluation includes PG19, WikiText, recall, and zero-shot tasks.Source: §3.1 · pp. 5–6

3.2 MAIN RESULTS

  • With E = 8, triadic GDN predicts better than the Transformer at 64k context, while larger states improve recall and keep RULER retrieval accurate at longer contexts.The Transformer’s key-value cache exceeds the recurrent model’s state beyond the indicated context lengths; language-modeling gains diminish as state grows.Source: §3.2 · pp. 6–7, Figure 3 · p. 7
  • At matched parameter counts, triadic GDN achieves lower perplexity on WikiText and every PG19 range at 2× and 4× state, and triadic sGLA has the lowest perplexity at both sizes.Triadic GDN also has the highest average recall score; Table 1 compares state-enlargement methods at matched parameter counts.Source: Table 1 · p. 6, §3.2 · p. 7

3.3 UPCYCLING TO TRIADIC LINEAR ATTENTION VIA CONTINUED PRETRAINING

  • Upcycling pretrained models from E = 1 to E = 8 improves both sequence mixers at both scales and recovers about half to three quarters of scratch-pretraining gains.The conversion copies forget gates across slices and initializes the added key/query projections and convolutions from scratch.Source: §3.3 · pp. 7–8, Table 2 · p. 7

3.4 HYBRID MODELS

  • In 3:1 GDN/GQA-8 hybrids, triadic E = 4 has the lowest perplexity at every context range and best zero-shot and NIAH averages, while using less memory than GQA-4 beyond about 4.6k tokens.It trails GQA-4 only on recall, within seed spread; Table 3 reports the hybrid comparison and state at 4k and 64k tokens.Source: §3.4 · p. 8, Table 3 · p. 8

3.5 TRAINING EFFICIENCY

  • At E = 8, triadic GDN is 5.1 times faster than the Transformer at 64k tokens, with 28%–30% overhead over vanilla GDN.Measurements use forward and backward time per block at batch size 4 on one H100; E = 8 is faster than the Transformer at contexts longer than 4k.Source: §3.5 · p. 8, Figure 4 · p. 8

3.6 ABLATIONS

  • At fixed joint key dimension, making the first and second key dimensions more similar slightly degrades perplexity; smaller key/query projections may still suit mixture-of-experts architectures.The ablation varies dk and E at dk · E = 1024 while matching parameters by adjusting MLP width.Source: §3.6 · p. 9, Table 4 · p. 9
  • For second-key and query activations, softplus and sigmoid consistently outperform SiLU and no activation in perplexity.The authors suggest signed entries may allow writes and reads to cancel when the second query aggregates slices.Source: §3.6 · p. 9, Table 4 · p. 9

4 DISCUSSION AND LIMITATIONS

State size is an important axis for improving linear RNNs, especially on recall-intensive and long-context tasks. Triadic linear attention improves perplexity and recall at two model scales, while its larger state slows training and some recall results still trail Transformers.

4 DISCUSSION AND LIMITATIONS

  • At 400M and 1.3B parameters, triadic linear attention consistently improves perplexity and recall and outperforms other state-enlargement methods.The reported gains are especially relevant to recall-intensive tasks and long contexts.Source: §4 · p. 9
  • Training is around 15% slower for E = 4 and around 30% slower for E = 8, even with optimized kernels.Some recall-intensive tasks still trail Transformers, but only with a far larger state size at long context.Source: §4 · p. 9

5 RELATED WORK

Related work situates triadic linear attention among linear-attention variants, larger-memory methods, test-time training, and higher-order associative memories. The paper distinguishes its single dense state from approaches that route among memory slots or use very large sparse memories.

Linear attention

  • Linear attention replaces softmax over attention logits with a kernel feature map and writes to a matrix state using a key–value outer product; triadic states can be viewed through an expanded key feature map.The feature map expands [k, k′] into k ⊗ k′; linear-attention updates coincide with the fast weight programmer up to an omitted normalization factor (Schlag et al., 2021).Source: §5 · p. 9
  • Modern Gated DeltaNet combines the delta rule with data-dependent forgetting, and later work generalizes its state transitions and parametrizations.The related-work passage traces forgetting from fixed decay to data-dependent gates, then describes identity-plus-low-rank transitions in later variants.Source: §5 · pp. 9–10

Test-time training

  • Test-time-training variants range from delta-rule gradient steps to multiple steps or optimal regression, while many other variants use substantially larger states.Because a gradient step can be interpreted as a linear-attention update (Irie et al., 2022), the authors suggest some TTT gains may reflect larger state capacity, particularly at long context.Source: §5 · p. 10

Increasing state size

  • State-enlargement alternatives include wider keys or values, more heads, memory slots with routing, and sparse product-key addressing.The paper uses dense variants as main baselines and contrasts routed approaches' generally higher state size and substantial memory-movement overhead with triadic attention's single dense state and efficient fused kernels.Source: §5 · p. 10

Higher-order associative memories

  • Higher-order associative memories raise capacity through higher-order interactions, with storable patterns scaling as d^(n−1) for d memory units and interaction order n.Schlag & Schmidhuber (2018) used a third-order tensor state for graph traversal on small-scale language tasks; this paper instead applies the construction to modern linear-attention layers trained at scale.Source: §5 · p. 10

6 CONCLUSION

The paper introduces triadic linear attention as a parameter-efficient way to enlarge the state of constant-time sequence mixers. Experiments report lower perplexity and better recall across two model scales, with custom kernels limiting training time for E = 8.

6 CONCLUSION

  • Triadic linear attention updates a third-order tensor state with the outer product of three vectors, increasing the capacity of constant-time sequence mixers.The construction adds little parameter overhead and admits efficient kernels that exploit its structure.Source: §6 · p. 10
  • Across 400M- and 1.3B-parameter language-modeling experiments, triadic Gated DeltaNet consistently lowers perplexity and improves recall.For E = 8, custom kernels keep training time within 1.3 times that of vanilla Gated DeltaNet.Source: §6 · p. 10
  • The authors present triadic linear attention as a starting point for larger-state linear-attention variants and, ultimately, pure linear-attention models at the frontier.Source: §6 · p. 10

A CHUNKWISE FORM AND GPU KERNELS

The appendix develops a chunkwise form of triadic linear attention with per-slice forgetting and the delta rule, then describes numerical-stability choices, input projections, and GPU kernels. Its kernels implement the chunkwise computation and are checked against FP64 references.

Chunkwise form

  • The chunkwise form uses C = 64-token chunks, separates each slice's joint-key interactions into d-dimensional key products and second-key terms, and reuses a C × C matrix across E slices.The transform is formed from d-dimensional keys and C^2E terms rather than d · E-dimensional joint keys.Source: A CHUNKWISE FORM AND GPU KERNELS · p. 17
  • The chunkwise delta rule applies its transform to the residual; with E = 1 and k′ = q′ = 1, the update reduces to the chunkwise form of Gated DeltaNet.Without forget gates, the per-slice decay matrix reduces to the causal mask.Source: A CHUNKWISE FORM AND GPU KERNELS · p. 17

Numerical stability

  • Per-slice scalar decay factors are computed as exponentials of non-positive differences of accumulated log gates, so they cannot overflow and need no clamp.Each slice decays by a single scalar, avoiding the secondary chunking required when decay varies along the joint key.Source: A CHUNKWISE FORM AND GPU KERNELS · p. 17
  • A cheaper midpoint-reference construction can overflow in FP32 and, after clamping, distort near-diagonal terms under strong within-chunk decay.In an FP64 simulation, the clamped construction's relative output error is 0.05 at decay e^-200 and 1.5 at e^-600; the authors therefore form R and R′ element by element.Source: A CHUNKWISE FORM AND GPU KERNELS · p. 17

Kernels

  • Figure 5 plots final-4k-token loss changes as preceding context grows from 4k to 64k, delta-strength distributions, and state-slice half-lives ranked within each head.Source: Figure 5 · p. 18

Input projections

  • Triadic layers use the same two projection GEMMs and one convolution as Gated DeltaNet, with wider projections for the second keys and queries.One GEMM produces q, k, v, k′, and q′, and a causal depthwise convolution processes them; SiLU is applied only to q, k, and v channels.Source: A CHUNKWISE FORM AND GPU KERNELS · p. 18

Kernels

  • The Hopper GPU kernels parallelize remaining chunk gradients after a reverse recurrence and produce bitwise-reproducible gradients without floating-point atomics.At E = 8, the forward kernel assigns state-slice products, input loading, and output computation across four warpgroups; output and gradients of all seven inputs were verified against FP64 references with relative ℓ2 error below 1%.Source: A CHUNKWISE FORM AND GPU KERNELS · p. 18

B ANALYSIS

On held-out PG19, Triadic GDN benefits more from added context than GDN, while its additional state slices span a broad range of memory timescales.

Long-context utilization

  • GDN begins to saturate by 16k context tokens, while Triadic GDN continues improving through 32k and gains consistently more from added context.The comparison uses three seeds and held-out PG19 books with at least 64k tokens.Source: B ANALYSIS · p. 18

Memory update dynamics

  • Mean β rises from 0.31 to 0.50; across eight slices, median half-life spans 0.20 to 1080 tokens, versus 5.1 tokens for GDN.The β mean is averaged across evaluated tokens, heads, and layers; half-life measures how many tokens it takes forget gates to halve a state contribution.Source: B ANALYSIS · pp. 18–19
  • The additional state slices have broadly distributed write strengths and timescales, suggesting they permit stronger memory updates while retaining information across a wider range of timescales.Triadic GDN’s per-token β distribution is broader than GDN’s.Source: B ANALYSIS · pp. 18–19

C MQAR SETUP

The MQAR setup tests whether models retrieve values for queried keys after accumulating randomly ordered key-value pairs, using a compact model and a broad sweep of sequence lengths and second-key dimensions.

MQAR sequence construction

  • Each example has 2N tokens: N key-value pairs in random order, followed by the corresponding N key queries, for which the model predicts values.Keys and values come without replacement from distinct 8k-token sets; queries are read-only and all read the same accumulated recurrent state.Source: C MQAR SETUP · p. 19

Model architecture

  • The model uses a vanilla triadic mixer with d_model = 128, four heads, D = 16, and E in {1, 2, 4, 8, 16}; E = 1 recovers ordinary linear attention.The per-head state has shape 16 × E × 16, and models range from 3.51M to 3.54M parameters.Source: C MQAR SETUP · p. 19

Training and evaluation

  • Models train for 50k steps on 12.5 billion tokens, sweeping N from 32 to 4096 and E in {1, 2, 4, 8, 16} across 25 seeds per cell.Reported test accuracy is the mean over seeds on a held-out set of 3k sequences.Source: C MQAR SETUP · p. 19

D LANGUAGE MODELING SETUP

The language-modeling setup applies triadic GDN and sGLA blocks in a 24-block architecture, then evaluates long-context language modeling, retrieval, recall, and zero-shot tasks.

Model architecture

  • Models use 24 pre-norm blocks with RMSNorm, a sequence mixer, and SwiGLU MLP; at E = 1, normalized second queries and keys reduce the triadic variant to the base mixer.The Transformer baseline uses grouped-query attention, QK-norm, and RoPE, with base theta 10k during pretraining and 2M for extension; all models use the Llama-2 tokenizer.Source: D LANGUAGE MODELING SETUP · pp. 19–20

Sequence mixer blocks

  • Each triadic modeling block uses a per-head D×E×D state, forgets state along the second-key axis, and reads by contracting both key axes with queries.The block applies a delta-rule write for GDN and a β-gated write for sGLA.Source: D LANGUAGE MODELING SETUP · p. 20

Evaluation

  • Evaluation covers PG19 and WikiText perplexity, average zero-shot accuracy across ten tasks, RULER retrieval, and six recall tasks.PG19 uses held-out books on 64k-token windows; RULER averages eight tasks and context lengths from 1k to 64k tokens.Source: D LANGUAGE MODELING SETUP · p. 20
Loading 2609.36529v1…