Switching Linear Attention

1 Introduction

• Softmax attention’s key-value cache grows linearly with sequence length, whereas linear attention’s single linear retrieval mapping limits context-dependent associative recall.Linear attention uses a recurrent state of size O(D2), where D is the embedding dimension.Source: §1 · p. 1
• SwiLA derives a mixture-of-linear-regressions recurrence from test-time regression using online expectation-maximization, while keeping a fixed-size recurrent state independent of sequence length.Each output dimension selects among components based on the input and prediction errors of the current state.Source: §1 · p. 2
• Across associative recall, in-context language learning, and language modeling benchmarks, SwiLA narrows the gap to softmax attention and surpasses it in several settings.The introduction characterizes its performance as strong across these settings.Source: §1 · p. 2

2 Background on Test-Time Regression

• In test-time regression, a sequence layer learns a regressor from input tokens in context and applies it to each query to retrieve an output.The regressor’s function class and optimization algorithm, together with the keys and values, determine the mapping.Source: §2 · p. 2
• The TTR framework describes sequence layers through their function class and optimizer, distinguishing linear attention, DeltaNet, and softmax attention by their regression models and fitting procedures.The regressor is learned online at inference time rather than during training.Source: §2 · p. 2

2.1 Limitations motivating Switching Linear Attention

• Linear attention compresses history into a fixed-size state and retrieves through one linear query mapping, while softmax attention uses a nonlinear kernel smoother but stores all KV pairs.The softmax KV cache grows with sequence length.Source: §2.1 · pp. 2–3
• SwiLA replaces the single linear regressor with a mixture, yielding a piecewise linear map while retaining a fixed-size recurrent state.Figure 1 illustrates switching regression and compares softmax-approximation test R2 against state size; its caption reports that SwiLA approximates softmax attention better than DeltaNet at every budget.Source: §2.1 · p. 3, Figure 1 · p. 3

3 Switching Linear Attention

• SwiLA combines a mixture of linear regressors with online expectation-maximization to derive its recurrent update and query retrieval.The regression model and optimizer are the two choices that specify the layer under the TTR framework.Source: §3 · p. 3, §3.2 · p. 4

3.1 Regression Model: Mixture of Linear Regressors

• Each output dimension selects among mixture components, giving an effective JD mixtures with weights that scale as JD2.Component membership is defined separately for each output dimension; each component has its own regression weights.Source: §3.1 · p. 3

3.2 Optimization Algorithm: Online EM

• Online EM updates the regression weights as each new key–value pair arrives, ascending the marginal log-likelihood using prediction errors and posterior responsibilities.After the state update, the query retrieves an output from the mixture of regressors.Source: §3.2 · p. 4

3.3 Regression Model Extension: Temporal SwiLA

• Temporal SwiLA uses a Markovian prior and a running component distribution mixed with a new input-dependent distribution to make mixture assignments persist across steps.This tractable approximation avoids the per-step marginalization and normalizing-constant computations of standard HMM filtering.Source: §3.3 · p. 4

3.4 Regression Model Extension: Gated SwiLA

• Gated SwiLA adds input-dependent state decay through Gaussian priors on mixture weights, with two formulations differing in whether gating precedes responsibility computation.The paper uses the formulation that applies the gate before computing prediction errors for Gated SwiLA.Source: §3.4 · p. 5

3.5 Improving the Optimization Algorithm: Load Balancing

• Without regularization, SwiLA can suffer expert collapse; its load-balancing loss encourages balanced expert use while retaining per-token specialization.The loss maximizes mutual information between tokens and expert assignments by encouraging high marginal entropy and low conditional entropy.Source: §3.5 · p. 5
• The method also injects learned input-dependent noise into expert logits during training.Noise is applied to key-side and query-side expert logits, and to gate logits in Temporal SwiLA.Source: §3.5 · p. 5

4 Related Work

• DeltaNet and SwiLA are used as linear-attention baselines, alongside Mixture-of-Memories and Sparse State Expansion.The paper identifies these layers as closely related methods.Source: §4 · p. 5

4.1 DeltaNet

• DeltaNet updates fast weights with a recurrence linear in the previous weights and retrieves through one linear query mapping, whereas SwiLA uses a nonlinear recurrence and switches among regressors.In TTR, the two layers correspond to different regressor function classes.Source: §4.1 · p. 6
• Both DeltaNet and SwiLA use SGD in the TTR framing, but optimize different objectives; DeltaNet’s chunkwise parallel algorithm could be applied within SwiLA’s Newton iterations.The paper presents this parallelization as a possibility, while current experiments use a sequential recurrent kernel.Source: §4.1 · p. 6

4.2 Mixture-of-Memories

• MoM and SSE route among memory states based on input, while SwiLA’s posterior-responsibility router also depends on the prior state, making its recurrence nonlinear.The cited passage links this distinction to nonlinear versus linear recurrences.Source: §4.2 · p. 6

5.1 Approximating Softmax Attention

• At every recurrent state budget, the best SwiLA configuration achieves higher test R2 than DeltaNet; mixtures and heads play complementary roles as capacity grows.At B=2, one-head H1J2 SwiLA outperforms two-head H2 DeltaNet; at B=4, H2J2 is an example of a balanced configuration.Source: §5.1 · p. 7

5.2 Synthetic Benchmarks

• On CDAR, Temporal SwiLA achieves near-perfect accuracy at the smallest recurrent state size; on RegBench, Temporal SwiLA and DN-MoM beat Transformer with 1,000 training PFAs, but Transformer leads with more data.Figure 2A covers four contexts and 32 key-value pairs per context; Figure 2B plots accuracy against training-set size.Source: Figure 2 · p. 8, §5.2 · p. 8

5.4 Throughput Comparison

• 4–5× higher end-to-end inference throughput than Transformer, but SwiLA is approximately 2× slower training and 1.25–1.5× slower inference than GDN, with larger memory use.The authors attribute the throughput advantage over Transformer to avoiding a KV cache and accommodating larger batch sizes.Source: §5.4 · p. 9

5.3 Pretraining Language Models

• Across commonsense and recall tasks, SwiLA variants outperform recurrent baselines; gated variants slightly surpass Transformer++ on commonsense, while Hybrid Temporal SwiLA approaches FoX on recall.FoX achieves the best average on recall-intensive tasks.Source: §5.3 · p. 9, Table 1 · p. 9, Table 2 · p. 10

5.5 Ablation Study

• Temporal persistence improves perplexity; load-balancing loss is critical for vanilla SwiLA but negligible for Temporal SwiLA, while increasing experts from 4 to 16 slightly degrades performance under a fixed state budget.The authors link the expert-count tradeoff to smaller query, key, and value projections per expert.Source: §5.5 · pp. 9–10

6 Conclusion

• SwiLA retains a sequence-length-independent recurrent state while representing a nonlinear query-to-output map, and performs strongly across recall, in-context learning, and language modeling benchmarks.It narrows the gap to softmax attention and surpasses it in several settings; the authors identify distillation into recurrent layers as a possible direction.Source: §6 · p. 10
• SwiLA’s fixed state cannot grow with context length, and its nonlinear recurrence currently requires sequential pretraining; parallelization needs further kernel optimization to yield speed-ups.The authors note that softmax attention also struggles with attention rot in long contexts.Source: §6 · p. 10
• The conclusion presents TTR as a way to design sequence layers by selecting regression models and update algorithms, with other function classes and softmax-to-recurrent distillation as future possibilities.The authors suggest local linear regression as another function class to explore.Source: §6 · p. 10

A.1 Design Space of Test-Time Regression

• In the TTR design space, linear attention is linear regression with a whitened design-matrix approximation, while softmax attention is a nonparametric local constant regressor.Adding decaying regression weights yields gated linear attention and state-space models; online SGD yields DeltaNet and related variants.Source: §A.1 · p. 19

A.2 Mixtures of Linear Regressions

• The mixture model assigns each observation a latent component, estimates component responsibilities, and alternates responsibility computation with parameter maximization through EM.The regression setup assumes a mixture of J linear Gaussian models.Source: §A.2 · p. 19

B.1 Derivation of SwiLA Recurrence

• SwiLA’s recurrence is derived in Appendix B.1 as a detailed development of equation (3.3).The supplied passage identifies the derivation’s scope but does not state its intermediate steps.Source: §B.1 · p. 20

B.2 Details of the SwiLA Layer

• The implementation projects inputs into queries, keys, and values, with optional SiLU nonlinearities; language-model pretraining applies SiLU to all three.Learning rates and mixture-membership priors also come from input projections, with softmax producing prior probabilities; low-rank projections reduce parameter count in language-model pretraining.Source: §B.2 · p. 20

B.3 Temporal SwiLA Details

• Temporal mixture priors combine the previous filtered posterior with a new input-dependent distribution, with learned g controlling persistence.When g is close to 0 the distribution stays similar to the previous timestep; when it is close to 1 the input suggests a new distribution.Source: §B.3 · pp. 20–21

B.4 Gated SwiLA Details

• When J = 1, the two gated formulations are equivalent after reparameterization, and their responsibilities drop out of the updates.The equivalence applies when α_td is positive.Source: §B.4 · p. 21
• Gating granularity determines whether SwiLA applies decay per key dimension, per output dimension, or per head, corresponding respectively to KDA, equation (3.7), and GDN.The general update allows decay to vary across both key and output indices.Source: §B.4 · p. 21

C Extended Related Work

• SwiLA extends TTR with mixture modeling; the authors describe it as the first TTR sequence layer to do so, while Local Linear Attention still requires O(TD) memory.The related-work discussion also distinguishes TTR’s focus on how recurrent updates affect in-context learning from approaches emphasizing pretraining or broader functional relationships.Source: C Extended Related Work · p. 22
• The discussion frames SwiLA’s mixture of regressors as state expansion and contrasts fixed-state designs with logarithmically growing state and software-optimization approaches.The authors also identify optimized software implementations as an orthogonal efficiency approach.Source: C Extended Related Work · p. 22

D Chunkwise Parallel Algorithms

• The paper combines parallel Newton iterations with DeltaNet’s chunkwise algorithm to parallelize SwiLA’s nonlinear recurrence in principle.The experiments instead use an optimized recurrent Triton kernel for sequential evaluation; the authors leave an optimized chunkwise implementation for future work.Source: D Chunkwise Parallel Algorithms · p. 23

D.1 Background on Parallelizing Recurrences

• When the approximate Jacobian has identity-minus-rank-one form, each fixed-point step can use DeltaNet’s chunkwise algorithm.The general parallel-Newton approach linearizes the nonlinear recurrence and evaluates the linearized system in parallel.Source: §D.1 · pp. 23–24

D.2 Linearizing the SwiLA Recurrence

• SwiLA’s per-component, per-output-dimension Jacobian is a rank-one perturbation of the identity, enabling chunkwise solution of each Newton iteration.The quasi-Newton Jacobians ignore cross-component interactions from the responsibility denominators; the cited work guarantees global convergence despite the approximate Jacobian.Source: §D.2 · pp. 24–25
• Quasi-Newton iterations have an approximate backward pass; fixed-point iterations in the backward pass or unnormalized responsibilities are proposed alternatives.Removing responsibility denominators is described as removing the source of cross-component interaction.Source: §D.2 · p. 25

D.3 Chunkwise Parallel Algorithm

• The chunkwise algorithm divides the sequence into chunks, resolves boundary states, then computes outputs within each chunk in parallel under a causal mask.The paper reports that its current parallel Triton kernels were effectively the same speed as the sequential kernel despite convergence in around 10 Newton iterations on synthetic inputs.Source: §D.3 · pp. 25–26, §D.2 · p. 25
• On the synthetic attention-regression task, SwiLA better approximates the attention mapping than DeltaNet at every recurrent-state budget.The figure compares test R2 against recurrent state size and parameter count, with trainable softmax attention shown by a dashed line.Source: Figure S1 · p. 27

E Experimental Details

• The experiments heavily utilize architecture implementations from Yang & Zhang (2024).The passage does not specify which implementations or experimental configurations are involved.Source: E Experimental Details · p. 27

• Models approximate causal softmax self-attention on length-512 sequences with embedding dimension 32, using fixed ground-truth attention projections as the target.The models receive only the input and learn their own Q/K/V projections.Source: §E.1 · p. 27
• Inputs are generated from 32 clusters, with each token formed by normalizing its cluster center plus Gaussian noise.The cluster centers are fixed; query, key, and value projections are random matrices.Source: §E.1 · p. 27
• The comparison includes learned causal multi-head softmax attention, DeltaNet, and SwiLA, with distinct projection and routing configurations.SwiLA uses separate β values across heads, output dimensions, and mixture components, and computes routing logits with SwiGLU projections and softmax.Source: §E.1 · p. 27

E.2 Synthetic Benchmarks

• CDAR extends multi-query associative recall to four contexts, each containing 32 distinct key–value pairs, followed by queries that retrieve the value for a context–key pair.Sequences have length 1024 and contain 128 associations in total.Source: §E.2 · pp. 27–28
• Temporal SwiLA yields the largest improvements on CDAR, and mean test accuracy across learning rates gives the same qualitative conclusions as Figure 2A.Figure S2 reports mean test accuracy versus state size for four contexts with 32 key-value pairs per context.Source: Figure S2 · p. 28, §E.2 · p. 29

E.3 Pretraining Language Models

• Language models are pretrained on the sample-100BT FineWeb-Edu subset with 4,096-token sequences for approximately 15B tokens on eight H100 GPUs.Training runs for 28,610 steps with a global batch size of 128 sequences.Source: §E.3 · p. 29
• The FineWeb-Edu language-model comparison uses 24-layer models with hidden dimension 1,024 and approximately 374M parameters, matching recurrent-state sizes across specified architectures.SwiLA variants use 131,072 recurrent-state scalars per layer, matching DeltaNet, GDN, and KDA.Source: §E.3 · p. 29

E.4 Throughput Comparison

• Training efficiency is profiled for six 1.3B models on one H100 at sequence lengths from 2,048 to 16,384, holding each update workload at 131,072 tokens.Throughput is reported in thousands of tokens per second, with peak allocated CUDA memory also measured.Source: §E.4 · p. 31, Figure S3 · p. 31
• Generation efficiency is measured for the same six models on one H100 using 2,048-token prompts, 128 generated tokens, and batch sizes from 1 to 128.The full generation pass includes prefill and decoding; batch sizes after an out-of-memory failure are omitted.Source: §E.4 · p. 31, Figure S4 · p. 32

E.5 Ablation Study

• Without load balancing, vanilla SwiLA’s key-value-side utilization entropy drops sharply and its dead expert fraction exceeds 40%.The router diagnostics separate query-side and key-value-side responsibilities.Source: §E.5 · p. 33, Figure S5 · p. 33
• Load balancing maintains high utilization entropy and near-zero dead experts, while Temporal SwiLA retains reasonable utilization even without load balancing.The ablation compares mixture-component counts, temporal persistence, and load-balancing loss.Source: §E.5 · p. 33
• Routing entropy stabilizes between 0.4 and 0.6 in all configurations except vanilla SwiLA without load balancing, indicating moderately selective rather than fully hard routing.The study tracks utilization entropy, routing entropy, and dead expert fraction separately for query-side and key-value-side responsibilities.Source: §E.5 · p. 33