Switching Linear Attention

Hyun Dong Lee; Xavier Gonzalez; Nicolas Zucchet; E. Kelly Buchanan; Emily B. Fox; Scott W. Linderman

cs.LGSep 30, 2026

Switching Linear Attention

Authors: Hyun Dong Lee, Xavier Gonzalez, Nicolas Zucchet, E. Kelly Buchanan, Emily B. Fox, Scott W. Linderman

Organizations: Stanford University · Unconventional AI. Work done at Stanford University.

Abstract

Designing expressive sequence layers with efficient inference remains a central challenge in modern machine learning. Standard softmax attention achieves excellent sequence modeling performance through rich nonlinear token interactions, but it requires a key-value cache that grows linearly with sequence length, limiting its scalability. Linear attention enables efficient recurrent computation with a constant memory footprint, yet its reduced expressivity often yields inferior modeling performance. We introduce Switching Linear Attention (SwiLA), a novel sequence layer that bridges this gap by enhancing representational capacity while retaining the fixed-size recurrent state of linear attention. We derive the SwiLA recurrence from the test-time regression framework, casting the state update rule as online expectation-maximization in a mixture of linear regressions model. At test time, each output dimension dynamically selects among multiple linear attention components based on the input. Across associative recall, in-context language learning, and language modeling benchmarks, SwiLA shows strong performance and narrows the gap to softmax attention, even surpassing it in several settings.

Figures & tables

Figure 1: Switching Linear Attention. (A) : When the input–output relationship varies with context, piecewise (”switching”) regression is more expressive than linear regression. (B) : SwiLA increases expressivity toward that of softmax attention while retaining, as in linear attention, a fixed-size recurrent state whose memory footprint is independent of sequence length. SwiLA’s state is a mixture of linear regressors. At each step, SwiLA selects which regressor to use based on the current input and previous state. (C) : Test $R^{2}$ vs. state size for SwiLA and DeltaNet on a synthetic softmax attention regression task, with softmax attention for reference. Labels mark the best SwiLA configuration ( H $=$ heads, J $=$ mixtures). At every budget, SwiLA better approximates softmax attention.

Panel of Figure 2 — Figure 2: Synthetic Benchmarks . ( 2A ) Context-Dependent Associative Recall with 4 contexts and 32 key-value pairs per context: test accuracy vs. state size. ( 2B ) RegBench: test accuracy as a function of training set size.

	Perplexity ( $\downarrow$ )		Accuracy ( $\uparrow$ , %)
Model (374M)	Wiki.	LMB.	LMB.	PIQA	Hella.	Wino.	ARC-E	ARC-C	SIQA	BoolQ	Avg.
Transformer++	28.2	37.4	33.5	66.0	32.4	52.6	56.0	24.1	38.4	59.1	45.3
FoX	28.7	50.3	30.8	65.2	32.2	52.5	57.2	23.2	37.9	58.0	44.6
DeltaNet	29.8	46.3	27.9	63.9	31.9	51.6	55.6	21.8	38.3	58.3	43.7
Gated DeltaNet	27.7	35.7	32.0	66.0	33.3	51.8	57.5	24.6	39.0	55.8	45.0
KDA	23.4	26.8	36.4	67.0	34.9	51.5	60.7	27.3	38.9	56.2	46.6

Table 1: Commonsense Reasoning . Best result per task in bold , second best underlined .

Model	SWDE	FDA	SQuAD	TriviaQA	NQ	DROP	Avg.
Transformer++	28.4	29.5	31.5	46.0	15.6	18.9	28.3
FoX	36.4	53.0	30.6	46.4	15.6	20.2	33.7
DeltaNet	7.9	5.4	23.4	41.7	10.5	14.9	17.3
Gated DeltaNet	9.1	5.5	25.2	44.4	12.9	18.0	19.2
KDA	13.9	5.5	25.5	46.8	12.5	17.6	20.3
DN-MoM	5.4	2.1	21.2	37.0	9.1	14.8	14.9

Table 2: Recall-Intensive Tasks . All scores use the contains metric (%,

\uparrow

). Best in bold , second best underlined . All inputs are truncated to 2K tokens.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Figure S1: Approximating Softmax Attention. Test

R^{2}

vs. recurrent state size (left) and parameter count (right) for SwiLA and DeltaNet on the synthetic attention regression task. The dashed line indicates trainable softmax attention. Labels mark the best SwiLA configuration ( H

=

heads, J

=

mixtures) at each budget. At every state budget, SwiLA better approximates the attention mapping.

Figure S2: Context-Depedent Associative Recall. CDAR with 4 contexts and 32 key-value pairs per context: test accuracy vs. state size.

Figure S3: Training throughput and peak memory of 1.3B models with 128K training tokens on a single H100 GPU (i.e., 1M training tokens on 8xH100 GPUs). Following the training throughput experiment of Yang et al. (2025) , we fix the total number of training tokens while varying the sequence length and batch size.

Figure S4: Inference throughput and peak memory of 1.3B models on a single H100 GPU. Following the inference throughput experiment of Gu and Dao (2023) , we fix the prompt length to 2048 and vary the batch size. Throughput and peak memory are measured over the full inference pass, combining prefill and decode. An X marks configurations where the model runs out of memory.

Figure S5: Switching Linear Attention router diagnostics. Utilization entropy, routing entropy, and dead expert fraction over the course of training, shown separately for the query-side responsibilities (top) and key-value-side responsibilities (bottom). Without load balancing, vanilla SwiLA (4

J

) collapses on the KV side, with low utilization entropy and a dead expert fraction near 40%; load balancing and temporal persistence each prevent this collapse.

Explore similar work

Linear Attention Architectures: Mechanisms, Trade-offs, and Cross-Layer Routing
Jul 8, 2026Tommaso Cerruti, Tim Rieder, George Rowlands +2Kimi Delta Attention Linear Attention
1ETH Zurich, D-INFK · 2ETH AI Center, ETH Zurich
Sliding-window beats linear attention
Aug 28, 2026Alexia Jolicoeur-Martineau, Rhea Sanjay Sukthanker, Pashmina Cameron +1Kimi Delta Attention Large Language Model Memory
Microsoft Applied Sciences Group (ASG) · Independent
Dynamic Linear Attention
Jun 9, 2026Xin Wang, Hui Shen, Boyuan Zheng +7Kimi Delta Attention Efficient Long-Context Inference
The Ohio State University · University of Michigan · ByteDance Seed

Jul 8, 2026, cs.LGJ/K move · Enter open · S save

Linear Attention Architectures: Mechanisms, Trade-offs, and Cross-Layer Routing

Tommaso Cerruti, Tim Rieder, George Rowlands, Lingfeng Jin+1

1ETH Zurich, D-INFK · 2ETH AI Center, ETH Zurich

Kimi Delta Attention Linear Attention Gated Deltanet Batch Trade-Off

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

Linear Attention Architectures: Mechanisms, Trade-offs, and Cross-Layer Routing

Sliding-window beats linear attention

Dynamic Linear Attention