Source-linked AI summary
Linear Attention Architectures: Mechanisms, Trade-offs, and Cross-Layer Routing
Tommaso Cerruti, Tim Rieder, George Rowlands, Lingfeng Jin, Imanol Schlag
TL;DR
Long-context modeling needs more efficient sequence mixers and lightweight ways to share information across recurrent layers. This paper compares attention architectures in a common framework and finds trade-offs among loss, throughput, stack design, and routing, with CLVR modestly improving matched DeltaNet and Gated DeltaNet runs.
Problem
Quadratic softmax attention motivates efficient long-context sequence mixers, while DeltaNet-style architectures raise the question of lightweight cross-layer information sharing that preserves linear-time structure.
Method
The paper unifies five attention mechanisms in recurrent-memory notation and evaluates lightweight cross-layer routing, including error and value routing for DeltaNet-style memories.
Results
Kimi Delta Attention with Muon achieves the best validation loss, pure Gated DeltaNet with AdamW the fastest normalized training, and CLVR modestly lowers loss for DeltaNet and Gated DeltaNet.
Takeaways & Limitations
Architecture choice should match the target operating point, balancing validation loss, throughput, optimizer, stack pattern, and metric semantics rather than selecting one universally dominant setting.
Takeaways & Limitations
The empirical comparisons use single runs with uneven hyperparameter coverage, and the limited downstream tasks may miss long-context or memory-intensive behaviors.
Abstract
from arXiv · showhide
Self-attention lets each token retrieve information from the full context, but its quadratic cost in sequence length limits training and inference at long context. This paper presents a comparative study of softmax attention and four recent recurrent linear-attention architectures: DeltaNet, Gated DeltaNet, Kimi Delta Attention, and Gated DeltaNet-2. We express these mechanisms in a common recurrent-memory notation, making explicit how they differ in expressivity, memory decay, erase and write control, training throughput, and implementation complexity. Our experiments center on 350M-parameter models trained for 15B tokens, and include optimizer and learning-rate comparisons, hybrid-versus-pure stack comparisons, sequence-length runtime measurements, larger DeltaNet runs at 1.3B and 3B parameters, and a small set of downstream evaluations. The reported speed results measure training throughput and iteration time; we do not provide an empirical inference-speed benchmark. Within the reported 350M-parameter, 15B-token sweep, Kimi Delta Attention with Muon reaches the lowest final validation loss, a pure Gated DeltaNet stack trained with AdamW has the highest normalized training throughput, hybrid stacks generally improve loss at a throughput cost, and Muon consistently lowers final validation loss relative to AdamW in the matched architecture settings we evaluate. We introduce and evaluate lightweight cross-layer routing mechanisms for DeltaNet-style memories. The most natural DeltaNet-inspired formulation, forwarding a lower layer's delta-rule write error into the next layer's value target, does not improve over matched baselines. Routing into the aligned hidden stream and forwarding the write value instead yields a modest improvement in the matched runs we report: Cross-Layer Value Routing (CLVR) lowers final validation loss for both DeltaNet and Gated DeltaNet.
1 Introduction
The paper studies recurrent linear-attention alternatives to expressive but quadratically costly softmax attention, organizing their mechanisms, trade-offs, and empirical behavior in a shared framework. It also introduces lightweight cross-layer routing for DeltaNet-style memories and identifies write values as more useful than write errors for routing.
- Motivation: Softmax attention enables token-to-token information exchange but requires pairwise comparisons between every input-token pair, producing quadratic cost.The introduction frames linear attention as a response to this efficiency challenge.
- Mechanism: A common recurrent-memory notation compares softmax attention, DeltaNet, Gated DeltaNet, Kimi Delta Attention, and Gated DeltaNet-2 through memory, residual, decay, erase, and write mechanisms.The framework isolates recurrent memory W and delta-rule residual r.
- Empirical behavior: The study maps trade-offs across validation loss, throughput, optimizer and learning-rate sensitivity, stack structure, sequence-length timing, scaling, and downstream behavior.The comparison covers both hybrid-versus-pure stacks and larger DeltaNet runs.
- Cross-layer routing: Cross-Layer Error Residuals forward a lower layer’s delta-rule write error into the next layer’s value target, but this formulation does not improve over matched baselines.The paper then routes into the aligned hidden stream and uses the layer’s write value instead.
- Scope: The paper’s aim is to make the design space legible rather than declare a single best architecture, while clarifying how design choices relate to observed properties and scaling behavior.It organizes the comparison around mechanism, trade-offs, and empirical behavior.
2 Background: From Softmax Attention to Linear Attention
This section contrasts expressive but quadratic-cost softmax attention with linear attention, which uses recurrent memory for constant per-token cost and linear total training cost. It then motivates DeltaNet-style delta-rule updates as error-correcting memory writers with increasingly explicit forgetting and editing control.
- Softmax attention: Softmax attention compares each query with every previous key, enabling full-context retrieval but requiring O(T^2) query-key interaction cost during training.Its normalized weighted averaging and exponential kernel provide soft selection and sharp attention distributions.
- The kernel view: Linear attention rearranges feature-map computations so each query interacts with one accumulated memory matrix instead of all previous tokens individually.The memory can be viewed as either a compressed record of key-value pairs or fast weights.
- Recurrent memory: Linear attention uses additive memory writes and a single matrix-vector product per output, giving constant per-token cost and linear total training cost in T.Without learned decay or erase, old associations remain active unless overwritten indirectly.
- From additive writes to delta-rule writes: DeltaNet replaces raw value writes with residual writes, turning memory updates into error-correcting corrections based on what the current memory already predicts.Later variants add scalar decay, channel-wise decay, or separate channel-wise erase and write gates.
3 Architectures
The section contrasts expressive but quadratic softmax attention with recurrent linear-attention architectures that update memory incrementally. DeltaNet, Gated DeltaNet, Kimi Delta Attention, and Gated DeltaNet-2 progressively add forgetting and finer-grained erase/write control, trading selectivity for flexibility.
- Standard softmax attention: Softmax attention retrieves selectively from the full previous context but incurs quadratic training cost in sequence length.Linear-attention variants replace explicit pairwise comparisons with incrementally updated recurrent memory states.
- DeltaNet: DeltaNet writes only the key-specific prediction error, improving selective associative storage without an explicit mechanism for globally clearing stale memory.Its update corrects what the memory already predicts rather than simply adding each new value.
- Gated DeltaNet: Gated DeltaNet adds token-dependent scalar decay before the delta correction, enabling memory clearance while partially sacrificing DeltaNet’s strictly selective update.The forgetting operation is state-wide, whereas the corrective write remains key-specific.
- Kimi Delta Attention: Kimi Delta Attention replaces scalar forgetting with channel-wise decay, allowing different transformed-key dimensions to be retained or cleared at different rates.Its active delta edit still uses one scalar β(i) for both removing old content and writing new values.
- Gated DeltaNet-2: Gated DeltaNet-2 decouples channel-wise erase and write gates, independently controlling key-side removal and value-side storage.It preserves Kimi Delta Attention’s channel-wise decay and reduces to earlier gated forms when its gates collapse appropriately.
4 Cross-Layer Routing
This section introduces lightweight cross-layer routing for DeltaNet-style memories using internal write signals, while preserving the host mixer, recurrent update, gating, and output read. Matched variants compare routing write residuals or write values, with hidden-stream injection addressing CLER’s value-space mismatch.
- Routing framework: Cross-layer routing adds a side-channel between recurrent-memory layers without changing the underlying mixer, recurrent update, or gating.Softmax layers can carry the side-channel in hybrid stacks but are not routing-capable.
- Cross-Layer Error Residuals: CLER forwards a lower layer’s write residual into the next routing-capable layer’s value target before its new residual is computed.The residual is scaled by a learned per-layer scalar and can pass through intervening softmax layers.
- Cross-Layer Error Residuals: The value-target injection in Eq. (27) does not improve over matched baselines, motivating routing into the shared hidden stream rather than independently learned value spaces.The reported explanation is a space mismatch between the routed residual and the receiver’s value geometry.
- Hidden-stream routing: CLER-H routes the write error into the hidden stream, whereas CLVR routes the write value through a zero-initialized per-layer projection into the residual stream.The zero initialization starts training exactly at the host baseline; CLER-H and CLVR are parameter-matched at every rank.
- Cross-Layer Value Routing: CLVR routes each linear-memory layer’s internal write value into the aligned residual stream, distinguishing it from softmax value-residual learning and Attention Residuals.It is defined for delta-rule memories and uses a learned, zero-initialized projection read by later layers and the output head.
5 Experimental Setup
The experiments use controlled 350M-parameter decoder-only language-model comparisons to isolate architecture, optimizer, and stack-composition effects. They define matched evaluation procedures for hybrid and pure stacks, learning-rate sensitivity, sequence-length scaling, and cross-layer routing limitations.
- Controlled setup: The main comparisons hold data, tokenizer, sequence length, batch size, precision, and hardware fixed, isolating architecture, optimizer, and stack-composition effects.The models use FineWeb-Edu, the LLaMA2 tokenizer, and architecture-specific layer counts and stack patterns within the 350M-parameter class.
- Stack composition: Hybrid stacks interleave linear-attention and softmax-attention layers, whereas pure stacks use linear-attention mixers throughout.Default DeltaNet, Gated DeltaNet, and Gated DeltaNet-2 hybrids place one softmax-attention layer every three layers, producing a 2:1 linear-to-softmax pattern.
- Metrics: Validation loss is held-out language-model cross-entropy with lower values better, while relative speed is normalized training throughput within the 350M-parameter, 15B-token sweep.Saturation loss is measured before final decay, and final loss at the recorded checkpoint after the full schedule.
- Optimizer and learning rate: The 350M-parameter, 15B-token comparison evaluates AdamW and Muon alongside learning-rate ablations, because preferred rates depend strongly on both optimizer and mixer.DeltaNet and Gated DeltaNet with AdamW prefer a learning rate near 10^-3, while other AdamW/Muon combinations cluster around 3 × 10^-4.
- Routing evaluation: Routing runs compare each routed model with a matched same-host baseline, while routing claims are restricted to DeltaNet and Gated DeltaNet hosts.The refined study includes 350M-parameter runs at 1B and 15B tokens and a larger Gated DeltaNet configuration, without implying a complete larger-scale sweep.
6 Results
Across the reported single-run experiments, Kimi Delta Attention with Muon achieved the lowest validation loss in the 350M-parameter sweep, while pure Gated DeltaNet with AdamW achieved the highest normalized training throughput. Hybrid stacks generally traded speed for loss improvements, and Cross-Layer Value Routing modestly improved matched DeltaNet and Gated DeltaNet baselines.
- 350M-parameter architecture comparison: Kimi Delta Attention with Muon in a hybrid stack achieved the best final validation loss, 2.273, among the reported 350M-parameter, 15B-token runs.DeltaNet with Muon in a hybrid stack reached 2.299, while Gated DeltaNet-2 with Muon in a hybrid stack reached 2.345.
- 350M-parameter architecture comparison: Pure Gated DeltaNet with AdamW was fastest at 100% normalized training throughput but had substantially higher final loss than the best run.Hybrid stacks usually improved loss relative to pure stacks while reducing part of linear attention’s speed advantage.
- Optimizer and learning-rate behavior: Muon improved matched final loss for every architecture family, with preferred learning rates around 3 × 10−4 for Muon and 10−3 for AdamW linear-attention variants.Standard softmax attention matched its Muon counterpart at 3 × 10−4, showing that learning-rate choice affects optimizer and architecture comparisons.
- Sequence-length scaling: At 32k tokens, iteration time was 3.37 seconds for softmax attention, 1.56 seconds for a Gated DeltaNet hybrid, and 0.96 seconds for pure Gated DeltaNet.From 4k to 32k tokens, growth factors were approximately 2.9×, 1.7×, and 1.1× respectively.
- Larger DeltaNet runs: At 1.3B parameters and 40B tokens, pure DeltaNet reached 2.063 versus 2.112 for pure CLER-DeltaNet at the same learning rate.At 3B parameters and 60B tokens, two better DeltaNet hybrid runs reached 1.955, while the 5 × 10−4 run reached 2.332.
- Cross-layer routing: CLVR lowered final validation loss relative to matched hosts, by −0.0103 for Gated DeltaNet and −0.0119 for DeltaNet at 350M parameters and 1B tokens.The gains remained negative but shrank to −0.0059 and −0.0016 at 350M/15B, and −0.0019 at 1.3B/40B; CLER’s value-space routing did not improve over baseline.
7 Discussion: Validation Loss, Throughput, and Architectural Trade-offs
The discussion frames the architectures as a multi-objective frontier rather than a single ranking, balancing validation loss, throughput, scaling, and implementation complexity. Routing results are mixed, while evidence limitations constrain broad conclusions.
- Architectural trade-offs: The DeltaNet design choices define a multi-objective frontier involving validation loss, training throughput, sequence-length scaling, and implementation complexity.The comparison separates memory reading, forgetting or erasing, and writing new value information.
- Architectural trade-offs: No single setting dominates every objective: Kimi Delta Attention with Muon achieves the best 350M-parameter, 15B-token validation loss but is among the slower entries.Pure recurrent stacks provide the strongest long-context iteration-time scaling, while hybrid stacks recover some validation-loss gap at a throughput cost.
- Evidence limitations: The comparisons are audited single-run results without seed-averaged rankings or standard deviations, with uneven hyperparameter coverage across variants.Downstream evaluation is limited to HellaSwag, PIQA, and WinoGrande.
- Cross-layer routing: CLER activates delta-rule write-error routing, but matched comparisons show no reliable validation-loss improvement despite nonzero residuals and trainable receiver-side coefficients.The discussion attributes the flat result to a basis mismatch between the routed delta-rule residual and the receiver representation.
8 Future Work
Future work should test whether CLVR’s small gains generalize across additional architectures, scales, training budgets, and tasks. It should also directly measure routing influence and extend evaluation from training throughput to inference and long-context quality.
- 8 Future Work: CLVR’s small gains should be tested on Kimi Delta Attention and Gated DeltaNet-2, especially because Gated DeltaNet-2 separates write values from erase and write residuals.A matched comparison would test whether routing the write value generalizes beyond DeltaNet and Gated DeltaNet.
- 8 Future Work: Additional repeated runs at larger token budgets should determine whether CLVR’s diminishing return becomes neutral or remains a small positive floor.This question arises from the diminishing return observed with longer training and in the larger Gated DeltaNet row.
- 8 Future Work: Future routing runs should log the projected routed signal, receiver residual streams before and after injection, and learned projection norms at matched checkpoints.These measurements could distinguish insufficient magnitude, an incorrect basis, and alignment that is redundant with the receiver’s local update.
- 8 Future Work: CLVR should be evaluated on synthetic key-value retrieval, passkey-style tasks, long-context QA, many-distractor in-context learning, and other associative-recall settings.These tasks may provide a clearer role for cross-layer recovery than average next-token prediction.
- 8 Future Work: The throughput study should extend from training iteration time to measured inference throughput, decoding memory footprint, and quality at long context length.This would add empirical inference-speed and long-context measurements to the reported training-throughput study.
9 Conclusion
The comparison reveals a multi-objective trade-off among validation loss, throughput, and sequence-length scaling rather than a single architecture ranking. Cross-layer routing succeeds only after replacing the natural error-residual formulation with routing into the aligned hidden stream.
- Overall comparison: The architectures exhibit a multi-objective trade-off among validation loss, throughput, and sequence-length scaling rather than a single ranking.The study compares softmax attention, DeltaNet, Gated DeltaNet, Kimi Delta Attention, and Gated DeltaNet-2 in a common recurrent-memory notation.
- Empirical trade-offs: Kimi Delta Attention with Muon in a hybrid stack reaches the best validation loss, while pure Gated DeltaNet with AdamW achieves the fastest normalized training speed.Hybrid stacks often improve loss, whereas pure recurrent stacks preserve the strongest long-context iteration-time scaling.
- Cross-layer routing: Cross-Layer Error Residuals, which forward a lower layer’s delta-rule write error into the next layer’s value target, do not improve over matched baselines.The study diagnoses this natural formulation as a space mismatch and instead routes the signal into the aligned hidden stream through a zero-initialized projection.
Appendix · A Supplementary 1B-Token Baselines
Table 8 reports an earlier 350M-parameter baseline sweep trained for approximately 1B tokens under the same controlled setup used in the routing experiments. The appendix provides supplementary context for the paper’s architecture comparisons.
- A Supplementary 1B-Token Baselines: Table 8 reports an earlier 350M-parameter baseline sweep trained for approximately 1B tokens.These baselines supplement the main 15B-token sweep.
- A Supplementary 1B-Token Baselines: The sweep uses the FineWeb-Edu/LLaMA2 setup.This preserves the stated data and model configuration.
- A Supplementary 1B-Token Baselines: Experiments use sequence length 4096.The sequence length matches the controlled routing experiments.
- A Supplementary 1B-Token Baselines: The experiments run on a one-node GH200 configuration.This is the hardware setting specified for the controlled comparison.
- Appendix: The main text emphasizes the 15B-token architecture sweep and larger 1.3B- and 3B-parameter DeltaNet runs.These runs provide the clearest validation-loss/throughput comparison in the paper.
- A Supplementary 1B-Token Baselines: The 1B-token baselines use the same setup as the controlled routing experiments.This makes the appendix sweep a supplementary reference for those experiments.
B Result Inclusion Criteria
The paper applies an auditability-based inclusion standard to distinguish main empirical results, supplementary context, and exploratory checks. Main results are comparable within defined experimental groups, while single-run gaps without corroboration remain suggestive.
- Inclusion standard: Main quantitative claims identify the model family, stack pattern, scale, token budget, optimizer, learning rate, metric definition, and checkpoint semantics.This is the stated standard for including results in the main empirical analysis.
- Statistical scope: Single-run results have no standard deviations, so small validation-loss gaps are treated as suggestive unless matched settings or downstream checks show the same direction.The standard is explicitly an auditability standard rather than a variance estimate.
- Main empirical results: The main empirical set includes the 350M-parameter, 15B-token architecture sweep, learning-rate ablation, sequence-length timing, larger DeltaNet runs, downstream evaluations, CLER ablations, and cross-layer routing comparisons.These results specify training scale and metric semantics and are comparable within their experimental groups.
- Supplementary context: The 350M-parameter, 1B-token baselines provide early baseline and optimizer context but do not determine the final validation-loss/throughput frontier.Their shorter training budget separates them from the main architecture sweep.
- Cross-layer routing results: CLER-H, CLVR, and Attention Residuals are main-text routing comparisons with matched baselines and final validation losses, whereas other routing controls are only indicative checks.The remaining controls include surprise-gated value routing, error-plus-value concatenation, output injection, and low-rank projection variants.