Paper deep dive
Switching Linear Attention
Hyun Dong Lee, Xavier Gonzalez, Nicolas Zucchet, E. Kelly Buchanan, Emily B. Fox, Scott W. Linderman
Intelligence
Status: not_run | Model: - | Prompt: - | Confidence: 0%
Entities (0)
Relation Signals (0)
No relation signals yet.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
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.
Tags
Links
- Source: https://arxiv.org/abs/2609.39034v1
- Canonical: https://arxiv.org/abs/2609.39034v1
Trouble viewing inline? Open PDF directly →
Full Text
107,248 characters extracted from source content.
Expand or collapse full text
Switching Linear Attention Hyun Dong Lee* Affiliation: Stanford University Xavier Gonzalez Affiliation: Stanford University† Nicolas Zucchet Affiliation: Stanford University E. Kelly Buchanan Affiliation: Stanford University Emily B. Fox Affiliation: Stanford University Scott W. Linderman Affiliation: 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. †footnotetext: ∗Correspondence to hdlee@cs.stanford.edu.†footnotetext: †Now at Unconventional AI. Work done at Stanford University.†footnotetext: ‡Code is available at switching-linear-attention. 1 Introduction Attention mechanisms (Vaswani et al., 2017) have emerged as the cornerstone of modern sequence modeling, powering state-of-the-art advances in language understanding (Brown et al., 2020), visual recognition (Dosovitskiy et al., 2021), and multimodal vision-language reasoning (Alayrac et al., 2022). The key innovation is dynamic information routing based on input-dependent similarity, enabling models to selectively retrieve relevant context from across a sequence. A useful perspective on this retrieval capability is that softmax attention implicitly fits a nonlinear mapping between inputs and outputs at inference time, a form of test-time regression (TTR) that underlies its expressive in-context learning (ICL) (Garg et al., 2022; Akyürek et al., 2023; von Oswald et al., 2023; Wang et al., 2025). However, this expressivity comes at a high computational cost. The key-value (KV) cache required for autoregressive generation grows linearly with sequence length, which becomes prohibitive for long-context applications. Linear attention (Katharopoulos et al., 2020) addresses this bottleneck by replacing softmax with a kernel feature map, reformulating attention as a linear recurrence with a state of size O(D2)O(D^2), where D is the embedding dimension. However, this compression fundamentally limits associative recall: at each timestep, linear attention retrieves values through a single linear mapping of the query, and cannot capture the diverse, context-dependent retrieval patterns that softmax attention achieves through its nonlinear mapping. Empirically, linear attention consistently underperforms on tasks requiring precise memory retrieval (Arora et al., 2024a; Arora et al., 2024b). To bridge this gap, we introduce Switching Linear Attention (SwiLA), a novel sequence layer that leverages the TTR framework of Wang et al. (2025), extending it to a mixture of linear regressors‡. SwiLA divides its recurrent state into J mixture components, each maintaining its own linear regressor of size O(D2)O(D^2) that specializes in different retrieval patterns. At each timestep, each of the D output dimensions selects a mixture component based on both a learned gating prior and the prediction errors of the current state, yielding a number of mixtures that scales exponentially as O(JD)O(J^D), despite having a state size that grows only polynomially as O(JD2)O(JD^2). We derive the update rule from online expectation maximization (Cappé and Moulines, 2009): the E-step computes posterior responsibilities based on prediction errors, and the M-step updates each component’s regression weights. During in-context learning, SwiLA learns an expressive piecewise linear mapping from inputs to outputs, while retaining, as in linear attention, a fixed-size recurrent state whose memory footprint is independent of sequence length. With experiments on associative recall, in-context language learning, and language modeling, we find that SwiLA performs strongly, narrowing the gap to softmax attention and even surpassing it in several settings. 2 Background on Test-Time Regression A central capability of effective sequence models is in-context associative recall: given a query, retrieve relevant information from previously observed tokens. Wang et al. (2025) formalize in-context associative recall as a two-step process of memorization and retrieval. They show that the memorization step can be cast as an online regression problem, which they call test-time regression (TTR). Memorization as regression. Given a sequence of KV pairs (k1,v1),…,(kT,vT)(k_1,v_1),…,(k_T,v_T) with kt∈ℝDk_t ^D and vt∈ℝDv_t ^D, an associative memory is a function mt:ℝD→ℝDm_t:R^D ^D such that mt(ki)≈vim_t(k_i)≈ v_i for i=1,…,ti=1,…,t. Finding such a map reduces to solving a regression problem: mt≈argmin∑i=1tm∈ℳ‖vi−m(ki)‖2,m_t≈ _m _i=1^t\|v_i-m(k_i)\|^2, (2.1) where ℳM is a chosen function class. Retrieval as function evaluation. Once the regressor mtm_t has been fit, memory retrieval is simply function application: given a query qtq_t, the output is ot=mt(qt).o_t=m_t(q_t). (2.2) A sequence layer that solves eq. 2.1 and applies eq. 2.2 at each timestep performs associative recall in its forward pass. This framework is called test-time regression (TTR): the regressor mtm_t is learned in-context from the input tokens and the structure of the recurrence. Commonly, the query qtq_t, key ktk_t, and value vtv_t are all projections of some input xtx_t. Design space. In TTR, a sequence layer is specified by the choice of the function class ℳM and the optimization algorithm for eq. 2.1. These choices, along with the particular keys and values shown at test-time, ultimately determine the regressor mtm_t. In the parametric case, mtm_t corresponds to the fixed-size weights of a function approximator, updated online as each new pair (kt,vt)(k_t,v_t) arrives. In the nonparametric case, we recover kernel regression. Importantly, mtm_t is not learned during training; it is learned in-context from the input tokens at inference time, by the very structure of the recurrence. Existing sequence layers arise as special cases. Linear attention (Katharopoulos et al., 2020) uses a linear model for mtm_t, fit with a whitened design matrix approximation. DeltaNet (Schlag et al., 2021; Yang et al., 2024b) also employs a linear model but uses SGD as its optimization algorithm. Softmax attention fits a nonparametric kernel smoother. We refer the reader to Wang et al. (2025) for a comprehensive taxonomy and to Appendix A.1 for further details. 2.1 Limitations motivating Switching Linear Attention TTR reveals a fundamental tension. Linear attention and its variants compress all history into a single fixed-size state, bounding the amount of context that can be stored. At each timestep, retrieval is a single linear mapping of the query, which cannot capture the diverse, context-dependent retrieval patterns that softmax attention achieves through its nonlinear mapping. Softmax attention avoids both limitations by storing all KV pairs and retrieving via a nonlinear kernel smoother, but at the cost of a KV cache that grows with sequence length. SwiLA defines a piecewise linear map that can approximate arbitrary nonlinear functions (fig. 1A) by replacing the single linear regressor with a mixture of linear regressors (fig. 1B). Thus, SwiLA addresses the retrieval limitation of linear attention while retaining a fixed-size recurrent state whose memory footprint is independent of sequence length. 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 R2R^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. 3 Switching Linear Attention We introduce Switching Linear Attention (SwiLA), a sequence layer based on mixtures of linear regressions. SwiLA retains the fixed-size recurrent state of linear attention, while improving its retrieval expressivity through a nonlinear mapping from query to output in eq. 2.2, as in softmax attention. SwiLA is a form of TTR with, • Regression model: mixture of linear regressors (Pearson, 1894; De Veaux, 1989) • Optimization algorithm: online expectation–maximization (EM) (Dempster et al., 1977; Cappé and Moulines, 2009). For the rest of the paper, we consider a sequence of keys ktk_t, values vtv_t, and queries qtq_t in ℝDR^D, projected from the inputs xt∈ℝdinx_t ^d_in. Further technical details are given in Appendix B. 3.1 Regression Model: Mixture of Linear Regressors Mixture modeling is a foundational approach for modeling functions that vary based on group membership (Bishop, 2006). Consider a dataset (kt,vt)t=1T\(k_t,v_t)\_t=1^T consisting of key–value pairs with kt∈ℝDk_t ^D and vt∈ℝDv_t ^D, assumed to be generated from a mixture of J linear Gaussian regression models. To improve expressivity in high-dimensional settings, we allow group membership to vary across the output dimension d. Specifically, let ztd∈1,…,Jz_td∈\1,…,J\ denote the group membership of data point t and output dimension d. Then, given a key ktk_t, the value vtdv_td is modeled as p(ztd|kt)=Cat(ztd|πjd(kt)),p(vtd|ztd=j,kt)=(vtd|wjd⊤kt, 1). p(z_td \ | \ k_t)=Cat (z_td \ | \ _jd(k_t) ), p(v_td \ | \ z_td=j,k_t)=N\! (v_td \ | \ w_jd k_t,\;1 ). (3.1) Here, πjd(kt) _jd(k_t) denotes the prior probabilities of the mixture components. The regression weights, wjd∈ℝDw_jd ^D, are specific to mixture component j and output dimension d. We denote W∈ℝJ×D×DW ^J× D× D as the collection of all weights. Since each output dimension independently selects its own mixture component, the effective number of mixtures grows exponentially as JDJ^D, whereas the number of weights scales polynomially as JD2JD^2. Under the probabilistic model in eq. 3.1, the marginal log-likelihood for timestep t and output dimension d is ℓtd=logp(vtd|kt)=log∑j=1Jπjd(kt)exp(−12(vtd−wjd⊤kt)2)+c. _td= p(v_td \ | \ k_t)= _j=1^J _jd(k_t)\, \! (- 12 (v_td-w_jd k_t )^2 )+c. (3.2) Complete modeling details are given in Section B.2. 3.2 Optimization Algorithm: Online EM The standard approach for fitting mixtures of linear regressors (3.1) is the EM algorithm. In the TTR setting, the weights W are estimated online: at each timestep t, the current estimates Wt−1W_t-1 are updated as the new pair (kt,vt)(k_t,v_t) arrives. Our choice of online EM as the optimizer is what gives rise to the SwiLA recurrence. It amounts to a one-step stochastic gradient update, where we ascend on the marginal log-likelihood (3.2) with respect to the ”fast-weights” W. In particular, the SwiLA recurrence for all mixture components j and output dimensions d is wtjd w_tjd =wt−1,jd+βtjd∇wt−1,jdℓtd =w_t-1,jd+ _tjd _w_t-1,jd _td =wt−1,jd+βtjdrtjdδtjdkt. =w_t-1,jd+ _tjd\,r_tjd\, _tjd\,k_t. (3.3) Here, βtjd _tjd is a learning rate and δtjd≔vtd−wt−1,jd⊤kt _tjd v_td-w_t-1,jd k_t is a prediction error, just as in the delta rule (Schlag et al., 2021; Yang et al., 2024b). The responsibilities rtjd≔πjd(kt)exp(−12δtjd2)∑j′πj′d(kt)exp(−12δtj′d2), r_tjd _jd(k_t)\, \! (- 12\, _tjd^2 ) _j _j d(k_t)\, \! (- 12\, _tj d^2 ), (3.4) are the posterior probabilities that the d-th dimension of vtv_t was generated by mixture component j. Once the state wtjdw_tjd has been updated, the query qtq_t then retrieves the output otdo_td according to the mixture of linear regressors, as otd=∑j=1Jπjd(qt)wtjd⊤qt, o_td= _j=1^J _jd(q_t)\,w_tjd q_t, (3.5) where πjd(qt) _jd(q_t) is also a learned projection from the input xtx_t. In summary, we derive the SwiLA recursion (3.3) and retrieval (3.5) from the TTR framework by choosing our model as a mixture of linear regressors and our optimizer as online EM. We build on this probabilistic perspective in the next two sections: Section 3.3 adds temporal persistence to the prior over mixture assignments, and Section 3.4 derives input-dependent gating from Gaussian priors on the mixture weights. 3.3 Regression Model Extension: Temporal SwiLA A natural benefit of SwiLA is that its mixture components can accommodate different ”contexts.” In many settings, context persists across consecutive inputs. For example, in an essay comparing the musician Harry Styles with the magician Harry Potter, the discussion of each person likely spans a full paragraph, and the value associated with ‘Harry’ should depend on that context. Mathematically, this persistence suggests that if input t has a high probability of belonging to mixture component j, then so should input t+1t+1. However, SwiLA as formulated in Section 3.1 does not have any persistence in the mixture assignments. The mixture assignment of each time step is memoryless, starting over with another draw from its prior π. To address this limitation, we extend SwiLA with a Markovian prior that introduces dependencies between ztdz_td and zt+1,dz_t+1,d. The standard probabilistic approach for such temporal persistence is a hidden Markov model (HMM), but such an approach would require marginalization and computation of normalizing constants at every time step. Instead, we use a computationally tractable approximation to HMM filtering. We maintain a running probability distribution over mixture components for each output dimension d, and at each time step, take a convex combination of this distribution with a new input-dependent distribution generated from the current input. With this extension, the mixture assignments in SwiLA can persist across consecutive time steps. See Section B.3 for full details on this temporal recurrence. 3.4 Regression Model Extension: Gated SwiLA Input-dependent state decay, or gating, has driven strong empirical gains in linear attention, as in GDN (Yang et al., 2025) and KDA (Team et al., 2025). We present two formulations of Gated SwiLA, which differ in whether the gate is applied after or before computing the responsibilities. See Section B.4 for full details. Under the TTR framework, state decay is equivalent to weight regularization of the fast weights (Wang et al., 2025). Accordingly, we place Gaussian priors on SwiLA’s mixture weights, wjd∼(0,λtjd−1I)w_jd (0, _tjd^-1I), which introduces a regularization term to eq. (3.2): ℓtd=log∑j=1Jπjd(kt)exp(−12(vtd−wjd⊤kt)2)−∑j=1Jλtjd2∥wjd∥22+c. _td= _j=1^J _jd(k_t)\, \! (- 12 (v_td-w_jd k_t )^2 )- _j=1^J _tjd2 w_jd _2^2+c. (3.6) A one-step stochastic gradient ascent on eq. (3.6) yields the recurrence wtjd=αtjdwt−1,jd+βtjdrtjdδtjdkt,αtjd≔1−βtjdλtjd,w_tjd= _tjd\,w_t-1,jd+ _tjd\,r_tjd\, _tjd\,k_t, _tjd 1- _tjd _tjd, (3.7) where we parameterize αtjd∈[0,1] _tjd∈[0,1] as a learned, input-dependent gate that is applied after the responsibilities are computed. Alternatively, following the implementation of the recurrent form of GDN (Yang et al., 2025; Yang and Zhang, 2024), which applies the gate before computing the prediction error, we obtain the recurrence wtjd=αtjdwt−1,jd+βtjdrtjd−δtjd−kt,w_tjd= _tjd\,w_t-1,jd+ _tjd\,r^-_tjd\,δ^-_tjd\,k_t, (3.8) where δtjd−≔vtd−αtjdwt−1,jd⊤ktδ^-_tjd v_td- _tjd\,w_t-1,jd k_t are the prediction errors at the decayed state and rtjd−r^-_tjd are the responsibilities computed from these errors. The two formulations coincide for J=1J=1 up to a reparameterization, but differ for J>1J>1. In practice, we use eq. (3.8) for Gated SwiLA. 3.5 Improving the Optimization Algorithm: Load Balancing Without regularization, SwiLA is susceptible to expert collapse, where a few mixture components capture all responsibility mass, and the remaining experts go unused. We employ two complementary strategies to prevent this. First, we add an auxiliary loss that adapts load balancing loss (Shazeer et al., 2017; Fedus et al., 2022). We maximize the mutual information I(token,expert)=Hmarg−HcondI(token;\,expert)=H_marg-H_cond between tokens and expert assignments, where HmargH_marg is the entropy of the time-averaged responsibilities and HcondH_cond is the mean per-token entropy. Maximizing HmargH_marg encourages balanced utilization across experts, while minimizing HcondH_cond encourages per-token specialization. The auxiliary loss is, ℒMI=−λmargHmarg+λcondHcond.L_MI=- _margH_marg+ _cond\,H_cond. (3.9) It is computed for both the key-side and query-side responsibilities and averaged. Second, following Shazeer et al. (2017), we inject learned input-dependent noise into the expert logits during training: π~jd(xt)=exp(θjd⊤xt+ϵtjd⋅softplus(θnoise,jd⊤xt))∑j′exp(θj′d⊤xt+ϵtj′d⋅softplus(θnoise,j′d⊤xt)),ϵtjd∼(0,1). π_jd(x_t)= \! ( _jd x_t+ _tjd·softplus( _noise,jd \,x_t) ) _j \! ( _j d x_t+ _tj d·softplus( _noise,j d \,x_t) ), _tjd (0,1). (3.10) Noise is injected into both the key-side and query-side expert logits, as well as into the gate logits in Temporal SwiLA (Section 3.3), each with its own learned projection. 4 Related Work We discuss two closely related linear attention sequence layers to SwiLA: DeltaNet (Schlag et al., 2021; Yang et al., 2024b) and Mixture-of-Memories (MoM) (Du et al., 2026) / Sparse State Expansion (SSE) (Pan et al., 2026). The similarity of these sequence layers to SwiLA motivates our use of them as our linear attention baselines in Section 5. We include a more expansive discussion of related works in Appendix C. 4.1 DeltaNet DeltaNet (Schlag et al., 2021) is a sequence modeling architecture that uses the delta rule for its linear attention sequence mixer: its state Wt∈ℝD×DW_t ^D× D updates according to eq. 4.1, and then WtW_t provides a linear mapping from query to output, as shown in eq. 4.2: Wt W_t =Wt−1+βt(vt−Wt−1kt)kt⊤ =W_t-1+ _t (v_t-W_t-1k_t )k_t (4.1) ot o_t =Wtqt =W_tq_t (4.2) On the one hand, there are clear differences between SwiLA and DeltaNet. The SwiLA recurrence (3.3) is nonlinear while the delta rule (4.1) is linear in its fast-weights Wt−1W_t-1. Furthermore, for retrieval, DeltaNet provides a single linear mapping from query to output (4.2), while SwiLA switches adaptively between mixtures of linear regressors (3.5). In TTR, DeltaNet and SwiLA correspond to different function classes for the regressor. On the other hand, in TTR, DeltaNet and SwiLA both use SGD as the optimizer. The delta rule optimizes ‖vt−Wtkt‖2\|v_t-W_tk_t\|^2, while SwiLA maximizes the marginal log-likelihood (3.2). This similarity is why the prediction error appears in both recurrences. Moreover, it means that DeltaNet’s chunkwise parallel algorithm (Yang et al., 2024b) could be applied within each Newton iteration to parallelize SwiLA across the sequence length (Appendix D). Chunkwise parallel algorithm SwiLA’s nonlinear recurrence in its fast weights WtW_t can theoretically be parallelized using parallel Newton iterations (Lim et al., 2024; Gonzalez et al., 2024; Gonzalez, 2026; Danieli et al., 2026). Parallel Newton iterations work by iteratively linearizing the recurrence, and then using parallel compute to evaluate the linearized system, and they are guaranteed to converge (Gonzalez, 2026). Previous implementations of these iterations relied on parallel associative scans (Blelloch, 1990), which incur a prohibitive O(TD2)O(TD^2) memory cost by materializing all hidden states (Yang, 2024). To overcome this, we observe that each Newton iteration applied to SwiLA shares the same rank-one update structure as the delta rule recurrence, enabling a novel chunkwise parallel Newton algorithm that avoids the parallel scan’s memory bottleneck. Although our current experiments rely on a fast recurrent (sequential) Triton kernel for SwiLA, the systems-level optimization of this chunkwise parallel approach presents an interesting direction for future work. We discuss parallelization further in Appendix D. 4.2 Mixture-of-Memories The most similar method to SwiLA is the Mixture-of-Memories (MoM) layer introduced by Du et al. (2026). MoM also splits the hidden state into J different memory states, and uses a learnable router to select which memories to activate based on the input. The Sparse State Expansion (SSE) layer of Pan et al. (2026) also uses a mixture of memories, albeit with sparse updates. Since the routers of these layers do not depend on the previous state, their recurrences are linear in their hidden states WtW_t. In contrast, we derive our SwiLA update from the TTR framework. Consequently, our “router” is based on the posterior responsibilities, which are functions not only of the input but also the prior state, rendering our recurrence nonlinear in WtW_t. Increasingly, theoretical (Siems et al., 2026; Merrill et al., 2026) and empirical (Farsang and Grosu, 2025; Zattra et al., 2025; Mishra et al., 2026) results demonstrate that nonlinear recurrences have superior expressivity to linear recurrences. 5 Experimental Results We evaluate SwiLA in three settings. Section 5.1 highlights SwiLA’s ability to approximate the input–output mapping of a single softmax attention head. Section 5.2 evaluates SwiLA on two synthetic benchmarks: context-dependent associative recall and in-context language learning. Section 5.3 scales to language models pretrained on FineWeb-Edu, with downstream evaluation on commonsense reasoning and recall-intensive tasks. Complementing these evaluations, Section 5.4 profiles training and inference efficiency, and Section 5.5 ablates SwiLA’s key design choices. For our language modeling evaluation, we compare SwiLA against a suite of ungated and gated baselines: DeltaNet (Schlag et al., 2021; Yang et al., 2024b), Gated DeltaNet (GDN) (Yang et al., 2025), KDA (Team et al., 2025), and MoM (Du et al., 2026) with DeltaNet-style (DN-MoM) and GDN-style (GDN-MoM) updates, as well as Transformer++ (Touvron et al., 2023) and FoX (Lin et al., 2025) for softmax attention variants. We further compare Hybrid GDN and Hybrid Temporal SwiLA, which interleave three recurrent layers with one softmax attention layer following Team et al. (2025). The recurrent state size is matched across all recurrent models in every experiment. See Appendix E for full experimental details. 5.1 Approximating Softmax Attention We test whether SwiLA can better approximate the input–output mapping of a single softmax attention head than DeltaNet, given the same recurrent state size. Setup. We construct a synthetic regression task in which the target function is causal softmax attention. In this experiment, we set D=32D=32. Inputs xt∈ℝDx_t ^D are drawn from 32 randomly sampled clusters in ℝDR^D. Using fixed, randomly sampled projection matrices θQ,θK,θV _Q, _K, _V, we compute queries qt=θQxtq_t= _Qx_t, keys kt=θKxtk_t= _Kx_t, and values vt=θVxtv_t= _Vx_t. The regression targets are the outputs of causal softmax attention with temperature τ=D/2τ= D/2: ot=∑i=1texp(qt⊤ki/τ)∑j=1texp(qt⊤kj/τ)vi.o_t= _i=1^t (\,q_t k_i/τ) _j=1^t (q_t k_j/τ)\,v_i. (5.1) We generate 65,536 training sequences of length 512 and report test R2R^2 on a held-out set of 2,048 sequences. We compare DeltaNet and SwiLA at matched state budgets B=H×JB=H× J, where H is the number of heads and J the number of mixtures per head (J=1J=1 recovers DeltaNet). All configurations use a fixed head dimension of D=32D=32, giving a total recurrent state size of B×D2B× D^2. For different factorizations (H,J)(H,J), we expand the key and value dimensions to keep D fixed, ensuring that both models maintain the same total recurrent state size. For each budget B∈1,2,4,8,16,32B∈\1,2,4,8,16,32\, we sweep all factorizations (H,J)(H,J) with H×J=BH× J=B and J≤16J≤ 16 for SwiLA, and set H=B,J=1H=B,\,J=1 for DeltaNet. For a given state budget, many additional configurations become available if the head dimension D is also allowed to vary. For instance, a budget of H×J×D2=4×1×322=4,096H× J× D^2=4× 1× 32^2=4,096 can equivalently be realized by H=4H=4 heads with J=4J=4 mixtures and keys and values projected down to dimension d=16d=16, since 4×4×162=4,0964× 4× 16^2=4,096. We fix D=32D=32 for this experiment. Results. Figure 1C plots test R2R^2 against recurrent state size. At every state budget, the best SwiLA configuration achieves a higher R2R^2 than DeltaNet. At B=2B=2, the one-head SwiLA (H1J2) outperforms the two-head DeltaNet (H2) by a clear margin, and this advantage persists as the budget increases. This indicates that allocating budget to mixtures within a head, rather than to independent heads, better approximates softmax attention. As the budget grows, the best SwiLA configurations tend to balance heads and mixtures (e.g., H2J2 at B=4B=4) rather than concentrating all budget in either dimension. Both forms of capacity—parallel heads and per-head mixtures—play complementary roles. In Figure S1, we plot all SwiLA configurations at each budget, where the right panel shows the number of model parameters on the x-axis. Overall, these results confirm that SwiLA provides a more expressive function class for approximating softmax attention than the independent-head decomposition used by DeltaNet. 5.2 Synthetic Benchmarks A Context-Dependent Associative Recall B RegBench 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. We evaluate SwiLA on two synthetic benchmarks: context-dependent associative recall and in-context language learning benchmarks. Context-Dependent Associative Recall (CDAR). Adapted from the associative recall task of Arora et al. (2024a), we design a variant that tests whether a model can maintain multiple separate associative memories and persist context information to select the correct one at retrieval time. In CDAR, the same key maps to different values depending on a context token. Recurrent models must therefore contend with interference across contexts and retain which context is active when answering queries. Crucially, since the context token appears only once at the start of each block and the query region is separated from the context blocks, the task requires temporal persistence of context awareness that cannot be resolved by the short convolutions commonly prepended to recurrent layers. Each example contains C context blocks followed by a query region. A context block begins with a context token cic_i and lists key-value pairs under that context. In the query region, each context’s queries are grouped contiguously: the context token appears once, followed by keys whose values the model must predict. For instance, with C=2C=2 contexts and 3 key-value pairs each: ,k2,v1,2,k1,v1,1,k3,v1,3|,k3,v2,3,k2,v2,2,k1,v2,1⏟context blocks|,k3,?,k1,?,k2,?|,k2,?,k3,?,k1,?⏟queries c_1,k_2,v_1,2,k_1,v_1,1,k_3,v_1,3\;|\;c_2,k_3,v_2,3,k_2,v_2,2,k_1,v_2,1_context blocks\;|\; c_2,k_3,?,k_1,?,k_2,?\;|\;c_1,k_2,?,k_3,?,k_1,?_queries The key set is shared across all contexts, but each (context, key) pair maps to a distinct value. The model must use the context token to determine which mapping is active and carry that information forward to the answer positions. Loss is computed only at answer positions. We use C=4C=4 contexts with 32 key-value pairs each (128 total associations), vocabulary size 8,192, and sequence length 1,024. Figure 2A shows the best test accuracy across learning rates and random seeds. Temporal SwiLA with a single head with dimension D=32D=32 and 4 mixtures (recurrent state size of 2122^12) is the smallest configuration that achieves near-perfect accuracy, suggesting that the mixtures with temporal priors are crucial. All other recurrent models (nearly) solve the task with one layer as the state size increases, while the Transformer requires two, consistent with recent findings that single-layer Transformers cannot perform associative recall beyond chance (Okpekpe and Orvieto, 2025). RegBench: In-Context Language Learning. RegBench (Akyürek et al., 2024) tests how efficiently different architectures learn to perform in-context language learning using probabilistic finite automata (PFAs), where they found that Transformers outperformed recurrent models. Each input packs 10–20 example strings from a randomly generated PFA (4-12 states, alphabet size 4-18) into a single sequence, and the model must learn the PFA’s next-token distribution in-context via next-token prediction. We compare five 8-layer models: SwiLA and Temporal SwiLA, DN-MoM, DeltaNet, and a Transformer. We report accuracy on a held-out test set of 500 PFAs using the best checkpoint by validation loss. As shown in Figure 2B, Temporal SwiLA and DN-MoM exhibit strong sample efficiency, outperforming the Transformer in the low-data regime (1,000 training PFAs), while the Transformer’s accuracy improves steeply with more data and ultimately reaches the highest accuracy. All models achieve similar performance at scale, with DeltaNet slightly trailing. 5.3 Pretraining Language Models We evaluate SwiLA on a realistic language modeling pipeline. All models have 374M parameters, share the same architecture skeleton (24 layers and 1024 hidden size), and are trained on 15B tokens from FineWeb-Edu (Penedo et al., 2024). SwiLA and its variants show strong results across all tasks, narrowing the gap to Transformer++ and even surpassing it on several benchmarks. Table 1: Commonsense Reasoning. Best result per task in bold, second best underlined. Perplexity (↓ ) Accuracy (↑ , %) 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 DN-MoM 40.0 84.5 24.1 62.6 30.4 49.2 52.2 22.8 36.6 62.0 42.5 GDN-MoM 35.3 56.9 27.5 64.6 31.6 50.2 54.2 23.6 37.2 53.3 42.8 SwiLA 28.8 43.0 30.4 66.6 32.0 51.5 56.1 22.8 38.1 60.1 44.7 Gated SwiLA 27.8 31.8 34.6 66.1 32.7 50.4 57.5 25.0 38.1 59.1 45.4 Temporal SwiLA 28.6 39.4 31.5 66.5 31.7 51.1 57.2 23.2 38.8 61.5 45.2 Gated Temporal SwiLA 27.8 31.7 34.5 65.8 32.9 51.8 56.5 24.5 38.5 58.4 45.4 Hybrid GDN 31.6 56.3 28.5 64.8 32.2 50.7 54.6 23.1 39.0 61.4 44.3 Hybrid Temporal SwiLA 27.8 40.0 32.4 66.1 32.6 50.2 57.2 23.3 37.7 61.1 45.1 Commonsense Reasoning. Following Gu and Dao (2024), we evaluate our model on a suite of commonsense reasoning tasks: Wikitext (Merity et al., 2016), LAMBADA (Paperno et al., 2016), PIQA (Bisk et al., 2020), HellaSwag (Zellers et al., 2019), WinoGrande (Sakaguchi et al., 2019), ARC-easy and ARC-challenge (Clark et al., 2018), SIQA (Sap et al., 2019), and BoolQ (Clark et al., 2019). Among the linear attention layers, KDA achieves the strongest commonsense results. Every SwiLA variant attains higher average accuracy than DeltaNet, DN-MoM and GDN-MoM, and the Gated SwiLA variants additionally exceed GDN, slightly surpassing Transformer++ on average and outperforming it on several individual benchmarks. Recall-Intensive Tasks. To measure the ability to perform context-based retrieval and comprehension, we evaluate on six recall-intensive tasks following Arora et al. (2024c): SWDE (Lockard et al., 2019; Arora et al., 2023), FDA (Arora et al., 2023), SQuAD (Rajpurkar et al., 2018), TriviaQA (Joshi et al., 2017), NQ (Kwiatkowski et al., 2019), and DROP (Dua et al., 2019). Table 2 shows that the softmax attention models achieve the highest scores, with FoX attaining the best average, benefiting from their full KV caches that grow linearly with context length. Every SwiLA variant outperforms DeltaNet and both MoM variants on all six recall benchmarks, and Gated SwiLA variants achieve the highest averages among all recurrent models, including KDA. Among the hybrid models, Hybrid Temporal SwiLA surpasses both Hybrid GDN and Transformer++, approaching the recall performance of FoX. These gains confirm that SwiLA’s mixture structure meaningfully improves associative recall compared to a single linear regressor, narrowing the gap to softmax attention while retaining a fixed-size recurrent state. 5.4 Throughput Comparison We profile the training and inference throughput and peak memory of six 1.3B-parameter models, a scale widely adopted for such analyses (Gu and Dao, 2023; Yang et al., 2025): Transformer, GDN, Hybrid GDN, SwiLA, Temporal SwiLA, and Hybrid Temporal SwiLA. For the hybrid models, we follow the 3:1 layer ratio of Team et al. (2025), interleaving three recurrent layers with one Transformer layer. See Section E.4 for details. Table 2: Recall-Intensive Tasks. All scores use the contains metric (%, ↑ ). Best in bold, second best underlined. All inputs are truncated to 2K tokens. 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 GDN-MoM 8.3 3.8 18.4 41.7 10.0 16.3 16.4 SwiLA 12.9 8.0 24.7 42.5 12.9 17.6 19.7 Gated SwiLA 10.2 10.1 26.3 45.2 13.6 17.5 20.5 Temporal SwiLA 13.6 11.6 23.8 42.8 13.3 17.2 20.3 Gated Temporal SwiLA 12.0 7.5 26.0 44.9 13.8 19.2 20.6 Hybrid GDN 25.0 38.2 34.9 43.7 16.0 19.1 29.5 Hybrid Temporal SwiLA 32.3 49.1 30.0 45.8 14.9 18.9 31.8 Overall, SwiLA achieves 4–5× higher end-to-end inference throughput than the Transformer baseline, as it does not require a KV cache and can therefore accommodate substantially larger batch sizes. Relative to GDN, SwiLA is approximately 2× slower at training time and 1.25–1.5× slower at inference time, and it has a larger memory footprint. We expect that part, though not all, of this gap can be narrowed through further engineering, such as custom CUDA kernels or the parallelization of the recurrence discussed in Appendix D. 5.5 Ablation Study We conduct ablation experiments that isolate the contribution of the number of experts, temporal persistence, and the load-balancing loss. We train 115M-parameter SwiLA variants on 5B tokens and report test perplexity on 100M held-out tokens in Figure S5, alongside routing diagnostics defined in Section E.5. We highlight three findings. First, temporal persistence consistently improves perplexity over vanilla SwiLA. Second, load balancing loss is critical for vanilla SwiLA, which otherwise suffers severe expert collapse, but has a negligible effect on Temporal SwiLA; temporal persistence itself appears to act as a natural regularizer against collapse. Third, under a fixed recurrent state size budget, more experts do not necessarily help: increasing the number of experts from 4 to 16 slightly degrades performance, since each expert must then use smaller query, key, and value projections, creating a tradeoff between expert count and per-expert projection capacity. 6 Conclusion We introduce Switching Linear Attention (SwiLA), a sequence layer that retains a recurrent state whose size is independent of sequence length as in linear attention, while representing a nonlinear map from queries to outputs as in softmax attention. Building on the test-time regression (TTR) framework (Wang et al., 2025), which unifies the design of sequence layers around associative recall, we derive SwiLA from a mixture of linear regressions via online expectation–maximization. Across associative recall, in-context language learning, and language modeling benchmarks, SwiLA performs strongly and narrows the gap to softmax attention, even surpassing it in several settings. Limitations While SwiLA increases the expressivity of standard linear attention, it retains a fixed state size. As a result, its capacity cannot grow with the context length, unlike the KV cache of softmax attention11 1 Though even softmax attention struggles with attention rot in long contexts (Hsieh et al., 2024).. Moreover, as SwiLA uses a nonlinear recurrence, we currently pretrain SwiLA with sequential evaluation across the sequence length, whereas both softmax and linear attention mechanisms are designed to allow for training that is parallelized across the sequence length. While we can parallelize the nonlinear recurrence of SwiLA using multiple Newton iterations, further optimizations are required to design custom kernels that yield speed-ups over optimized kernels for sequential evaluation. Future Outlook SwiLA demonstrates the promise of TTR: we can design sequence layers with desired in-context learning properties by choosing appropriate regression models and update algorithms. We anticipate TTR to open up a range of possibilities between the efficiency of linear attention and the expressivity of softmax attention. For example, while SwiLA shows the promise of mixture modeling as occupying a useful middle ground between these two extremes, other function classes such as local linear regression (Zuo et al., 2026) may prove useful in fully exploring this space. Moreover, SwiLA’s enhanced ability to approximate softmax attention (Figure 1C) while retaining a fixed-state size appears promising for distillation—converting pretrained softmax layers into recurrent layers for efficient inference (Bick et al., 2024; Wang et al., 2024; Zhang et al., 2024a). Acknowledgments H.D.L. is supported by the Ketterer-Vorwald Stanford Interdisciplinary Graduate Fellowship. N.Z. is supported by Postdoc.Mobility grant P500-2_235376 from the Swiss National Science Foundation. E.B.F. is supported in part by ONR Grant N00014-22-1-2110 and the Stanford Institute for Human-Centered Artificial Intelligence (HAI), and is a Biohub, San Francisco, Investigator. S.W.L. is supported by grants from the NIH (U01NS136507, R01NS131987, R01NS113119, RF1MH133778, R01AG097491, R01NS130789), the NSF (2440859), and the Simons and McKnight Foundations. We thank the members of the Linderman Lab and the Dynamode Lab for feedback throughout this project. We are also grateful for GPU support from the Stanford Research Computing Center’s Sherlock cluster, from VESSL AI, and from Marlowe (Kapfer et al., 2025), Stanford University’s GPU-based Computational Instrument, supported by Stanford HAI and Stanford Research Computing. References Akyürek et al. (2023) E. Akyürek, D. Schuurmans, J. Andreas, T. Ma, and D. Zhou What learning algorithm is in-context learning? Investigations with linear models. In International Conference on Learning Representations (ICLR), Cited by: §1. Akyürek et al. (2024) E. Akyürek, B. Wang, Y. Kim, and J. Andreas In-context language learning: architectures and algorithms. arXiv preprint arXiv:2401.12973. Cited by: §5.2. Alayrac et al. (2022) J. Alayrac, J. Donahue, P. Luc, A. Miech, I. Barr, Y. Hasson, K. Lenc, A. Mensch, K. Millican, M. Reynolds, et al. Flamingo: a visual language model for few-shot learning. In Advances in Neural Information Processing Systems, Vol. 35, p. 23716–23736. Cited by: §1. Arora et al. (2024a) S. Arora, S. Eyuboglu, A. Timber, I. Johnson, M. Poli, J. Zou, A. Rudra, and C. Ré Zoology: measuring and improving recall in efficient language models. In International Conference on Learning Representations, Cited by: Appendix C, §E.2, §E.2, §1, §5.2. Arora et al. (2024b) S. Arora, S. Eyuboglu, M. Zhang, A. Timber, S. Alberti, D. Zinber, J. Zou, C. Re, and C. Ré Simple linear attention language models balance the recall-throughput tradeoff. In International Conference on Machine Learning, Cited by: §1. Arora et al. (2024c) S. Arora, A. Timalsina, A. Singhal, B. Spector, S. Eyuboglu, X. Zhao, A. Rao, A. Rudra, and C. Ré Just read twice: Closing the recall gap for recurrent language models. arXiv preprint arXiv:2407.05483. Cited by: §5.3. Arora et al. (2023) S. Arora, B. Yang, S. Eyuboglu, A. Narayan, A. Hojel, I. Trummer, and C. Ré Language models enable simple systems for generating structured views of heterogeneous data lakes. arXiv preprint arXiv:2304.09433. Cited by: §5.3. Beck et al. (2025) M. Beck, K. Pöppel, P. Lippe, and S. Hochreiter Tiled flash linear attention: more efficient linear rnn and xLSTM kernels. In Advances in Neural Information Processing Systems (NeurIPS), Cited by: Appendix C. Behrouz et al. (2024) A. Behrouz, P. Zhong, and V. Mirrokni Titans: learning to memorize at test time. arXiv preprint arXiv:2501.00663. Cited by: §A.1. Bellen and Zennaro (1989) A. Bellen and M. Zennaro Parallel algorithms for initial-value problems for difference and differential equations. Journal of Computational and Applied Mathematics 25 (3), p. 341–350. Cited by: §D.1. Bick et al. (2024) A. Bick, K. Li, E. Xing, J. Z. Kolter, and A. Gu Transformers to SSMs: Distilling quadratic knowledge to subquadratic models. In Neural Information Processing Systems (NeurIPS), Cited by: §6. Bishop (2006) C. M. Bishop Pattern recognition and machine learning. Springer. Cited by: §3.1. Bisk et al. (2020) Y. Bisk, R. Zellers, R. L. Bras, J. Gao, and Y. Choi PIQA: reasoning about physical commonsense in natural language. In Thirty-Fourth AAAI Conference on Artificial Intelligence, Cited by: §5.3. Blelloch (1990) G. E. Blelloch Prefix Sums and Their Applications. Technical report Technical Report CMU-CS-90-190, Carnegie Mellon University, School of Computer Science. Cited by: §D.1, §D.1, Appendix D, §4.1. Brown et al. (2020) T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al. Language models are few-shot learners. In Advances in Neural Information Processing Systems, Vol. 33, p. 1877–1901. Cited by: §1. Cappé and Moulines (2009) O. Cappé and E. Moulines On-line expectation–maximization algorithm for latent data models. Journal of the Royal Statistical Society Series B: Statistical Methodology 71 (3), p. 593–613. Cited by: §1, 2nd item. Clark et al. (2019) C. Clark, K. Lee, M. Chang, T. Kwiatkowski, M. Collins, and K. Toutanova BoolQ: Exploring the surprising difficulty of natural yes/no questions. In Proceedings of the 2019 conference of the north American chapter of the association for computational linguistics: Human language technologies, volume 1 (long and short papers), p. 2924–2936. Cited by: §5.3. Clark et al. (2018) P. Clark, I. Cowhey, O. Etzioni, T. Khot, A. Sabharwal, C. Schoenick, and O. Tafjord Think you have solved question answering? Try ARC, the AI2 reasoning challenge. ArXiv abs/1803.05457. Cited by: §5.3. Danieli et al. (2026) F. Danieli, P. Rodriguez, M. Sarabia, X. Suau, and L. Zappella ParaRNN: Unlocking parallel training of nonlinear rnns for large language models. In International Conference on Learning Representations (ICLR), Cited by: 2nd item, Appendix D, §4.1. Danieli et al. (2023) F. Danieli, M. Sarabia, X. Suau, P. Rodríguez, and L. Zappella DeepPCR: Parallelizing Sequential Operations in Neural Networks. In Advances in Neural Information Processing Systems (NeurIPS), Cited by: §D.1. Dao et al. (2022) T. Dao, D. Fu, S. Ermon, A. Rudra, and C. Ré FlashAttention: Fast and memory-efficient exact attention with IO-awareness. In Advances in neural information processing systems, Cited by: Appendix C. Dao and Gu (2024) T. Dao and A. Gu Transformers are SSMs: generalized models and efficient algorithms through structured state space duality. In International Conference on Machine Learning (ICML), Cited by: §A.1. De Veaux (1989) R. D. De Veaux Mixtures of linear regressions. Computational Statistics & Data Analysis 8 (3), p. 227–245. Cited by: 1st item. Dempster et al. (1977) A. P. Dempster, N. M. Laird, and D. B. Rubin Maximum likelihood from incomplete data via the EM algorithm. Journal of the royal statistical society: series B (methodological) 39 (1), p. 1–22. Cited by: 2nd item. Dosovitskiy et al. (2021) A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, et al. An image is worth 16x16 words: transformers for image recognition at scale. In International Conference on Learning Representations, Cited by: §1. Du et al. (2026) J. Du, W. Sun, D. Lan, J. Hu, and Y. Cheng Mom: linear sequence modeling with mixture-of-memories. In International Conference on Learning Representations (ICLR), Cited by: §E.3, §E.3, §4.2, §4, §5. Dua et al. (2019) D. Dua, Y. Wang, P. Dasigi, G. Stanovsky, S. Singh, and M. Gardner DROP: a reading comprehension benchmark requiring discrete reasoning over paragraphs. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), p. 2368–2378. Cited by: §5.3. Farsang and Grosu (2025) M. Farsang and R. Grosu Scaling up liquid-resistance liquid-capacitance networks for efficient sequence modeling. In Advances in Neural Information Processing Systems (NeurIPS), Cited by: §4.2. Fedus et al. (2022) W. Fedus, B. Zoph, and N. Shazeer Switch transformers: scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research 23 (120), p. 1–39. Cited by: §3.5. Finn et al. (2017) C. Finn, P. Abbeel, and S. Levine Model-agnostic meta-learning for fast adaptation of deep networks. In International Conference on Machine Learning (ICML), Cited by: Appendix C. Garg et al. (2022) S. Garg, D. Tsipras, P. S. Liang, and G. Valiant What can transformers learn in-context? A case study of simple function classes. In Neural Information Processing Systems (NeurIPS), Cited by: §1. Gelada et al. (2025) C. Gelada, J. Buckman, S. Zhang, and T. Bach Scaling context requires rethinking attention. arXiv preprint arXiv:2507.04239. Cited by: Appendix C. Gonzalez et al. (2025) X. Gonzalez, E. K. Buchanan, H. D. Lee, J. W. Liu, K. A. Wang, D. M. Zoltowski, C. Ré, and S. W. Linderman A unifying framework for parallelizing sequential models with linear dynamical systems. arXiv preprint arXiv:2509.21716. Cited by: §D.1. Gonzalez et al. (2024) X. Gonzalez, A. Warrington, J. T. Smith, and S. W. Linderman Towards scalable and stable parallelization of nonlinear RNNs. Advances in Neural Information Processing Systems 37, p. 5817–5849. Cited by: §D.1, §D.2, Appendix D, §4.1. Gonzalez (2026) X. Gonzalez Unifying optimization and dynamics to parallelize sequential computation: a guide to parallel Newton methods for breaking sequential bottlenecks. Ph.D. Thesis, Stanford University. External Links: Link Cited by: 2nd item, §D.1, §D.2, Appendix D, §4.1. Gu and Dao (2023) A. Gu and T. Dao Mamba: linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752. Cited by: Figure S4, §5.4. Gu and Dao (2024) A. Gu and T. Dao Mamba: linear-time sequence modeling with selective state spaces. In Conference on Language Modeling (COLM), Cited by: §5.3. Gu (2025) A. Gu On the tradeoffs of state space models and transformers. External Links: Link Cited by: Appendix C. Guo et al. (2026) H. Guo, S. Yang, T. Goel, E. P. Xing, T. Dao, and Y. Kim Log-linear attention. In International Conference on Learning Representations (ICLR), Cited by: Appendix C. Ha et al. (2017) D. Ha, A. Dai, and Q. V. Le Hypernetworks. In International Conference on Learning Representations (ICLR), Cited by: Appendix C. Han et al. (2024a) D. Han, Y. Pu, Z. Xia, Y. Han, X. Pan, X. Li, J. Lu, S. Song, and G. Huang Bridging the divide: reconsidering softmax and linear attention. Advances in Neural Information Processing Systems 37, p. 79221–79245. Cited by: Appendix C. Han et al. (2024b) D. Han, T. Ye, Y. Han, Z. Xia, S. Pan, P. Wan, S. Song, and G. Huang Agent attention: on the integration of softmax and linear attention. In European conference on computer vision, p. 124–140. Cited by: Appendix C. Hinton and Plaut (1987) G. E. Hinton and D. C. Plaut Using fast weights to deblur old memories. In Proceedings of the ninth annual conference of the Cognitive Science Society, p. 177–186. Cited by: Appendix C. Hsieh et al. (2024) C. Hsieh, S. Sun, S. Kriman, S. Acharya, D. Rekesh, F. Jia, Y. Zhang, and B. Ginsburg RULER: what’s the real context size of your long-context language models?. In Conference on Language Modeling (COLM), Cited by: footnote 1. Irie and Gershman (2025) K. Irie and S. J. Gershman Fast weight programming and linear transformers: From machine learning to neurobiology. Transactions on Machine Learning Research (TMLR). Cited by: Appendix C. Irie et al. (2021) K. Irie, I. Schlag, R. Csordás, and J. Schmidhuber Going beyond linear transformers with recurrent fast weight programmers. In Neural Information Processing Systems (NeurIPS), Cited by: Appendix C. Joshi et al. (2017) M. Joshi, E. Choi, D. S. Weld, and L. Zettlemoyer TriviaQA: A large scale distantly supervised challenge dataset for reading comprehension. In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), p. 1601–1611. Cited by: §5.3. Kapfer et al. (2025) C. Kapfer, K. Stine, B. Narasimhan, C. Mentzel, and E. Candès Marlowe: stanford’s gpu-based computational instrument. Zenodo. External Links: Document, Link Cited by: Acknowledgments. Katharopoulos et al. (2020) A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret Transformers are RNNs: fast autoregressive transformers with linear attention. In International Conference on Machine Learning, p. 5156–5165. Cited by: §A.1, Appendix C, §1, §2. Kwiatkowski et al. (2019) T. Kwiatkowski, J. Palomaki, O. Redfield, M. Collins, A. Parikh, C. Alberti, D. Epstein, I. Polosukhin, J. Devlin, K. Lee, et al. Natural questions: A benchmark for question answering research. Transactions of the Association for Computational Linguistics 7, p. 453–466. Cited by: §5.3. Lim et al. (2024) Y. H. Lim, Q. Zhu, J. Selfridge, and M. F. Kasim Parallelizing non-linear sequential models over the sequence length. In International Conference on Learning Representations, Cited by: 2nd item, §D.1, Appendix D, §4.1. Lin et al. (2025) Z. Lin, E. Nikishin, X. He, and A. Courville Forgetting transformer: softmax attention with a forget gate. In International Conference on Learning Representations, Vol. 2025, p. 69704–69738. Cited by: §E.3, §5. Liu et al. (2024) B. Liu, R. Wang, L. Wu, Y. Feng, P. Stone, and Q. Liu Longhorn: state space models are amortized online learners. arXiv preprint arXiv:2407.14207. Cited by: §A.1. Lockard et al. (2019) C. Lockard, P. Shiralkar, and X. L. Dong Openceres: when open information extraction meets the semi-structured web. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), p. 3047–3056. Cited by: §5.3. Martin and Cundy (2018) E. Martin and C. Cundy Parallelizing linear recurrent neural nets over sequence length. In International Conference on Learning Representations (ICLR), Cited by: §D.1. Merity et al. (2016) S. Merity, C. Xiong, J. Bradbury, and R. Socher Pointer sentinel mixture models. External Links: 1609.07843 Cited by: §5.3. Merrill et al. (2026) W. Merrill, H. Jiang, Y. Li, and A. Sabharwal Why are linear RNNs more parallelizable?. arXiv preprint arXiv:2603.03612. Cited by: §4.2. Mishra et al. (2026) M. Mishra, S. Tan, I. Stoica, J. Gonzalez, and T. Dao M2^2 RNN: Non-Linear RNNs with Matrix-Valued States for Scalable Language Modeling. arXiv preprint arXiv:2603.14360. Cited by: Appendix C, §4.2. Okpekpe and Orvieto (2025) D. Okpekpe and A. Orvieto Revisiting associative recall in modern recurrent models. arXiv preprint arXiv:2508.19029. Cited by: §5.2. Pan et al. (2026) Y. Pan, Y. An, Z. Li, Y. Chou, R. Zhu, X. Wang, M. Wang, J. Wang, and G. Li Scaling linear attention with sparse state expansion. In International Conference on Learning Representations (ICLR), Cited by: §4.2, §4. Paperno et al. (2016) D. Paperno, G. Kruszewski, A. Lazaridou, N. Pham, R. Bernardi, S. Pezzelle, M. Baroni, G. Boleda, and R. Fernández The LAMBADA dataset: Word prediction requiring a broad discourse context. In Proceedings of the 54th annual meeting of the association for computational linguistics (volume 1: Long papers), p. 1525–1534. Cited by: §5.3. Pearson (1894) K. Pearson Contributions to the mathematical theory of evolution. Proceedings of the Royal Society of London 54 (326-330), p. 329–333. Cited by: 1st item. Penedo et al. (2024) G. Penedo, H. Kydlíček, A. Lozhkov, M. Mitchell, C. Raffel, L. Von Werra, T. Wolf, et al. The FineWeb datasets: Decanting the web for the finest text data at scale. Advances in Neural Information Processing Systems 37, p. 30811–30849. Cited by: §E.3, §5.3. Peng et al. (2025) L. Peng, A. Chattopadhyay, L. Zancato, E. Nunez, W. Xia, and S. Soatto Gated kalmanet: a fading memory layer through test-time ridge regression. arXiv preprint arXiv:2511.21016. Cited by: Appendix C. Pöppel et al. (2025) K. Pöppel, M. Beck, and S. Hochreiter FlashRNN: I/O-aware optimization of traditional RNNs on modern hardware. In International Conference on Learning Representations (ICLR), Cited by: Appendix C. Rajpurkar et al. (2018) P. Rajpurkar, R. Jia, and P. Liang Know what you don’t know: unanswerable questions for squad. In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers), p. 784–789. Cited by: §5.3. Sakaguchi et al. (2019) K. Sakaguchi, R. L. Bras, C. Bhagavatula, and Y. Choi WinoGrande: an adversarial Winograd schema challenge at scale. arXiv preprint arXiv:1907.10641. Cited by: §5.3. Sap et al. (2019) M. Sap, H. Rashkin, D. Chen, R. Le Bras, and Y. Choi Social IQa: commonsense reasoning about social interactions. In Proceedings of the 2019 conference on empirical methods in natural language processing and the 9th international joint conference on natural language processing (EMNLP-IJCNLP), p. 4463–4473. Cited by: §5.3. Schlag et al. (2021) I. Schlag, K. Irie, and J. Schmidhuber Linear transformers are secretly fast weight programmers. In International Conference on Machine Learning (ICML), Cited by: §A.1, §2, §3.2, §4.1, §4, §5. Schmidhuber (1992) J. Schmidhuber Learning to control fast-weight memories: an alternative to dynamic recurrent networks. Neural Computation 4 (1), p. 131–139. Cited by: Appendix C. Shazeer et al. (2017) N. Shazeer, A. Mirhoseini, K. Maziarz, A. Davis, Q. Le, G. Hinton, and J. Dean Outrageously large neural networks: the sparsely-gated mixture-of-experts layer. arXiv preprint arXiv:1701.06538. Cited by: §3.5, §3.5. Shih et al. (2023) A. Shih, S. Belkhale, S. Ermon, D. Sadigh, and N. Anari Parallel sampling of diffusion models. In Advances in Neural Information Processing Systems (NeurIPS), Cited by: §D.1. Siems et al. (2026) J. Siems, R. Grazzi, K. Kalinin, H. Ballani, and B. Rahmani Learning state-tracking from code using linear RNNs. arXiv preprint arXiv:2602.14814. Cited by: §4.2. Smith et al. (2023) J. T.H. Smith, A. Warrington, and S. W. Linderman Simplified state space layers for sequence modeling. In International Conference on Learning Representations (ICLR), Cited by: §D.1. Song et al. (2021) Y. Song, C. Meng, R. Liao, and S. Ermon Accelerating feedforward computation via parallel nonlinear equation solving. In International Conference on Machine Learning (ICML), Cited by: §D.1. Sun et al. (2025) Y. Sun, X. Li, K. Dalal, J. Xu, A. Vikram, G. Zhang, Y. Dubois, X. Chen, X. Wang, S. Koyejo, T. Hashimoto, and C. Guestrin Learning to (learn at test time): RNNs with expressive hidden states. In International Conference on Machine LEarning (ICML), Cited by: Appendix C. Tandon et al. (2025) A. Tandon, K. Dalal, X. Li, D. Koceja, M. Rød, S. Buchanan, X. Wang, J. Leskovec, S. Koyejo, T. Hashimoto, C. Guestrin, J. McCaleb, Y. Choi, and Y. Sun End-to-end test-time training for long context. arXiv preprint arXiv:2512.23675. Cited by: Appendix C, Appendix C. Tang et al. (2024) Z. Tang, J. Tang, H. Luo, F. Wang, and T. Chang Accelerating parallel sampling of diffusion models. In International Conference on Machine Learning (ICML), Cited by: §D.1. Team et al. (2025) K. Team, Y. Zhang, Z. Lin, X. Yao, J. Hu, F. Meng, C. Liu, X. Men, S. Yang, Z. Li, et al. Kimi linear: an expressive, efficient attention architecture. arXiv preprint arXiv:2510.26692. Cited by: §B.4, §E.3, §E.3, §3.4, §5.4, §5. Touvron et al. (2023) H. Touvron, T. Lavril, G. Izacard, X. Martinet, M. Lachaux, T. Lacroix, B. Rozière, N. Goyal, E. Hambro, F. Azhar, et al. Llama: open and efficient foundation language models. arXiv preprint arXiv:2302.13971. Cited by: §E.3, §5. Vaswani et al. (2017) A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin Attention is all you need. In Advances in Neural Information Processing Systems, Vol. 30. Cited by: Appendix C, §1. von der Malsburg (1981) C. von der Malsburg The correlation theory of brain function. Internal Report Technical Report 81-2, Department of Neurobiology, Max Planck Institute for Biophysical Chemistry, Goettingen. Cited by: Appendix C. von Oswald et al. (2023) J. von Oswald, E. Niklasson, E. Randazzo, J. Sacramento, A. Mordvintsev, A. Zhmoginov, and M. Vladymyrov Transformers learn in-context by gradient descent. In International Conference on Machine Learning, p. 35150–35174. Cited by: §1. von Oswald et al. (2026) J. von Oswald, N. Scherrer, S. Kobayashi, L. Versari, S. Yang, M. Schlegel, K. Maile, Y. Schimpf, O. Sieberling, A. Meulemans, et al. MesaNet: sequence modeling by locally optimal test-time training. In International Conference on Learning Representations (ICLR), Cited by: Appendix C, Appendix C. Von Oswald et al. (2024) J. Von Oswald, M. Schlegel, A. Meulemans, S. Kobayashi, E. Niklasson, N. Zucchet, N. Scherrer, N. Miller, M. Sandler, M. Vladymyrov, et al. Uncovering mesa-optimization algorithms in transformers. arXiv preprint arXiv:2309.05858. Cited by: Appendix C, Appendix C. Wang et al. (2026) B. Wang, M. Zecchin, and O. Simeone Distributed Associative Memory via Online Convex Optimization. In ICASSP 2026 - 2026 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), Cited by: Appendix C. Wang et al. (2024) J. Wang, D. Paliotta, A. May, A. M. Rush, and T. Dao The Mamba in the Llama: distilling and accelerating hybrid models. In Neural Information Processing Systems (NeurIPS), Cited by: §6. Wang et al. (2025) K. A. Wang, J. Shi, and E. B. Fox Test-time regression: A unifying framework for designing sequence models with associative memory. arXiv preprint arXiv:2501.12352. Cited by: §A.1, Appendix C, §1, §1, §2, §2, §3.4, §6. Yang et al. (2025) S. Yang, J. Kautz, and A. Hatamizadeh Gated delta networks: improving Mamba2 with delta rule. In International Conference on Learning Representations (ICLR), Cited by: §A.1, §B.4, Figure S3, §E.3, §3.4, §3.4, §5.4, §5. Yang et al. (2024a) S. Yang, B. Wang, Y. Shen, R. Panda, and Y. Kim Gated linear attention transformers with hardware-efficient training. In International Conference on Machine Learning (ICML), Cited by: §A.1. Yang et al. (2024b) S. Yang, B. Wang, Y. Zhang, Y. Shen, and Y. Kim Parallelizing linear transformers with the delta rule over sequence length. Advances in neural information processing systems 37, p. 115491–115522. Cited by: §A.1, 1st item, §D.1, §D.2, §D.3, Appendix D, Appendix D, §E.3, §E.3, §2, §3.2, §4.1, §4, §5. Yang and Zhang (2024) S. Yang and Y. Zhang FLA: a Triton-based library for hardware-efficient implementations of linear attention mechanism. External Links: Link Cited by: Appendix C, Appendix E, §3.4. Yang (2024) S. Yang DeltaNet explained: a gentle and comprehensive introduction to the deltanet. Note: https://sustcsonglin.github.io/blog/2024/deltanet-2/blogpost Cited by: §D.3, Appendix D, Appendix D, Appendix D, §4.1. Zadouri et al. (2026) T. Zadouri, M. Hoehnerbach, J. Shah, T. Liu, V. Thakkar, and T. Dao FlashAttention-4: algorithm and kernel pipelining co-design for asymmetric hardware scaling. arXiv preprint arXiv:2603.05451. Cited by: Appendix C. Zattra et al. (2025) R. Zattra, G. Baggio, U. Casti, A. Ferrante, and F. Ticozzi Context-selective state space models: feedback is all you need. arXiv preprint arXiv:2510.14027. Cited by: §4.2. Zellers et al. (2019) R. Zellers, A. Holtzman, Y. Bisk, A. Farhadi, and Y. Choi HellaSwag: can a machine really finish your sentence?. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, Cited by: §5.3. Zhang et al. (2024a) M. Zhang, S. Arora, R. Chalamala, A. Wu, B. Spector, A. Singhal, K. Ramesh, and C. Ré Lolcats: on low-rank linearizing of large language models. arXiv preprint arXiv:2410.10254. Cited by: §6. Zhang et al. (2024b) M. Zhang, K. Bhatia, H. Kumbong, and C. Ré The hedgehog & the porcupine: expressive linear attentions with softmax mimicry. In International Conference on Learning Representations (ICLR), Cited by: Appendix C. Zhang and Yang (2025) Y. Zhang and S. Yang Flame: flash language modeling made easy. External Links: Link Cited by: §E.3. Zhang et al. (2026) Y. Zhang, W. Sun, X. Sun, W. Ding, R. Xie, H. Wang, J. Chen, J. Liu, S. Wang, Y. Zhang, et al. A survey of linear attention: algorithm, theory, application, and infrastructure. Authorea Preprints. Cited by: Appendix C. Zuo et al. (2026) Y. Zuo, Y. Yin, Z. Zeng, A. Li, B. Zhu, and Z. Wang Local linear attention: an optimal interpolation of linear and softmax attention for test-time regression. In International Conference on Learning Representations (ICLR), Cited by: Appendix C, §6. Appendix A Extended Background A.1 Design Space of Test-Time Regression The test-time regression framework of Wang et al. (2025) recovers a broad class of existing sequence layers through different design choices. Linear attention (Katharopoulos et al., 2020) corresponds to linear regression with a whitened design matrix approximation. Adding decaying regression weights yields gated linear attention and state-space models (Yang et al., 2024a; Dao and Gu, 2024). Learning the regression in an online streaming manner with stochastic gradient descent produces DeltaNet (Schlag et al., 2021; Yang et al., 2024b), with further variations recovering Longhorn (Liu et al., 2024), Gated DeltaNet (Yang et al., 2025), and Titans (Behrouz et al., 2024). Softmax attention arises as a nonparametric local constant regressor. A.2 Mixtures of Linear Regressions Consider a dataset (xt,yt)t=1T\(x_t,y_t)\_t=1^T with xt∈ℝDx_t ^D and yt∈ℝDy_t ^D, assumed to be generated from a mixture of J linear Gaussian regression models. Let zt∈1,…,Jz_t∈\1,…,J\ denote a latent component indicator with mixing weights π=(π1,…,πJ)π=( _1,…, _J), where πj>0 _j>0 and ∑j=1Jπj=1 _j=1^J _j=1. Conditioned on zt=jz_t=j, the output yty_t is generated according to yt|xt,zt=j∼(Wjxt,Σj),y_t \ | \ x_t,z_t=j (W_jx_t, _j), (A.1) where Wj∈ℝD×DW_j ^D× D is the regression matrix and Σj∈ℝD×D _j ^D× D is the noise covariance, often restricted to the isotropic form Σj=σj2ID _j= _j^2I_D. The complete parameter set is Θ=πj,Wj,Σjj=1J =\ _j,W_j, _j\_j=1^J. Given parameters Θ , inference over the latent variables is performed by computing the responsibilities rtm=πm(yt,Wmxt,Σm)∑j′=1Jπj′(yt,Wj′xt,Σj′),r_tm= _m\,N(y_t;\,W_mx_t,\, _m) _j =1^J _j \,N(y_t;\,W_j x_t,\, _j ), (A.2) which represent the posterior probability that observation (xt,yt)(x_t,y_t) was generated by component j. The model parameters and responsibilities can be estimated using the Expectation–Maximization (EM) algorithm, which alternates between computing the responsibilities (E-step) and maximizing the expected complete-data log-likelihood with respect to Θ (M-step). Appendix B Model Specification Details B.1 Derivation of SwiLA Recurrence In this appendix, we work out in more detail the derivation of the SwiLA recurrence in eq. 3.3. To make bookkeeping easier, we define πj _j :=πjd(kt)exp(−12δtjd2) := _jd(k_t)\, \! (- 12\, _tjd^2 ) (B.1) Z Z :=∑jπj := _j _j (B.2) Taking derivatives, it follows that ∂πj∂wt−1,jd ∂ _j∂ w_t-1,jd =πjd(kt)exp(−δtjd2/2)δtjdkt = _jd(k_t) (- _tjd^2/2) _tjdk_t =πjδtjdkt = _j _tjdk_t ∂Z∂wt−1,jd ∂ Z∂ w_t-1,jd =∂πj∂wt−1,jd = ∂ _j∂ w_t-1,jd Thus, we obtain ∂ℓtd∂wt−1,jd ∂ _td∂ w_t-1,jd =∂πj/∂wt−1,jdZ = ∂ _j∂ w_t-1,jdZ =rtjdδtjdkt.■ =r_tjd _tjdk_t. B.2 Details of the SwiLA Layer In the practical implementation of SwiLA, we obtain many quantities of interest as learnable projections of the input xt∈ℝdinx_t ^d_in. In keeping with standard practice, we generate keys, queries, and values as projections of the input xtx_t, according to qt q_t =θQxt = _Qx_t kt k_t =θKxt = _Kx_t vt v_t =θVxt. = _Vx_t. These projections can also have nonlinearities applied to them, such as the sigmoid linear unit SiLU(x)≔xσ(x)SiLU(x) xσ(x). For language model pretraining (Section 5.3), we apply SiLUSiLU to the queries, keys, and values. Up to reshaping, the learning rate β comes from a linear mapping from the inputs, followed by a sigmoid function to keep the learning rate in [0,1][0,1]. Similarly, our priors on the keys π(kt)π(k_t) and queries π(qt)π(q_t) for the mixture memberships come from linear projections applied to the inputs xtx_t, followed by the application of softmax over the mixture-component dimension (to make sure the priors are valid probabilities). In language model pretraining, we use a low-rank projection to reduce the number of parameters used by these projections. B.3 Temporal SwiLA Details We provide more details on how we encourage temporal persistence in our priors over mixture component assignment πd(kt) _d(k_t) and πd(qt) _d(q_t). In our implementation, the temporal recurrence for both these priors is a convex combination of the previous filtered posterior and a new input-dependent distribution. The relative weighting is controlled by g, which is a learned function of the input xtx_t, dictating the amount of persistence of stickiness in our distribution over mixture component assignment. The range of g is [0,1][0,1]: if g is close to 1, then we use a new probability distribution suggested by the input, while if g is close to 0, at time t we use a very similar probability distribution to that used at time t−1t-1. To make this recursion formal, for time t we define r¯td r_td to be our probability distribution over mixture components for the keys, and ρ¯td ρ_td to be our probability distribution over the mixture components for the queries. Using these definitions, we can then define the recurrences as gtdk g_td^k =σ((θgk)⊤xt+bgk),gtdq=σ((θgq)⊤xt+bgq) =σ (( _g^k) x_t+b_g^k ), g_td^q=σ (( _g^q) x_t+b_g^q ) (B.3) r¯tjd r_tjd =(1−gtdk)rt−1,jd+gtdk(πjdk(kt)), =(1-g^k_td)\,r_t-1,jd+g^k_td\, ( _jd^k(k_t) ), (B.4) ρ¯tjd ρ_tjd =(1−gtdq)ρ¯t−1,jd+gtdq(πjdq(qt)), =(1-g_td^q)\, ρ_t-1,jd\;+\;g_td^q\,\! ( _jd^q(q_t) ), (B.5) where πdk(kt) _d^k(k_t) and πdq(qt) _d^q(q_t) come from the learnable projections followed by softmaxes discussed in Section B.2. Finally, at time t, we use r¯tjd r_tjd as our prior for the key mixture assignment in eq. 3.4, and we use ρ¯tjd ρ_tjd as our prior for the query mixture assignment in eq. 3.5. B.4 Gated SwiLA Details We provide more details on the equivalence of the two gating formulations in Section 3.4 when J=1J=1, and on the granularity at which the gate can be applied. Equivalence of the two gating formulations when J=1J=1. Eqs. (3.7) and (3.8) differ only in whether the prediction errors and responsibilities are evaluated at the undecayed or the decayed state. When J=1J=1, we obtain rt1d=rt1d−=1r_t1d=r^-_t1d=1, so the responsibilities drop out of both updates. Dropping the index j and writing each as an affine map of wt−1,dw_t-1,d, eq. (3.7):wtd . \ eq:post-decay: w_td =(αtdI−βtdktkt⊤)wt−1,d+βtdvtdkt, = ( _tdI- _td\,k_tk_t )w_t-1,d+ _tdv_tdk_t, (B.6) eq. (3.8):wtd . \ eq:pre-decay: w_td =αtd(I−βtdktkt⊤)wt−1,d+βtdvtdkt. = _td (I- _td\,k_tk_t )w_t-1,d+ _tdv_tdk_t. (B.7) With αtd>0 _td>0, re-parameterizing (αtd,βtd,vtd)↦(αtd,βtd/αtd,αtdvtd)( _td, _td,v_td) ( _td,\, _td/ _td,\, _tdv_td) makes the two equivalent. Granularity of the gate. Section 3.4 places an isotropic Gaussian prior wjd∼(0,λtjd−1I)w_jd (0, _tjd^-1I) on each row of the state. More generally, we may place a prior wjdd′∼(0,λtjdd′−1)w_jdd (0, _tjdd ^-1), where d′d indexes the key dimension, so that wjdd′w_jdd is the weight connecting key dimension d′d to output dimension d within mixture component j. One step of gradient ascent on the resulting objective yields wtjd=αtjd⊙wt−1,jd+βtjdrtjdδtjdkt,αtjdd′≔1−βtjdλtjdd′,w_tjd= _tjd w_t-1,jd+ _tjd\,r_tjd\, _tjd\,k_t, _tjdd 1- _tjd _tjdd , (B.8) where ⊙ is the Hadamard product. Collecting the fast weights into Wtj∈ℝD×DW_tj ^D× D with (Wtj)dd′=wtjdd′(W_tj)_d =w_tjdd , such that rows index output dimensions and columns index key dimensions, the granularity of the decay is controlled by restricting how λtjdd′ _tjdd varies over its indices. Tying the precision across the output index, λtjdd′=λtjd′ _tjdd = _tjd , gives a right diagonal scaling Wt−1,jDiag(atj)W_t-1,jDiag(a_tj), i.e. one gate per key dimension, as in KDA (Team et al., 2025). Tying it across the key index, λtjdd′=λtjd _tjdd = _tjd, gives a left diagonal scaling Diag(αtj)Wt−1,jDiag( _tj)W_t-1,j, i.e. one gate per output dimension, which is eq. (3.7). The restriction λtjdd′=λtj _tjdd = _tj recovers the scalar per-head decay of GDN (Yang et al., 2025). Appendix C Extended Related Work Test-time regression SwiLA is motivated by the test-time regression (TTR) framework of Wang et al. (2025), which allows for the principled design of sequence layers based on ICL desiderata. TTR is related to many other important concepts in sequence modeling, including mesa-optimization, test-time training (T), and fast-weights. Mesa-optimization (Von Oswald et al., 2024) also shows that effective sequence layers optimize a learning objective during ICL, but often emphasizes the role of pretraining in unlocking such capabilities; whereas TTR focuses on how a chosen recurrent update affects ICL regardless of pretraining. T also considers the recurrent state to be the parameters of a function, but often considers more general functional relationships (like MLPs) or more general losses, such as reconstruction (Sun et al., 2025) or next-token prediction (Tandon et al., 2025), and is pretrained to encourage metalearning (Finn et al., 2017). The fast-weights paradigm (von der Malsburg, 1981; Hinton and Plaut, 1987; Schmidhuber, 1992; Irie et al., 2021; Irie and Gershman, 2025) also distinguishes between pretrained “slow” weights and test-time “fast” weights, but tends to focus more on the fast weights being generated by hypernetworks (Ha et al., 2017). Other papers have built on TTR with different design goals in mind. For example, Von Oswald et al. (2024); von Oswald et al. (2026) construct a linear attention sequence layer they call a mesa-layer, which uses correct Newton updates for linear regression (instead of the whitened design matrix approximation used in linear attention (Katharopoulos et al., 2020)) to fully optimize the regression objective. Peng et al. (2025) introduce Gated KalmaNet Attention (GKA), which draws inspiration from Kalman filtering. Zuo et al. (2026) introduced Local Linear Attention (LLA), which drew on TTR to interpolate between linear and softmax attention via local linear regression; however, like softmax attention, LLA still requires a memory store that scales as O(TD)O(TD) memory, whereas SwiLA uses a fixed state size like linear attention. To our knowledge, SwiLA is the first sequence layer to extend the TTR framework by specifically incorporating mixture modeling; though Wang et al. (2026) use a collection of TTR agents which communicate with each other via a fixed weight matrix. Efficient attention The main motivation for this paper is that standard softmax attention (Vaswani et al., 2017) is highly expressive, but can be computationally expensive: it suffers from a KV cache that grows with the sequence length. In contrast, linear attention mechanisms (Zhang et al., 2026) compress context to a fixed-dimensional state, and so are computationally cheaper—but this fixed state often results in too much compression of the context, limiting long-range modeling capabilities (Arora et al., 2024a; von Oswald et al., 2026; Tandon et al., 2025). This dilemma has been noted by many in the literature (Zhang et al., 2024b; Han et al., 2024a; Han et al., 2024b). One mechanism to increase the expressivity of recurrent architecture is state expansion, namely to make the fixed state-size bigger (Gu, 2025; Gelada et al., 2025). One perspective on the benefits of SwiLA is that the mixture of regressors is a form of state expansion. Another approach is to relax the constraint of fixed state size, and instead allow the state to grow with the sequence length, albeit slower than linearly (as the KV cache does). For example, Guo et al. (2026) introduce log-linear attention, a sequence layer with state size that grows logarithmically in the sequence length. Finally, an orthogonal approach to efficient attention is to optimize the software implementations of attention (Dao et al., 2022; Zadouri et al., 2026) or linear attention (Yang and Zhang, 2024; Beck et al., 2025), with the goal of making them as performant as possible on GPUs. Such approaches also include optimizing the software implementations of sequential evaluations of nonlinear RNNs (Pöppel et al., 2025; Mishra et al., 2026). Appendix D Chunkwise Parallel Algorithms DeltaNet admits a chunkwise parallel algorithm (Yang et al., 2024b; Yang, 2024) that parallelizes its linear recurrence (4.1) over the sequence length. The SwiLA recurrence (3.3) is nonlinear in its hidden state, which makes it harder to parallelize. Nevertheless, SwiLA can still be parallelized, at least theoretically, using parallel Newton iterations (Lim et al., 2024; Gonzalez et al., 2024; Danieli et al., 2026; Gonzalez, 2026). However, most previous implementations of parallel Newton iterations use the parallel associative scan (Blelloch, 1990), colloquially known as the pscan. A pscan-based implementation would materialize the full state at every time step, which scales very badly for matrix-valued hidden states (Yang, 2024). One of the novel contributions of this paper is that we show that we can parallelize the nonlinear SwiLA recurrence with parallel Newton iterations; but where every Newton iteration is parallelized with the chunkwise parallel algorithm used for DeltaNet (Yang et al., 2024b; Yang, 2024). Nonetheless, in all our experiments, we use an optimized recurrent Triton kernel for SwiLA (i.e. we evaluate the nonlinear recurrence sequentially). The purpose of this appendix is to show that a chunkwise parallel implementation is also available in principle. An optimized implementation of this novel chunkwise parallel algorithm for nonlinear recurrences is an interesting direction for future work. D.1 Background on Parallelizing Recurrences To develop our novel chunkwise parallel algorithm for nonlinear recurrences, we combine two ideas: • the chunkwise parallel algorithm for DeltaNet (Yang et al., 2024b); and • parallel Newton iterations for nonlinear recurrences (Lim et al., 2024; Danieli et al., 2026; Gonzalez, 2026) DeltaNet’s chunkwise parallel algorithm Consider a linear recurrence of the form Wt=AtWt−1+bt. W_t=A_t\,W_t-1+b_t. When the transition matrices AtA_t are unstructured, parallel evaluation via associative scans (Blelloch, 1990) requires materializing and multiplying dense matrices, which can be prohibitively costly. Yang et al. (2024b) observe that when AtA_t has the structure of a rank-one perturbation of the identity At=I−βtktkt⊤, A_t=I- _t\,k_tk_t , (D.1) the recurrence admits an efficient ”chunkwise” algorithm. This algorithm gets its name because the sequence is divided into ”chunks” of size C. Within each chunk, the dependency among timesteps is captured by a strictly lower triangular matrix Λ∈ℝC×C ^C× C with entries Λs,t=−βsks⊤ktif s>t,0otherwise, _s,t= cases- _s\,k_s k_t&if s>t,\\ 0&otherwise, cases (D.2) and intra-chunk states are recovered via forward substitution through this triangular dependency structure. Boundary states are propagated sequentially across the N≔T/CN T/C chunk boundaries, but this cost is amortized over the chunk size. The algorithm avoids materializing full D×D× D transition matrices, operating instead on C×C× C matrices of key inner products, and is well suited to modern GPU hardware. Parallel Newton iterations for nonlinear recurrences Many sequence models involve nonlinear state updates wt=ft(wt−1)w_t=f_t(w_t-1), which seem like they must be evaluated sequentially. Recent work (Song et al., 2021; Danieli et al., 2023; Lim et al., 2024; Gonzalez et al., 2025; Gonzalez, 2026) shows that such recurrences can be parallelized by recasting sequence evaluation as a fixed-point problem over the full trajectory. Starting from an initial guess for the entire trajectory, often denoted by w(0)1:Tw^(0)_1:T, fixed-point methods iteratively converge to the true nonlinear rollout w1:Tw_1:T. A broad class of methods, including Newton, quasi-Newton, and Picard iterations, can be written in the unified form wt(i+1)=ft(wt−1(i))+At(i)(wt−1(i+1)−wt−1(i)), w^(i+1)_t=f_t\! (w^(i)_t-1 )+A_t^(i)\! (w^(i+1)_t-1-w^(i)_t-1 ), (D.3) where At(i)A_t^(i) approximates the Jacobian ∂ft/∂wt−1(wt−1(i))∂ f_t/∂ w_t-1(w_t-1^(i)). Because each iteration defines a linear dynamical system, the full sequence can be evaluated using a parallel scan in (logT)O( T) depth (Blelloch, 1990; Martin and Cundy, 2018; Smith et al., 2023). Moreover, these fixed-point iterations are guaranteed to converge in at most T iterations (Bellen and Zennaro, 1989; Shih et al., 2023; Tang et al., 2024; Gonzalez et al., 2024). Our core insight in combining these two approaches is that if Jacobian approximation At(i)A_t^(i) in the parallel Newton iterations in eq. D.3 has the identity-minus-rank-one form of the delta rule eq. D.1, then each fixed-point step can be parallelized with DeltaNet’s chunkwise parallel algorithm. D.2 Linearizing the SwiLA Recurrence To apply parallel Newton iterations to SwiLA, we must first linearize its recurrence (3.3). For each mixture component j and output dimension d, we can write the SwiLA recurrence as wtjd=f(wt−1,jd)≔wt−1,jd+βtjdrtjdδtjdkt.w_tjd=f(w_t-1,jd) w_t-1,jd+ _tjd\,r_tjd\, _tjd\,k_t. To linearize this recurrence, we start by computing the Jacobians ∂f∂wt−1,jd∈ℝD×D ∂ f∂ w_t-1,jd ^D× D. We start this process with some intermediate computations using shorthands defined in Section B.1, obtaining ∂δtjd∂wt−1,jd ∂ _tjd∂ w_t-1,jd =−kt⊤ =-k_t ∂rtjd∂wt−1,jd ∂ r_tjd∂ w_t-1,jd =Z∂πj/∂wt−1,jd−πj∂Z/∂wt−1,jdZ2 = Z ∂ _j∂ w_t-1,jd- _j ∂ Z∂ w_t-1,jdZ^2 =(Z−πj)(πjδtjdkt⊤)Z2 = (Z- _j)( _j _tjdk_t )Z^2 =(1−rtjd)rtjdδtjdkt⊤ =(1-r_tjd)r_tjd _tjdk_t Using these derivatives and the product rule, it follows that ∂rtjdδtjd∂wt−1,jd ∂ r_tjd _tjd∂ w_t-1,jd =rtjd∂δtjd∂wt−1,jd+δtjd∂rtjd∂wt−1,jd =r_tjd ∂ _tjd∂ w_t-1,jd+ _tjd ∂ r_tjd∂ w_t-1,jd =−rtjdkt⊤+rtjd(1−rtjd)δtjd2kt⊤ =-r_tjdk_t +r_tjd(1-r_tjd) _tjd^2k_t =−rtjd(1−δtjd2(1−rtjd))kt⊤. =-r_tjd (1- _tjd^2(1-r_tjd) )k_t . Thus, it follows that ∂f∂wt−1,jd=I−gtjdktkt⊤, ∂ f∂ w_t-1,jd=I-g_tjd\,k_tk_t , (D.4) where gtjd=βtjdrtjd[1−δtjd2(1−rtjd)]. g_tjd= _tjd\,r_tjd [1- _tjd^2(1-r_tjd) ]. (D.5) Thus, the Jacobian of the per-expert, per output dimension SwiLA recurrence is a rank-one perturbation of the identity, which is exactly the structure exploited by the chunkwise algorithm of Yang et al. (2024b). Applying the Newton linearization, the i-th iteration solves the linear recurrence wtjd(i+1)=(I−gtjd(i)ktkt⊤)wt−1,jd(i+1)+ktptjd(i), w^(i+1)_tjd= (I-g^(i)_tjd\,k_tk_t )\,w^(i+1)_t-1,jd+k_t\,p^(i)_tjd, (D.6) where the effective input term is ptjd(i):=gtjd(i)kt⊤wt−1,jd(i)+βtjdrtjd(i)δtjd(i). p^(i)_tjd:=g^(i)_tjd\,k_t w^(i)_t-1,jd+ _tjd\,r^(i)_tjd\,δ^(i)_tjd. (D.7) This effective input term follows because we know from eq. D.3 that bias term btb_t takes the form ft(wt−1(i))−Atwt−1(i)f_t(w_t-1^(i))-A_tw_t-1^(i). In this setting, At=I−gtjd(i)ktkt⊤A_t=I-g_tjd^(i)k_tk_t , and so the bias term is given by bt(i) b_t^(i) =wt−1,jd(i)+βtjdrtjd(i)δtjd(i)kt−(I−gtjd(i)ktkt⊤)wt−1(i) =w_t-1,jd^(i)+ _tjdr_tjd^(i) _tjd^(i)k_t- (I-g_tjd^(i)k_tk_t )w_t-1^(i) =kt(gtjd(i)kt⊤wt−1(i)+βtjdrtjd(i)δtjd(i)).■ =k_t (g_tjd^(i)k_t w_t-1^(i)+ _tjdr_tjd^(i) _tjd^(i) ). Because eq. (D.6) is a linear recurrence with identity-minus-rank-one transitions, each Newton iteration can be solved efficiently using the chunkwise algorithm. Note, however, that this is a quasi-Newton iteration because these Jacobians are only approximate: they ignore the cross mixture component interactions coming from the denominator of the responsibilities in eq. 3.4. However, Gonzalez et al. (2024) proves that quasi-Newton iterations are still guaranteed to converge globally, regardless of the form of the Jacobian approximation. A limitation of quasi-Newton iterations, however, is that their backwards pass is only approximate (Gonzalez, 2026). We can surmount this by running fixed-point iterations in the backwards pass as well; or, alternatively, by modifying the SwiLA recurrence so that the Jacobians above are exact (for example, by using ”unnormalized” responsibilities in eq. 3.4, i.e. removing the denominators, which is the only source of cross mixture component interaction). We have attempted these approaches: in configurations inspired by the language modeling pretraining settings (batch size of 16, sequence length T=4096T=4096, 8 heads, D=64D=64, and J=4J=4 mixture components), these various approaches converge in around 10 Newton iterations with synthetic inputs. However, despite this rapid convergence in Newton iterations, we so far did not optimize the parallel Triton kernels efficiently, so we were effectively the same speed as our sequential Triton kernel. Future work should be able to further accelerate such chunkwise parallel approaches for nonlinear recurrences. D.3 Chunkwise Parallel Algorithm To maximize hardware utilization, we implement the linearized recurrence using a chunkwise parallel algorithm following Yang et al. (2024b). We try to follow the notation in Yang (2024). In particular, we divide the sequence length into chunks of size C. We use [n][n] to indicate boundary states and □[n]c=□nC+c _[n]^c= _nC+c to denote states within a chunk. Within each Newton iteration i, the computation proceeds in three steps. 1. Intra-chunk precomputation. For all timesteps t in parallel, we compute the Jacobian coefficients gtjd(i)g^(i)_tjd and effective inputs ptjd(i)p^(i)_tjd, and form the strictly lower triangular dependency matrix Λ[n],j(i)∈ℝC×C ^(i)_[n],j ^C× C with entries (Λ[n],j(i))s,t=−g[n],j(i),s(k[n]s)⊤k[n]tif s>t,0otherwise. ( ^(i)_[n],j )_s,t= cases-g^(i),s_[n],j\,(k^s_[n]) k^t_[n]&if s>t,\\ 0&otherwise. cases (D.8) 2. Boundary state resolution. For each chunk n, we compute the aggregate transition and history matrices via forward substitution through the unit lower triangular system (I−Λ[n],j(i))(I- ^(i)_[n],j): M[n],j(i) M^(i)_[n],j =(I−Λ[n],j(i))−1diag(g[n],j(i))K[n], =(I- ^(i)_[n],j)^-1\,diag(g^(i)_[n],j)\,K_[n], (D.9) H[n],j(i) H^(i)_[n],j =(I−Λ[n],j(i))−1vec(p[n],j(i),1,…,p[n],j(i),C). =(I- ^(i)_[n],j)^-1\,vec(p^(i),1_[n],j,…,p^(i),C_[n],j). (D.10) These aggregates allow the boundary states to be updated sequentially across chunks: w[n+1],j(i+1)=w[n],j(i+1)+K[n]⊤(H[n],j(i)−M[n],j(i)w[n],j(i+1)). w^(i+1)_[n+1],j=w^(i+1)_[n],j+K _[n] (H^(i)_[n],j-M^(i)_[n],j\,w^(i+1)_[n],j ). (D.11) 3. Parallel output generation. Given the boundary states, the outputs within each chunk are computed in parallel: O[n],j(i+1)=K[n]w[n],j(i+1)+(K[n]K[n]⊤⊙L)(H[n],j(i)−M[n],j(i)w[n],j(i+1)), O^(i+1)_[n],j=K_[n]\,w^(i+1)_[n],j+ (K_[n]K _[n] L ) (H^(i)_[n],j-M^(i)_[n],j\,w^(i+1)_[n],j ), (D.12) where L∈ℝC×CL ^C× C is a strictly lower triangular causal mask ensuring autoregressive validity. Appendix E Experimental Details For all experiments, we heavily utilized architecture implementations in Yang and Zhang (2024). E.1 Approximating Softmax Attention Figure S1: Approximating Softmax Attention. Test R2R^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. Each model receives input sequences x∈ℝT×Dx ^T× D, where T=512T=512 and D=32D=32, and must approximate the output of causal softmax self-attention with fixed, ground-truth projections θq,θk,θv _q, _k, _v. Models receive only x and must learn their own Q/K/V projections. Data Generation At each sequence position, a cluster index is sampled uniformly from C=32C=32 clusters. The input is xt=L2normalize(μc+0.5ϵt)x_t=L2normalize( _c+0.5\, _t), where μc∼(0,I32) _c (0,I_32) are fixed cluster centers and ϵt∼(0,I32) _t (0,I_32). The query, key, and value projections θq,θk,θv∈ℝ32×32 _q, _k, _v ^32× 32 are random matrices, with each entry sampled from a standard normal distribution. The attention temperature is τ=D/2τ= D/2, and the sequence length is T=512T=512. We generate 65,536 training, 2,048 validation, and 2,048 test sequences. Models. • MHA. Standard causal multi-head softmax attention with learned Wq,Wk,WvW_q,W_k,W_v (bias enabled) and output projection WoW_o (no bias). Scaling factor 1/dk1/ d_k. • DeltaNet. Learned projections with SiLU activation followed by L2 normalization on Q/K. SiLU activation on V. Scalar β per head via a linear layer ℝdin→ℝHR^d_in ^H followed by sigmoid. RMSNorm on the output, then an output projection. • SwiLA. Same Q/K/V projection scheme as DeltaNet (SiLU + L2 normalization on Q/K, SiLU on V). Separate β per head, output dimension, and mixture component. Routing logits are computed via SwiGLU projections (inner dimension set to 256) with softmax activation. RMSNorm on the output, then an output projection. MI auxiliary loss coefficient λ=0.001λ=0.001. We do not employ input-dependent noise injection. E.2 Synthetic Benchmarks Context-Dependent Associative Recall (CDAR) Context-Dependent Associative Recall (CDAR) is a synthetic sequence modeling task that extends Multi-Query Associative Recall (MQAR) from Arora et al. (2024a) to multiple contexts. Each input sequence has length 1024 and consists of two zones: a context zone containing C=4C=4 distinct contexts, each defining 32 unique key–value pairs (128 total associations), followed by a query zone in which the model must retrieve the correct value for a given (context, key) pair. The vocabulary size is 8192. Query positions within the query zone are sampled according to a power-law distribution p(i)∝ia−1p(i) i^\,a-1 with exponent a=0.01a=0.01, following the Zoology benchmark setup (Arora et al., 2024a). We generate 100,000 training examples and 3,000 test examples. All queries for each context are grouped contiguously, context tokens are not repeated within query groups, and non-query positions are filled with zeros. Training uses a batch size of 64. All models use a GPT-NeoX-style pre-norm block structure following the Zoology framework (Arora et al., 2024a): each block consists of a sequence mixer followed by a state mixer (MLP, hidden multiplier 4), with LayerNorm, residual connections, and dropout on each sub-layer. The sequence mixer is the model-specific layer (Temporal SwiLA, DeltaNet, Softmax Attention, or MoM), while the state mixer is shared across all models. All models use a short convolution (kernel size 2) before the sequence mixer. We detail the hyperparameters that differ across models: • Temporal SwiLA. dmodel∈32,64,128d_model∈\32,64,128\, nlayers=1n_layers=1, H=1H=1 head, K=4K=4 mixture components, and key/value expansion factors ek=ev=0.5e_k=e_v=0.5. We used SwiGLU projections for π and β, with inner dimensions 256 and 32, respectively, and input-dependent noise injection. We set the MI auxiliary loss coefficient λ=0.0λ=0.0. • SwiLA. dmodel∈32,64,128,256d_model∈\32,64,128,256\, nlayers∈1,2n_layers∈\1,2\, H=1H=1 head, K=4K=4 mixture components, and key/value expansion factors ek=ev=0.5e_k=e_v=0.5. We used SwiGLU projections for π and β, with inner dimensions 256 and 32, respectively, and input-dependent noise injection. We set the MI auxiliary loss coefficient λ=0.0λ=0.0. • DeltaNet. dmodel∈32,64,128,256d_model∈\32,64,128,256\, nlayers∈1,2n_layers∈\1,2\, H=4H=4 heads, ek=ev=2.0e_k=e_v=2.0. • Softmax Attention. dmodel∈32,64,128,256d_model∈\32,64,128,256\, nlayers∈1,2n_layers∈\1,2\, H=4H=4 heads, RoPE positional encoding with θ=10,000θ=10,000. • DeltaNet-MoM. dmodel∈32,64,128,256d_model∈\32,64,128,256\, nlayers∈1,2n_layers∈\1,2\, M=4M=4 memories with top-k=2k=2 routing, H=1H=1 head, ev=1.0e_v=1.0. Figure S2: Context-Depedent Associative Recall. CDAR with 4 contexts and 32 key-value pairs per context: test accuracy vs. state size. All models are trained with AdamW and a grid search over four learning rates log-spaced from 3.16×10−43.16× 10^-4 to 10−210^-2, i.e. 3.16×10−4, 10−3, 3.16×10−3, 10−2\3.16× 10^-4,\;10^-3,\;3.16× 10^-3,\;10^-2\. All models are trained for 50 epochs. For each configuration and learning rate, we run three random seeds. We report the maximum accuracy across the learning-rate sweep and three random seeds in Figure 2A because learning on these algorithmic tasks typically exhibits sharp, step-like transitions, and our primary objective is to assess whether an architecture can learn the target algorithm rather than its optimization speed or robustness. For completeness, we report the mean test accuracy across learning rates in Figure S2, which yields the same qualitative conclusions, with the largest improvements obtained by Temporal SwiLA. RegBench: In-Context Language Learning All models have 8 layers, dmodel=256d_model=256, MLP inner dimension 1,0241,024 (4×dmodel4× d_model), embedding dropout 0.1, and no residual dropout. SwiLA, Temporal SwiLA, DN-MoM, and DeltaNet all prepend a short convolution of width 4. We compare: SwiLA (2 heads, 2 mixtures, expand_k=expand_v=0.5 expand\_k= expand\_v=0.5); Temporal SwiLA (same as SwiLA with a temporal gate); DN-MoM (1 head, 4 memories, top-2 routing, dhead=64d_head=64, expand_v=1.0 expand\_v=1.0, DeltaNet backend); DeltaNet (4 heads, expand_k=expand_v=1.0 expand\_k= expand\_v=1.0); and Transformer (4 heads with RoPE). The recurrent state size is matched across SwiLA, DN-MoM, and DeltaNet at 16,384=4×64×6416,384=4× 64× 64 per layer. For both SwiLA and Temporal SwiLA, we do not use input-dependent noise injection and set the MI loss coefficient to 0.005. We use AdamW with learning rate 2.5×10−42.5× 10^-4 and weight decay 0.10.1. The learning rate follows a cosine schedule with 10% linear warmup (from 10−610^-6) and minimum learning rate 2.5×10−52.5× 10^-5. Batch size is 32 and models are trained for 200 epochs. We report test accuracy using the checkpoint with the lowest validation loss. E.3 Pretraining Language Models All models are pretrained on a subset of the FineWeb-Edu dataset (Penedo et al., 2024) (sample-100BT subset) using the 32K-Llama2 tokenizer. Training sequences have length 4,0964,096. All models are trained with identical optimization hyperparameters chosen following Yang et al. (2024b). We use the AdamW optimizer with a learning rate of 3×10−43× 10^-4 and weight decay of 0.010.01. The learning rate follows a cosine decay schedule with 954 warmup steps and a minimum learning rate ratio of 0.10.1 (i.e., final learning rate 3×10−53× 10^-5). Gradient norms are clipped to 1.01.0. Training runs for 28,610 steps with a global batch size of 128128 sequences, totaling approximately 15B tokens. Training is conducted on a node of 8 NVIDIA H100 GPUs. All thirteen models—Transformer++ (Touvron et al., 2023), DeltaNet (Yang et al., 2024b), GDN (Yang et al., 2025), two MoM variants (DeltaNet-MoM and GDN-MoM) (Du et al., 2026), KDA (Team et al., 2025), FoX (Lin et al., 2025), SwiLA, Temporal SwiLA, Gated SwiLA, Gated Temporal SwiLA, Hybrid GDN, and Hybrid Temporal SwiLA—share the same architecture skeleton: 24 layers and a hidden dimension of 1,024, targeting approximately 374M total parameters. Feed-forward intermediate dimensions are adjusted slightly across architectures to match the total parameter count. Following the default configurations in Zhang and Yang (2025), we use 16 heads with head dimension 64 for Transformer++. DeltaNet, GDN, KDA, and FoX use 8 heads with head dimensions of 128. SwiLA, Temporal SwiLA, Gated SwiLA, and Gated Temporal SwiLA use 8 heads with 4 mixtures per head and 64-dimensional key and value projections. This gives a recurrent matrix state of 8×4×64×64=131,0728× 4× 64× 64=131,072 scalars per layer, matching the 8×128×128=131,0728× 128× 128=131,072 state scalars used by DeltaNet, GDN, and KDA. For all SwiLA variants, we disable input-dependent router and temporal gate noise. We set the load balancing loss coefficients to 0.0010.001 for SwiLA, 0.00050.0005 for Temporal SwiLA, 0.00050.0005 for Gated SwiLA, and 0.000250.00025 for Gated Temporal SwiLA. For Gated SwiLA variants, we tie the precision across the output index, i.e., one gate per key dimension, as in KDA (Team et al., 2025). Hybrid Temporal SwiLA uses the same 0.00050.0005 coefficient as Temporal SwiLA. Hybrid GDN and Hybrid Temporal SwiLA contain six full causal-attention layers at layers 4, 8, 12, 16, 20, and 24, with the remaining 18 layers using GDN and Temporal SwiLA, respectively. The full-attention layers use 8 heads with head dimension 128. For GDN-MoM, we use GDN for the recurrent memory updates and employ three routed memories with top-k=2k=2, together with one shared memory. Each memory uses 8 heads with head dimension 64, giving a total main recurrent state size of 4×8×64×64=131,0724× 8× 64× 64=131,072 scalars per layer. For DeltaNet-MoM, we trained the following three architectural configurations and report the one with the best average performance on the commonsense benchmark (best configuration in bold): 1. 374.1M parameters: 4 heads, head dimension 64, 7 routed memories with top-k=4k=4, one shared memory, and an MLP intermediate dimension of 2,634. 2. 374.2M parameters: 1 head, head dimension 128, 7 routed memories with top-k=4k=4, one shared memory, and an MLP intermediate dimension of 3,412. 3. 374.1M parameters: 2 heads, head dimension 128, 3 routed memories with top-k=2k=2, one shared memory, and an MLP intermediate dimension of 3,326. Each of the three configurations maintains 131,072 main recurrent-state scalars per layer. Following Du et al. (2026), we use shared memory and adjust the MLP intermediate dimension to match the parameter counts of the other models. E.4 Throughput Comparison Training throughput and memory. 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. We benchmark complete training updates on a single NVIDIA H100 for the six 1.3B models: Transformer, Gated DeltaNet (GDN), Hybrid GDN, SwiLA, Temporal SwiLA, and Hybrid Temporal SwiLA. At sequence lengths 2,048, 4,096, 8,192, and 16,384 we use microbatch sizes 64, 32, 16, and 8, respectively, holding the workload fixed at 131,072 tokens per update. A timed update comprises the forward pass and loss computation, the backward pass, global gradient-norm clipping, and an AdamW step. For each configuration, three warm-up updates are followed by three timed updates with CUDA synchronization. Throughput is 131,072 divided by the mean wall-clock update time, reported in thousands of tokens per second, and peak memory is PyTorch’s maximum allocated CUDA memory after warm-up. Generation throughput and memory. 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. We measure end-to-end autoregressive generation for the same six models on a single NVIDIA H100. Each input is a synthetic 2,048-token prompt extended by 128 generated tokens using greedy decoding. We sweep batch sizes from 1 to 128 in powers of two. For each model, batch sizes beyond the first out-of-memory failure are omitted. One warm-up call is followed by three timed calls to the full generation routine, with CUDA synchronized around the timed loop. Reported values are means over the three calls, where the wall-clock time of a generation call includes both prefill and decoding. Peak memory is PyTorch’s maximum allocated CUDA memory after warm-up. E.5 Ablation Study We conducted ablation experiments that isolate the contribution of the number of mixture components, temporal persistence, and the load balancing loss. We trained variants of 115M SwiLA on 5B tokens and evaluated test perplexity on 100M tokens, reported in the legend of Figure S5. To analyze routing dynamics, we tracked three statistics during training, computed separately for the key-value-side responsibilities and the query-side responsibilities. Utilization entropy measures the entropy of the marginal expert usage distribution, averaged over batch and sequence positions; it equals 1 when experts are used uniformly and is lower when usage is concentrated on a few experts. Routing entropy measures the average per-token entropy of the routing distribution; it equals 1 when each token spreads its responsibility uniformly across all experts and approaches 0 when each token commits to a single expert. Dead expert fraction counts the fraction of experts whose marginal usage falls below 10% of uniform usage (i.e., below 0.1/J0.1/J). Without load balancing, vanilla SwiLA suffers significant expert collapse: KV-side utilization entropy drops sharply, and the dead expert fraction exceeds 40%. Adding the load-balancing auxiliary loss resolves this, maintaining high utilization entropy and near-zero dead experts throughout training. Temporal SwiLA is notably more robust, maintaining reasonable utilization even without load balancing. Across all configurations except for the vanilla SwiLA without load balancing, routing entropy stabilizes in the 0.4–0.6 range, indicating moderately selective but not fully hard routing. 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 (4J) 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.