Local Linear Attention: An Optimal Interpolation of Linear and Softmax Attention For Test-Time Regression
Authors: Yifei Zuo
Yutong Yin
Zhichen Zeng
Ang Li
Banghua Zhu
Zhaoran Wang
Venue: ICLR 2026 Poster (2026)
Keywords: Attention Mechanism; Sequence Modeling; Test-Time Training; Local Linear Regression; Associative Memory; Hardware-Efficient Attention
Primary area: foundation or frontier models, including LLMs

Abstract: Transformer architectures have achieved remarkable success in various domains. While efficient alternatives to Softmax Attention have been widely studied, the search for more expressive mechanisms grounded in theoretical insight—even at greater computational cost—has been relatively underexplored. In this work, we bridge this gap by proposing Local Linear Attention (LLA), a novel attention mechanism derived from nonparametric statistics through the lens of test-time regression. First, we show that LLA offers theoretical advantages over Linear and Softmax Attention for associative memory via a bias-variance trade-off analysis. Next, we address its computational challenges and propose two memory-efficient primitives to tackle the $\Theta(n^2d)$ and $\Theta(nd^2)$ complexity. We then introduce {FlashLLA}, a hardware-efficient, blockwise algorithm that enables scalable and parallel computation on modern accelerators. In addition, we implement and profile a customized inference kernel that significantly reduces memory overheads. Finally, we empirically validate the advantages and limitations of LLA on test-time regression, in-context regression, associative recall and state tracking tasks. Experiment results demonstrate that LLA effectively adapts to non-stationarity, outperforming strong baselines in test-time training and in-context learning, and exhibiting promising evidence for its scalability and applicability in large-scale models.

Published as a conference paper at ICLR 2026
LOCAL LINEAR ATTENTION: AN OPTIMAL INTERPO- LATION OF LINEAR AND SOFTMAX ATTENTION FOR BEST-TIME REGRESSION
Yifei Zuo
Yutong Yin
Northwestern University
Northwestern University
Evanston, IL 60208, USA yi fei Zue229920 northwestern.ed
Evanston, IL 60208, USA
yutongyin2028euu northwestern..dd
Zhicheng Zeng University of Washington Seattle, WA 98195, USA
Ang Li
Banghua Zhu
University of Washington Seattle, WA 98195, USA anglizQuw.en
University of Washington Seattle, WA 98195, USA banghuaß uw. edu
Sczeng@uw.edu
Zhaoran Wang
Northwestern University Evanston, IL 60208, USA
zhaoranwang amauiigg@@maii.mmmmmm com
ABSTRACT
Transformer architectures have achieved remarkable success in various domains. While efficient alternatives to Softmax Attention have been widely studied, the search for more expressive mechanisms grounded in theoretical insight-even at greater computational cost has been relatively underexplored. In this work, we bridge this gap by proposing Local Linear Attention (LLA), novel attention mechanism derived from nonparametric statistics through the lens of test-time re- gression. First, show that L.LA offers theoretical advantages over Linear and Softmax Attention for associative memory via a bias-variance trade-off analysis. Next, we address its computational challenges and propose two memory-efficiet primitives to tackle the (n²d) and ((nd²) complexity. We then introduce Flash- LLA, a hardware-efficieer
blockwise algorithm that enables scalable and paral- lel computation on modern accelerators. In addition, we implement and profile a customized inference kernel that significantly reduces memory overheads. Finally, we empirically validate the advantages and limitations of L.IA on test-time regres- sion, in-context regression, associative recall and state tracking tasks. Experiment results demonstrate that LLA effectively adapts to non-stationarity, outperform- ing strong baselines in test-time training and in-context learning, and exhibiting promising evidence for its scalability and applicability in large-scale models.
INTRODUCTION
baasee architectures have dominated modern AI systems and powered breakthroughs across various fields. The core computation primitive-Softme Attention (Vaswani et al., (017), adaptively processes and aggregates contextual information. Recent research on architecture inno- vation has proposed numerous variants of attention mechanism such Linear Attention (LA) (Yang)
Transformer-
ct al. | 2024b 2025 Siems et al. 2025 Sun al. 2023 Ma et al. 20244] Liu et al.
2024 Katharopou- los et al. 2020 Yang et al. 2024a) and state space models (SSMs) (Gu et al. 2022a Gu & Dao
2024 Dao & Gu 2024 Poli et al. 2202).. While these methods offer remarkable efficiency improve-
ments long sequences, they often incur performance penalty compared to Softmax Attention
(Bick et al., [205 Katharopoulos et al. 2020). Meanwhile, many of the design choices such as
gating and forgetting factor Yang al.
2024a 2025
Lin et al. 2025
Gers al. 2000) are often

---

Published as a conference paper at ICLR 2026
guided by heuristics empirical results, lacking principled understanding. Conversely, the search for more expressive attention mechanisms, even at an additional computational cost, has been rela- tively under-explored. Recently, test-time regression (Wang et al. 2025) unifies the design choices of different attention variants, indicating that the attention mechanism can be viewed as test-time optimizer for layer-specific regression problems. A natural question arises: can systematically improve the Sofimax Attention mechanism from the perspective of test-time regression?
In this work, we propose local linear attention (LLA), an upgrade Softmax Attention derived from local linear regression. Our contributions are summarized as follows:
We systematically analyze the design space of attention mechanisms within the test-time regression framework and propose LLA. We provide theoretical comparison of LLA with Softmax Attention and LA family from the perspectives of bias-variance trade-off and demonstrate its provable advantage in associative recall capability.
We provide a detailed discussion on the computation of LLA and overcome the ((n2d) and O(nd²) memory complexity with two optimizations, where is the sequence length and is the dimension. We then introduce FlashLLA, a blockwise and hardware-efficien algorithm to parallelize the computation on modern accelerators. In addition, we implement and profile customized inference kernel with significant memory reduction.
We conduct extensive experiments on synthetic tasks including test-time regression, in- context regression, associative recall and state tracking to validate the advantages and lim- itations of LLA compared to strong baselines.
1.1 RELATED WORKS.
Linear Attention and State Space Models. Due to the quadratic computational cost and linear memory consumption of the softmax attention mechanism for autoregressive sequence modeling, efficient attention mechanisms such as LA (Yang et al.) 2024b 2025; Siems et al., 2025 Sun et al.,
2023
Ma et al. 2024 Liu al.
Yang al., 2024a) and SSMs (Gu al. 2022a; Gu & 2024 Poli et al. 2023) were proposed in search of more efficient alterna-
2024
Dao
2024; Dao
Gu
lives for long-context sequence generation. These methods maintain constant sized hidden state 2020) during decoding and update it like a linear RNN. This state update
KKaharopoulos et al.
behavior also referred fest-w programming (Schmidhuber 1992 Schlag et al.) 202T). Es- sentially, MesaNet (von Oswald et al. 2025) is LA variant that preconditions the hidden state and achieves optimal regression objectives among linear models.
In-Context Learning by Optimization. growing body of work suggests that attention mecha- nism implicitly performs optimization algorithms to achieve in-context learning (Garg et al. 2022
Akyürek et al. 2023 von Oswald et al. 2023 Kirsch et al., 20244 Zhang et al., [2024] Mahankali
et al. 2024 Ahn et al.) 2023] Dai et al. 2023). For example, Mesa optimization (von Oswald
et al., 2023) suggested that the attention layer inherently performs or approximates optimization
steps during the forward pass. This behavior is particularly evident in LA and SSMs, as the hidden state update can be interpreted as performing gradient descent to solve a linear regression objective (Wang et al. 2025 Liu et al.) 2024). MesaNet (von Oswald et al.) 2025) is one-step convergent algorithm for such a problem as the objective accepts a closed-form solution.
Hardware Efficient Attention. Prior works have focused on efficient attention implementation on modern hardwares to alleviate memory and computation overheads. FlashAttention (Dao et al.) 2022 Dao 2024) performs a block-wise online softmax to reduce I/O latency. Several other ap- proaches, including NSA (Yuan et al. 2025), SeerAttention (Gao et al. 2024). MoBA (Lu et al.) 2025). and Block Sparse Attention (Xiao) 2202), leverage sparsity to lower the effective computa-
tional cost while preserving GPU utilization. Flash Linear Attention (Yang et al. (024a) provides
hardware- -freandd formulation of linear attention through chunk-wise computation.
1.2 NOTATION.
We use upper-case letters to denote matrices and lower-case case letters to denote vectors. For a matrix X, denote its Frobenius norm as ||||| and the Hadamard product @. For vector Jr, denote its
2

---

Published as a conference paper at ICLR 2026
Euclidean norm as ||2||22 Furthermore, define rsum(X) X1 for matrix X and bcast(() x1 for vector z, where a vector of ones. We use the abbreviation brs brsum(r) = bcast (rsum())))()))
2 BEYOND LOCAL CONSTANT ESTIMATE
In this section, we first revisit the test-time regression interpretation for the attention mechanism (Wang et al. 2025). Then, we analyze the associative recall capacity and show the inherent limita- tions of LA and Softmax Attention. Lastly, introduce the formulation for LLA.
2.1 ATTENTION AS TEST-TIME REGRESSION
In test-time regression framework, the attention mechanism is interpreted as layer-specific re- gression solver. The goal is to approximate an unknown regression function f R + R using historical key-value pairs. To be specific, given a hypothesis space F and a position 1 I S n, an estimator E F is fitted on the dataset D; = {(k), E IRd )== where the attention keys kj serve as the features and attention values as the labels. The prediction is made at a query E IRd, which is treated as test data point.
Linear Attention as Parametric Regression. Parametric model constrains the function class F to a set of functions defined by a finite-dimensional parameter E O. The most fundamental instantiation is the linear regression, which sets {fo(z) Wx + b | 0 (W,b), W pdxd b E Rd.. We omit the intercept b for simplicity. At each position i, the parameter W2 is estimated by solving the following least square problem on the training dataset D,,
min C(WDD) =
pulluy - Wky + AIIWI
(1)
where Y/ E is weighting factor and X 2 is the ridge regularization penalty. For appropriate regularization, objective equation [ admits closed-form optimal solution. MesaNet (von Oswald
al. 2025 2024) hardcodes this solution for the case V/s where the prediction given by,
Bresa (g:) WMlsss li
Σ kikk
Σ kigk + +X1)
qr.
Across different position i, the weight WHess can be updated recurrently by maintaining two statis- ties S; and H, with (dd2) memory. This allows MesaNet to be interpreted as a linear RNN with two recurrent states. To avoid the expensive matrix inversion, vanilla LA brutally approximates the precondition matrix H, R in equation 2 leading to suboptimal solution of the least square problem. It can also be shown that LA variants and SSMs such as GLA (Yang et al. 2024a), RetNet (Sun et al., 2023). RWKV (Peng et al. 2024 2025) and Mamba (Gu & Dao 2024 Dao Gu 2024) can be derived by the approximation H; x with different weighting schemes Rrj. Besides exact solutions, vanilla LA and variants such as DeltaNet (Yang et al.) 2024b) and Gated DeltaNet (Yang let al.] 2025) can be interpreted performing one step of first order stochastic gradient descent weight W. We refer readers to (Wang et al. (2025) and tables in (Peng et al. 2025 Yang et al. 2025 for detailed derivations and how each model is implemented.
Softmax Attention is Non-Parametric. Non- Non-parametric regression makes minimal structural sumptions on the function class F. A canonical example is the kernel regression, where the estima- tor f is defined directly from the data in way that depends on the query point. Specifically, for query vector g:. let F(g)) denote the local function class around 9;- The local regression objective is defined as:
min L(f:D,) (EF(g)
=
Σ ==1
Wis 110 - f(k))11
(3)
3

---

Published as a conference paper at ICLR 2026
where wi = Kn k,) E R is a query-dependent weight that measures the locality of the training
point ky to the query qs. The simplest instantiation the constant model, where F(g:) = {fo (z) =
E IR Vx E ]. Solving objective equation] 3 with this function class yiclds:
f(aa)
8ij
Wij
(4)
Way'
Consider an RBF kernel Way exp(e 1| /h12/hh) with bandwidth = 2vd, the estimator equation 4 exactly recovers the softmax attention when QK normalization (Dehghani et al.) 2023
Wortsman et 2023 Team 2024) is applied, since in this case Wij exp(y ky Vd) and the common constant factor cancels is in the division. It is also known as the Badaraya-Watsod (NW) [964]Bierens) QK normalization is not strictly necessary to represent practical Softmax Attention as the additional term can naturally serve as positional encoding (Press et al. 2022), and the scale effectively tunes the bandwidth in data-dependent manner.
kernel regression (Nadaraya 1964, Watson 1964, Bierens 1988) statisties literature. We note that
2.2 LEARNING BEHAVIOR OF ASSOCIATIVE RECALL
Attention mechanisms are often evaluated by the associative memory capacity (Zhong al. 2025 Behrouz et al. 2026, Ramsauer et al., 202T). Specifically, given a training set of key-value pairs retrieve the value U associated with k, when queried at ky. This objective can be exactly captured by the Mean Square Error (MSE). For example, the retrieval error in vanilla LA is given by,
[[kg, "2) = the model expected
MSE Σ
Skkk 0/11==
Σ 11C - 1)03 + Σ ku 379
(5)
The first term in the summation is the signal bias and can be avoid by QK normalization. The second term is the interference from other key-value pairs. Deltaformer (Zhong al.) (2025) quantitatively analyze this eиror by the inverse of signal-to-noise ratio SNR which is essentially normalized version of MSE. In classic literature, MSE decomposition allows us to analyze the approximation and generalization error of the model and the corresponding bias- trade-off.
Irreducible Approximation Error of Global Linear Model. Recall that global linear models correspond to solving a least square problem equation over the hypothesis space F = {fe(z) = Wx+bj. When the ground truth function f is not global linear, any estimator f J will suffer from contrast, the local constant
a non-vanishing approximation error due to model misspecification. model is nonparametric estimator that does not impose structural assumptions the function class except for smoothness or regularity. Consequently, the approximation error vanishes asymptotically with proper assumptions. In fact, have the following separation result between global linear (GL.) and local constant (NW) estimators:
Proposition 2.1. Let (X,, Yi) be ii.dd, X, E Rdd supported on a bounded set D C Red, and Y, s(X,) + E Rdy with Ele, X,] 0 and Els² X.) = (X,)). Let JeL denote global- linear estimator and fow the local constant (NW) estimator with optimal bandwidth. Under mild assumptions, if is not globally linear, then
E 1 f(x))
E =/ [1fxw(() - f(ae)) dz O(n -3/(d+3)).
Irreducible Boundary Bias of Local Constant Model. Despite the appealing convergent prop- erty of local constant model, suffers when predicting near the boundary the data support, par- ticularly with symmetric kernels like RBF. This phenomenon becomes more pronounced in high- dimension and likely to occur more frequently in autoregressive prediction. Local polynomial re- gression is standard remedy in nonparametric statisties to address this issue. In Section 2.3 will introduce LLA as a natural adaptation local linear regression. In fact, we have the following separation result between local constant (NW) and local linear (LI.) estimators:
4

---

Published as a conference paper at ICLR 2026
Under the setting of Proposition 2.1 let fiw and fLL denote local-constant and local-lines linear estimators with their respective optimal bandwidths. Under mild assumptions, if has sufficiently large normal gradient along the boundary of D. then
Proposition 2.2.
E | -f))( 12(n 33/d433)) E Jf [ O(n 4/(4+44))
f(x)]1
The proof proposition 2.1 and 2.2 are provided in Appendix and respectively.
2.3 LOCAL LINEAR ATTENTION
For a query do instantiate the local function class F(a) (fe(x) b W(x a.) (W,6), W E gdxd b [.. The regularized local linear regression objective is,
Formulation.
min C(f:D.) (EFla.)
Σ wiy lle b W q:)l + Al
(6)
where Wyj query-dependent kernel weights and À 2 ridge penalty. This objective also admits a closed-form solution for the intercept b and weight W. Importantly, at test time, the predic- tion is only made at (4:)) = b.. Thus suffices to derive the formulation for the intercept. Define Zij kj qi and the following query-specific statistics,
Wi
Wii Ry
p4 Σ Wijza ER,
Σ, Σ
Σ wizzizy
Al Rdxd
Also denote Pı Di²tu E Rd. The optimal intercept can be computed as follows,
- z1,Pi
b = Σ My
81f Wij W; 14 Ps
(8)
Similar to maintaining H, in MesaNet, LLA requires capturing the precondition matrix E,, How- ever, a key difference is that H; is built on global statistics that are independent of the query, whereas Σ, is constructed from features centered around the specific query y for each position S n. Consequently, LLA requires KV cache of size ((nd) similar to Softmax Attention, rather than constant-size recurrent states as in the LA family.
LLA Interpolates Linear and Softmmx Attention. A more interpretable form of equation 8 can be obtained by decomposing the prediction into two components. Suppose that the weight matrix W; given prior solving equation 6. then the optimal intercept can expressed as.
6,
Σ (y Wyky) W.g.
Wij
8ij
== Way
The first term a local constant regression to predict the residuals my Wiky, while the second term linear prediction based on W. The formulation recovers LLA if W, is obtained by optimally solving equation6 However, by allowing suboptimal estimation, state similar to LA. This decomposition reveals how LLA interpolates between Linear and Softmax Attention and provides template for designing algorithms.
construct W, recurrent
3 PRACTICAL ALGORITHM
In this section, provide a detailed discussion of the computation involved in LLA. We high- light two major challenges in naive implementations and develop a practical block-wise algorithm FlashL.L. LLA that scales efficiently modern accelerators such GPUs.
5

---

Published as a conference paper at ICLR 2026
3.1 MEMORY EFFICIENT PRIMITIVES
Avoid Pairwise Materialization. The first bottleneck is the evaluation of vectors Zij ki qa for every 1 < j < i I n, which requires ((n²d) memory to materialize. This pairwise difference is later used in the formulation of μu and Σ; defined in equation 7 as well as the inner product Pi in equation 88. For both cases, explicit materialization of 2ij can be avoid by algebraically separating the contributions of ky and q the final result. Specifically, the statistics μ4; and Eg can be reformulated in terms of intermediate quantities that are independent of the query and can then be transformed to recover the original centered statisties:
p4 Kigky E GRd
Σ,
Σ Wiyhy AI E Ra
(10)
[4 P4 W.q =. N.g q:14: + w.q.
(11)
The computation in equation 10 and equation 11 oonly requires vectors kj and qi individually, reduc- ing the memory cost to ((nd). The same principle applies to computing inner products of the form P.j.; kj H2 q5 for any vector x; E RR. The matrix operator relative matrix multiplication (relmm) for this optimization as follows,
re mm X, Q, K) == XK - brsum( X @ Q)-
(12)
This operator is invoked once in the forward computation with x, = Ps+ but appears multiple times in the backward with other variables. Further details on the backward are provided in Appendix D
Matrix-Free Inversion via Conjugate Gradients. The second bottleneck arises in solving linear systems of the form E; -x, I; for some vector x; E Rd.. Directly inverting M, for every 1 < i < incurs a prohibitive ((nd)) memory. Following the approach in MesaNet (von Oswald et al. 2025); exploit the num-offrank-ood conjugate gradient (CG) method (Hestenes & Stiefel 1952). The key insight is that CG only evaluate the matrix-vector product E,p for a search direction p Rd without explicit matrix materialization:
structure of Σ, and solve the linear system iteratively using the
E,p = Σ wis (k] - (g p)A, (0, p)q: + (wig) p)q: + Xp.
(13)
j=1
Each term only involves inner products or weighted sums over keys and the query, both of which be computed efficiently using batched matrix multiplication with ((nd) memory. This CG operation of Σ; is invoked once in the forward computation with x; = and twice in the backward computation with other variables. Further details on the CG algorithm are provided in
3.2 PARALLEL FORM AND BLOCKWISE ALGORITHM
Matrix Formulation. We first express the key components of the LLA forward pass in matrix form. Let Q, K, V Raxd be the query, key, and value matrices respectively for a given layer and head. This function applies a causal mask to the input tensor using the trill operator that preserves the lowar-tria matrix. Then the output E Raxd can be computed as follows,
W trii(exp(QK)) /h)), M = WK K - braum (W)
(14)
R CGSolve (M, M,Q,K,KKK X),
8 rsuu (W) rsum [ O R)
(15)
0 - 1- relmm( IR, Q, K)
c W V,
(16)
(6)
where W Ruxn is the matrix kernel weight way. M goxd stores the first-order statisties μis R E Ruxd contains the solution to the linear systems Σ 14; for every 1 < i S n, and CGSolve(-) invokes the CG algorithm that construct Σ; implicitly and solve these systems in parallel. The division and subtraction are performed sementtwisee The single-head computation can be naturally extended to multi-head the same way in standard multi-head attention mechanisms.
Blockwise Algorithm. Denote B,, Be as the block size for queries and keys/values along the sequence length dimension and r, as the block index. Denote Q, E B, xd and K,, V. B as the block-wise representations. The forward pass of FlashLLA summarized in Algorithm

---

Published as a conference paper at ICLR 2026
Algorithm 1 FlashLLA Forward Pass
Require: Matrices Q, K, V in HBM, block sizes Br, Be, regularization X, bandwidth h.
1: Divide Q into [n/ B,1 blocks of size B, and K, V into Ba] blocks of size B
2: Divide output O, It into [n/B,] blocks of size Br.
3: for r = to mn/ do
4: Load Q, from HBM to SRAM.
5: Initialize on-chip: MM,") 0 R BB,dd w/ t OERRR, mr
(0) t 00 Rr
6: for to [n/ B.] do
7: Load Kg from HBM to SRAM.
Compute W = Q. h and m = max (c-1) rowmaxxW)).
8:
Compute a, exp(m, (e-1) m), W exp(W boast(m) (m)) and update m. (0)
9:
Compute w/ = at @ wr (0-1) rs)
10:
Compute M bcast (c) M (c-1) WKK-
11:
12:
end for
Initialize on-chip: 0 0 O R 0 Raxd
13:
Compute M, ( brsum W) 0 Qr.
14:
Compute R CGSolve(M,, Q,, K, M (last) (last)
15:
A).
Compute 6, = (2aast) rsum M, @ R.).
16:
17:
for to fn/Bal do
18:
Load Ke, Ve from HBM to SRAM. _ /h- beast beast(m..ast)) ()
19:
Compute W exp
20:
Compute S (1 - re relmmm RG Qr, Ko)) o W/bcast(6.).
Compute 0 = (c-1) + SVe.
21:
22:
end for
23:
end for
Since the statistics equation 7 for each query are computed inde- pendently, the forward pass of ILIA can naturally made parallel for batched queries. Therefore, the algorithm proceeds by iterat- ing over query blocks r. Within each query block, the iteration over key/value blocks has three passes. (i) The first pass (line 12 corresponds to accumulating the statistics M,, wp in an online fash-
1254
Yensay Gnat (MB)
Nate
Baachh staee
ion. Similar online softmax YYe) 2023 Milakov & Gimelshein 2018), we maintain a running maximum my to ensure numerical stability when computing the kernel weights. This trick is valid the computation equation i homogeneous in Way- (ii) The second pass (line 15) encapsulated in the CGSolve ve (-) operator (see Ap- pendix C for details). (iii) The third pass (line 17 22 computes the final output O, using the pre-come results and the values V.. To save computation in backward pass, we also store the intermediate R, and denominator 8, alongside the output into HBM.
133
Seguenee Leugh
FlashLLA re- duces working set memory to (ndd). The figure shows the profiling result for 128, OOM points are omitted.
Figure 1:
We implement and benchmark the algorithm T in custom Triton kernel (~5000 lines of Python) across a range of dimensions and batch sizes on a single NVIDIA H200 GPU. [iguurt[demoosttrut the quadratic dependency and quickly rums out of memory for naïve method. In contrast, the block- wise Tr i on kernel significantly reduces the working set and scales linearly with sequence length, making it hardware-efficies and feasible for long-context and large batch training inference. We also provide additional analysis and latency profiling Appendix
EMPIRICAL RESULTS
Test-Time Regression on Non-Stationary Data. We first devise a synthetic piecewise-linee winear re- gression task to isolate the test-time adaptation capabilities of different attention mechanisms di- rectly without training the query, key and value projections. Each sample is length- sequence par- titioned into T.SS contiguous segments with S being the segment size. For each c (1...., T./S).

---

Published as a conference paper at ICLR 2026
Sryment Size 64
Segment Size 256
Seyment Size 512
Segment Size 1024
10P
LosoyDo Don
10
10
250
Posttin
Mons
Roxdom
Figure 2: Test-time regression performance on piecewise-linear task. The figures demonstrate position-wise MSE for d = 64 with S E {64, 256, 512, 1024]. Results are averaged over 10,000 independently sampled sequences: LLA outperforms other baselines and benefits from more segment data; MesaNet excels only before the first shift. The y-axis uses logarithmic scale.
keys k; E Rd are drawn from a segment-specific distribution Pa supported on a distinct cone in the input space (see Appendix construction details). The corresponding values v; are generated by a segment-specific linear function v; - AAAii+e te, for i E (c-1))++ +l..... cS} where A, ~ N ((,, I) and E; (0, 52 This design ensures the generated data [ki, e)) has non-stationary input distribution and conditional mapping [kk = Ack.
We evaluate single layer of each candidate model, including LLA, Softmax Attention, vanilla LA (Katharopoulos et al.] 2020)) MesaNet (von Oswald et al.) 2025) as well as a random predictor. As this test-time only evaluation, we exclude mechanisms that require training to adapt. We set L = 1024 and sweep over different segment sizes S and input dimensions d. Performance measured by the position-wise MSE & = [k=())) v for i E {L..... L} to capture the adaptation capability along the se- quence. We also investigate the scaling behavior by evaluating the MSE ratio Σ
0001
Mina Baoken
Stn
" Mfodel Σز= CyL for each model compared to LLA across different values of d and S. The results are summa- rized in Figure 2 McsaNet achieves the best performance in the first segment where all the data is drawn from one linear mapping, as it effectively learns the optimal global linear function. However, it has limited capacity to adapt to distribution shifts and its perfor- mance degrades quickly in subsequent segments, where the data distribution does not match the learned global mapping.
10
30
Dicoennke
The advantage of LLA scales with the data di- mension d. Both axes use logarithmic scale.
Figure 3:
Meanwhile, LLA continues to improve within each segment whereas Softmax Attention does not benefit from more in-distribution data. Moreover, the advantage of LLA scales favorably with data dimensionality (Figure 3), indicating the potential for adaptation to larger models and datasets.
In-Context Regression Non-Stationary Data. We next evaluate the models' ability to perform in-context regression on non-stationary, piecewise-lineat data. The data generation process follows the same principle as in the test-time regression task. The data points {Ir Ed= } are generated from segment-specific distributions P. and the target is given by yi = A + E Rd for e {(c- 1)$ + 1, 1,....cc c,,,, where As ~ N xdz /dz) and G; ~ No I. Query x' E Rd. is randomly sampled from the segment distributions Pa. Each in-context regression prompt X is constructed by concatenating L shuffled input-target pairs with L' queries:
X
Φ1
32
LL
I'L. 0
E (R(L+4))((LL))
(17)
21
12
yL
In contrast to the test-time regression setting, the query, key and value projections are parameterized.
The model fe : td, + RR is trained to predict the target Y. where the label to each query is generated by 5/$ = A.r/. Specifically, for a dataset Dersin = (xx()) y } 1o=1; we minimize the MSE loss C(6: D Dtrain) on the query tokens and report the test error L(0": Dreat) after training.
We compare LLA against several strong baselines, including Softmax Attention, Mamba (Gu & Dao. 2024
Dao & Gu 2024). GLA (Yang et al. 2024a). Hyena (Poli et al. 2023) and Gated DeltaNet.
8

---

Published as a conference paper at ICLR 2026
Sequrrce Length Sequenee Length 128
Sequeno 250
Recall Accuracy
1.00
100
1LA
0.78
Test Loss
Accumey
0..00
0.25
128
51341
512.20
280--44
286612
Let-22
L2b-10
1280
64-10
64-1
Segment
Siegraen
Segment
Sib
Sequesee Leugth KY Number)
(a) Test Error of In-Context Regression.
(b) Test Accuracy of Associative Recall.
Figure 4: and for models with d = 128 and 2 attention heads. Each point represents the best performance achieved across training hyperparameters, averaged over 3 random seeds.
Yang et al., 2024b; 2025). We fix the dimension d,. dly 32 and query number L' 16 for all sweeps over segment sizes and evaluate two-layer models without MLPs. The results in Figure +| demonstrate similar trends as in the test-time regression task, where LLA consistently outperforms other baselines across all configurations, particularly with smaller segment sizes.
In-Context Associative Recall. In-context recall is a fundamental capability language models which requires the model retrieve relevant information from the context based on the query. We adopt the MQAR task in Zoology (Arora et al. 2023) to evaluate this ability. Specifically, given two alphabets A, and set of value pairs (k,, v:) E A₃ A,,, the set {A; + e defines many-to-one key-value association. The model is prompted with sequence of key-value pairs and then queried with keys sampled from the context predict the corresponding value.
As indicated in WWang et al. 2025). ) a single short convolution layer is sufficient to solve next-token recall tasks. Therefore, we disable short convolution in all baselines ensure fair comparison. We set || U A.ll 8k and sweep over different sequence lengths and number of KV pairs. The test recall accuracy results in Figure 1b indicate that the advantages of LLA can be effectively trans- ferred to discrete token prediction tasks. Additionally, we also observe different learning dynamics between LLA and Gated DeltaNet in this task. The results are discussed in Appendix F.4
We then test models' state-tracking ability by permutation state- tracking task. Given an initial assignment of items to positions, sequence of swap instructions, and query positions. The model is trained to predict the item at each query position after all swaps. Each example is constructed as
Permutation
P1 (1, P2= 02.
PN ON
$1, 82: ..... is is
91=01 ... IN= ON
initial state
swap instruction
query answer
Here Pa E [I,...,N] N} denotes position n; on E A is the item initially assigned to n; each (is, E .N....,N2 is a swap in- struction exchanging the items at positions i, and j.: and q is the queried position. The target the item at position & after applying all S swaps. We include explicit delimiter tokens #, =, and for structure.
State Trundding Test Accuracy
We draw the number of swaps as S Uniform(N/) N/3) and set |A| = 8k for each example. The results in Figure show that LLA achieves test accuracy on par with Softmax Attention across N. This outcome is expected from a wetplexxitttthorrtttee perspec- tive: constant-depth Softmax Attention is no more expressive than constant-depth threshold circuits TC and has limited ability to real- ize unbounded-deptth state-traas tracking as N grows (Hahn 2020 MMr-
Figure 5: Test Accuracy of State Tracking. Results show the best score averaged over random seeds.
rill & Sabharwal 2023). By Eqs. and 8, LLA augments Softmax Attention with a query-specific first-order correction computed via a constant number of parallel algebraic passes (weighted sum, inner product, inverse), which adds at most constant extra circuit

---

Published as a conference paper at ICLR 2026
layer. Its performance therefore matches the theoretical limits of Softmax Attention, explaining the results in Figure 5
5 LIMITATIONS AND FUTURE DIRECTIONS
High Computation and I/O Intensity. Despite the significant reduction memory consumption, ILA's computational cost remains substantially higher than that of Softmax Attention, primarily due to the matrix inversion involved in the computation. Exploring approximations to reduce the computation an important direction for future work. Furthermore, while FlashLLA achieves the same (nd + n²) I/O complexity as in FlashAttention when the number of CG iterations is set as a constant (which is sufficient in practice), the constant factor is still higher due to the additional reads and writes required by the iterative solver. Incorporating hardware awwarre--wware optimizations, such as sliding windows or sparsity, could further reduce I/O complexity.
Kernel Development and Evaluation on LLMs. This work evaluates LLA synthetic and moderate-scake tasks; its efficacy on large language models remains an ongoing question. Training LLMs with LLA using PyTorch implementation infeasible due to its high computational and memory complexity. Therefore significant engineering efforts are required to stabilize and optimize the forward and backward kernel. Additionally, the numerical sensitivity of the matrix inversion poses a challenge for developing low precision kernels without sacrificing performance.
Efficient Interpolation of Linear and Softmax Attention. As shown equation 9 LLA pro- vides optimal interpolation between Linear and Softmax Attention in solving the regression ob- jective. And the formulation also provides a template to design algorithm for better computational efficiency while still retain strong estimation capabilities and potentially even improve upon the cir- cuit complexity of Softmax Attention. For instance, future work could explore the integration of state-of-the-att Linear Attention architectures such DeltaNet and Mamba using this template.
REPRODUCIBILITY STATEMENT
The proof' of theoretical results 2.1 and 2.2 are provided in Appendix and respectively. We also provide additional analysis and discussion on the implementation and profiling results in Ap- pendix In Appendix F we provide detailed instructions, configurations and additional experimen- tal results for reproducing all the experiments in Section
REFERENCES
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to imple- ment preconditioned gradient descent for in-context learning. In Advances in Neural Information Processing Systems 37 (NeurIPS 2023), 2023.
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. In International Conference on Learning Representations (ICLR), 2023.
Simran Arora, Sabri Eyuboglu, Aman Timalsina, Isys Johnson, Michael Poli, James Zou, Atri Rudra, and Christopher Ré. Zoology: Measuring and improving recall in efficient language mod- cls, 2023. https: http:::///xxvxoaaeaa abs/ 04927
Ali Behrouz, Meisam Razaviyayn, Peilin Zhong, and Vahab Mirrokni. It's all connected: A jour- ney through test-time memorization, attentional bias, retention, and online optimization. In 14th International Conference on Learning Representations, ICLR 2026, 2026. URL https:
//openzeview.cr wet/forune wttp::/eeooooaaaa/gg..a........ ddd.dd==22yymmmom0m
Aviv Bick, Eric Xing, and Albert Gu. Understanding the skill gap in recurrent language mod- mechanism. In Aarti Singh, Maryam Fazel, Daniel Hsu, Simon Lacoste-Julien, Felix Berkenkamp. Tegan Maharaj, Kin Wagstaff, and Jerry Zhu (eds.), Proceedings of the 42nd International Conference Machine Learning, volume 267
els: The role of the aadeer-ad-aggreee
10

---

Published as a conference paper at ICLR 2026
of Proceedings of Machine Learning Research, pp. 4324 4344. PMLR, 13-19 Jul 2025, URL httts::/
wttp:/iiiiinmm.....e.... wwt7/bick22222 htt...ppe..s//2////bc.... htmml
J. Bierens. The nadaraya-watson kernel regression function estimator. Technical Report 1988-58, Faculty of Economics and Business Administration, Vrije Universiteit Amsterdam, Amsterdam, 1988.
Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Shuming Ma, Zhifang Sui, and Furu Wei. Why can gpt learn in-context? language models implicitly perform gradient descent as meta-optimizers, 2023. URL https::/aaxii...../.// org/ abs/ 2212. 10559
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. In 12th International Conference Learning Representations, ICLR 2024, 2024. URL https://
openreview. wwweeoouummae www.ummiddm2222aa m2n2Xyh3Ec
Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality, 2024. URL https://arxiv. org/ abs/ 2405. 21060
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashatten- tion: Fast and memory-efficient exact attention with i/o-awareness. In Advances in Neural Information Processing Systems 35, pp. 16344-16359. Curran Associates, Inc.,
2022. URL httpps://proceeediggs.ott neuzips...e.urissc.c.cs..sss cc/pppe files/pape//2022/ File/ Eage2777e2220000032233eeeeeeeee2
pdf
Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer, Andreas Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, et al. Scaling vision transformers to 22 billion parameters. In Proceedings of the 40th International Conference on Machine Learning (ICML 2023), pp. 7480-7512, PMLR, 2023.
Jianqing Fan, Irêne Gijbels, Tien-Chung Hu, and Li-Shan Huang. A study of variable bandwidth selection for local polynomial regression. Statistica Sinica, pp. 113-127, 1996.
Yizhao Gao, Zhichen Zeng, Dayou Du, Shijie Cao, Peiyuan Zhou, Jiaxing Qi. Junjie Lai, Hayden Kwok-Hay So, Ting Cao, Fan Yang, and Mao Yang. Secrattention: Learning intrinsic spanse attention in your llms. arXiv preprint arXiv:2410.133: 13277 2024. /Rtp:ts:://axxi
abs/2410.13277
Shivam Garg. Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in- context? case study of simple function classes. In Advances in Neural Information Processing Systems 35 (NeurIPS 2022), pp. 30583-30598. 2022.
Felix A. Gers, Jürgen Schmidhuber, and Fred Cummins. Learning to forget: Continual prediction with Istm. Neural Computation, 12(10))2451-247 2000.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. In First Conference on Language Modeling (COLM), Jul 2024. URL Available on RevnRe
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. In 10th International Conference on Learming Representations, ICLR 2022, 2022a. httpthtsss///ppp/et///s http:/wwwwwweeeeeeeeeeeeeeeeeeee -ow.
Albert Gu, Isys Johnson, Aman Timalsina, Atri Rudra, and Christopher Ré. How to train your hippo: State space models with generalized orthogonal basis projections, 2022b. URL https:
wwrxiv.org/abie 12007
Michael Hahn. Theoretical limitations of self-attention neural sequence models. Transactions of the Association for Computational Linguistics, 8: 8:156-171, 2020.
Magnus R. Hestenes and Eduard Stiefel. Methods of conjugate gradients for solving linear systems. Journal of Research of the National Bureau of Standards, 49(6):00-443 1952.
11

---

Published as a conference paper at ICLR 2026
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers mns:: Fast autoregressive transformers with linear attention. In Proceedings of the 37th Inter- national Conference on Machine Learning (ICML 2020), volume 119, pp- 5156-5165. PMLR, 2020.
Louis Kirsch, James Harrison, Jascha Sohl-Dickstein, and Luke Metz. General-purpose in-context learning by meta-learning transformers, 2024. URL httpas//arxxivoorgg/bb////21
04458
Zhixuan Lin, Evgenii Nikishin, Xu Owen He, and Aaron Courville. Forgetting transformer: Softmax attention with a forget gate, 2025. URL BRRLttts:: :tap:://aaxxa //wpxaaaivvv////////////////aivv org/ abs 2tt.b///5533.1000
Bo Liu, Rui Wang, Lemeng Wu, Yihao Feng, Peter Stone, and Qiang Liu. Longhorn: State space models are amortized online learners. In Proceedings of the 41st International Conference on Machine Learning (ICML 2024), 2024.
Enzhe Lu, Zhejun Jiang, Jingyuan Liu, Yulun Du, Tao Jiang, Chao Hong, Shaowei Liu, Weiran He, Enming Yuan, Yuzhi Wang. Zhiqi Huang. Huan Yuan, Suting Xu, Xinran Xu, Guokun Lai, Yanru Chen, Huabin Zheng, Junjie Yan, Jianlin Su, Yuxin Wu, Neo Y. Zhang, Zhilin Yang. Xinyu Zhou, Mingxing Zhang, and Jiezhong Qiu. Moba: Mixture of block attention for long-context llms. In Proceedings of the 38th Conference Neural Information Processing Systems (NeurIPS 2025), Spotlight, 2025. URL
net forum? R
Xuezhe Ma. Xiaomeng Yang, Wenhan Xiong, Beidi Chen, Lili Yu, Hao Zhang, Jonathan May, Luke Zettlemoyer, Omer Levy, and Chunting Zhou. Megalodon: Efficient Ilm pretraining and inference with unlimited context length. In Advances in Neural Information Processing Systems 37 (NeurIPS 2024), 2024.
Arvind Mahankali, Tatsunori B. Hashimoto, and Tengyu Ma. One step of gradient descent is prov- ably the optimal in-context learner with one layer of linear self-attention. In Proceedings of the 12th International Conference on Learning Representations (ICLR), 2024.
William Merrill and Ashish Sabharwal. The parallelism tradeoff: Limitations of log-precist precision trans- formers. Transactions of the Association for Computational Linguistics, 11:531-545 2023.
Maxim Milakov and Natalia Gimelshein. Online normalizer calculation for softmax, 2018. URL http::///aaaiv..o/aa///////////
E. A. Nadaraya. On estimating regression. Theory of Probability and Its Applications, 91):141-142,
1964.
Bo Peng, Daniel Goldstein, Quentin Anthony, Alon Albalak, Eric Alcaide, Stella Biderman, Eu- gene Cheah, Xingjian Du, Teddy Ferdinan, Haowen Hou, et al. Eagle and finch: Rwky with matrix-valued states and dynamic recurrence, 2024. URL pettps://arxvv.org// 2404.000 05892
Bo Peng, Ruichong Zhang, Daniel Goldstein, Eric Alcaide, Xingjian Du, Haowen Hou, Jiaju Lin, Jiaxing Liu, Janna Lu, William Merrill, al. Rwkv-7 "goose" with expressive dynamic state evolution, 2025. httlsssss:///xxiit
org/ wrtp:///ggaaaabaabee.. 14456
Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y. Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models. In Proceedings of the 40th International Conference on Machine Learning (ICML 2023), volume 202, pp. 28044828007. PMLR, 2023.
Ofir Press, Noah A. Smith, and Mike Lewis. Train short, test long: Attention with linear biases ables input length extrapolation. In Proceedings the 10th International Conference on Learning Representations (ICIR), 2022.
Hubert Ramsauer, Bernhard Schäfl, Johannes Lehner, Philipp Seidl, Michael Widrich, Thomas Adler, Lukas Gruber, Markus Holzleitner, Milena Pavlović, Geir Kjetil Sandve, Victor Greiff, David Kreil, Michael Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. Hopfield networks is all you need. In International Conference on Learning Representations (ICLR), 2021.
12

---

Published as a conference paper at ICLR 2026
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers secretly fast weight programmers. In Proceedings of the 38th International Conference on Machine Learning (ICML 2021), volume 139, pp.
Virtual only, 2021. PMLR.
Jürgen Schmidhuber. Learning to control weight memories: An alternative to dynamic recurrent networks. Neural Computation, 4(1):131-139 1992. doi: 10.1162/neco.19/1 199..4...999......1....
Julien Siems, Timur Canstensen, Arber Zela, Frank Hutter, Massimiliano Pontil, and Riccardo Grazzi. Deltaproduct: Improving tracking in linear mns via houscholder products, 2025.
URL
wtt. http://ggaaaa/////// 10297
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and
Furu Wei. Retentive network: A successor to transformer for large language models, 2023. URL
/ittps:/aariii www.orgg/bbe.. abs httaaabb///333/0 230720007.330
Chameleon Team. Chameleon: Mixed-modal early-fusion foundation models. arXiv preprint
DrXi::2405.09818, 2024.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Infor- mation Processing Systems 30 (NIPS 2017), pp. 5998-6008, Curran Associates, Inc., 2017.
Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordy- intsev, Andrey Zhmoginov. and Max Vladymyrov. Transformers learn in-context by gradient descent. In Proceedings of the 40th International Conference Machine Learning (ICML), volume 202 of Proceedings of Machine Learning Research, pp. 35151-35174, PMLR, 2023.
Johannes von Oswald, Maximilian Schlegel, Alexander Meulemans, Seijin Kobayashi, Eyvind Niklasson, Nicolas Zucchet, Nino Scherrer, Nolan Miller, Mark Sandler, Blaise Agüera y Ar- cas, al. Uncovering mesa-optimization algorithms in transformers, 2024. URL https:
//axxiv.orq/abs//20000
Johannes von Oswald, Nino Scherrer, Seijin Kobayashi, Luca Versari, Songlin Yang, Maximil- ian Schlegel, Kaitlin Maile, Yanick Schimpf, Oliver Sieberling, Alexander Meulemans, et al. Mesanet: Sequence modeling by locally optimal test-time training, 2025. URL https:
//axxiv.org/abb/2500..o 05233
Ke Alexander Wang, Jiaxin Shi, and Emily B. Fox. Test-time regression: unifying framework for designing sequence models with associative memory, 2025. URL httlhttts//aaxiitt org/ abs/2500 12352
G. S. Watson. Smooth regression analysis. Sankhyā: The Indian Journal of Statistics, Series A, 26 (4):359-372, 1964.
Mitchell Wortsman, Peter J. Liu, Lechao Xiao, Katie Everett, Alex Alemi, Ben Adlam, John D. Co- Reyes, Izzeddin Gur, Abhishek Kumar, Roman Novak, et al. Scall-scale proxies for large-scale transformer training instabilities. arXiv preprint arXiv:2309. 14322, 2023.
Guangxuan Xiao. Statisties behind block sparse attention.
com/
blog//blcck-sparse-atttt
html] 2025.
Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with harfwite-efficient training. In Proceedings of the 41st International Conference on Machine Learning (ICMI. 2024), 2024a.
Songlin Yang. Bailin Wang. Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear trans- formers with the delta rule sequence length. In Advances in Neural Information Processing Systems 37 (NeurIPS 2024), 2024b. doi: 10.52202/079017-3666
Songlin Yang, Jan Kautz, and Ali Hatamizadeh. Gated delta networks: Improving mamba2 with delta rule. In Proceedings of the 13th International Conference on Learning Representations (ICLR 2025), 2025.
13

---

Published as a conference paper at ICLR 2026
Zihao Ye. From online softmax to flashattention. Course notes for "ML for ML Systems" (CSE 599M). University of Washington, Spring 2023, May 2023. URL http:::/pss/ttt http//ouussssceooeeooe cs. washington edu/ wttda//ooursseeeeeeeeeee http:///sppppp///eeeeeet wlwsattt.aaatatnaaaaaaaamaaaattt pdf Online; Uni- versity of Washington.
Jingyang Yuan, Huazuo Gao, Damai Dai, Junyu Luo, Liang Zhao, Zhengyan Zhang, Zhenda Xie, Yuxing Wei, Lean Wang, Zhiping Xiao, Yuqing Wang. Chong Ruan, Ming Zhang, Wenfeng Liang, and Wangding Zeng. Native sparse attention: Hlrdware--lligne and natively train- able sparse attention. In Wanxiang Che, Joyce Nabende, Ekaterina Shutova, and Moham- mad 'aherr Pilehvar (eds.), Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 23078-23091 Vienna, Austria, Jul URL
2025. Association for Computational Linguisties. doi: htt.::///pecaaatoooooo..../////.
Ruiqi Zhang, Spencer Frei, and Peter I. Bartlett, Trained transformers learn linear models in-context. Journal of Machine Learning Research, 25(49):1-55, 2024.
Shu Zhong, Mingyu Xu, Tenglong Ao, and Guang Shi. Understanding transformer from the per- spective of associative memory, 2025. //apxis://xxiii..// wrw.vvoorrrggooooo.. abs /205. 1 9488
14

---

Published as a conference paper at ICLR 2026
A APPENDIX: PROOF OF PROPOSITION 2.1
Let (X,, Yi) be i.i.d. with X; E IRd supported on bounded domain D C R.d with density p. and
Y = f(X,) + E R4, EE, X = 0, E X = (X,,.
The global linear estimator at x E is
fcl(2) = 01,x1=
where 0 gd+1 minimizes the empirical squared loss over the global affine class.
Let K Fd 7 [0, oc) be a bounded, compactly supported, radially symmetric kernel with (u) du For symmetric p.d. bandwidth matrix H H with a constant condition- number upper bound K1, write
(u) |/|- K(H-1/2
| + 0,
n|H +
The NW estimator at x E is
FNW(x) 1 WWY
1 11
where W ==
- x)). x)). = (L,...,1) R and Y = (Yi,..., ..)....,,,......... E Ruxd,
To make the analysis casier, we assume the following smoothness requirements.
Assumption A.1. The domain D has C² boundary (defined as 0D) with principal curvatures uni- formly bounded by K2-
Assumption A.2. The density function p of X satisfies p E (D) For all dimensions
the function I satisfies I3 E (D),, and the variance function x² satisfies o3 E C(D..
We also assume that the kernel has easy-to-handle support.
Assumption A.3. K is radial, compactly supported in the unit ball Bdd = {uERR || 1) sK)
Throughout, E[] denotes expectation with respect the randomness of the training sample S, {(X,, K.) while (-) dæ denotes the Lebesgue integral over the spatial domain D. Because
dy
E mf. 1f(x)) s(0) = Σ E (f(0)
ΣE
(5(x))f560)) fz(2)) -
dz.
the output dimension dy only induces summation across components and does not affect the order; hence, without loss of generality, we take dy = below.
A. INTEGRAL ERROR ESTIMATION OF GLOBAL LINEAR REGRESSION
We consider the global affine class
G { (0 = 3o ATx 0 (Bo, E tdt
Lemma A.1. If & G, there exists a constant Ab > such that for every n,
E B [[ (fan.())-(xx)) f(x))² dx 2 Ab.
= {h D with h(x) dx < oo} endowed with inner product (hi, 122 = SD (1))22(2)) dr. The set G is a dines-aimansiooal linear subspace and hence closed in L2, dx). By the Projection Theorem, the LD, oothhoggl projection of f onto & exists and is unique:
Proof. Let
g E arg min gEg
|| -01l(>=0 = arg min
gEg /
- g(x))² dr.
15

---

Published as a conference paper at ICLR 2026
Set Ab = 9 f $ 9, then gt + in (D) and thus 4b 0.
Fix an arbitrary realization Sn = (X,, Y) Since (= E g, the optimality of gt yields
(o(x) - f(x)) dx AD. J (fonl(2; f(x)) dæ 2 inf gEg
This inequality holds for every sample Sn. Taking expectation over the training data proves the claim.
L
A.2 POINT-WISE ERROR ESTIMATION OF LoCaL CONSTANT REGRESSION
In this section we will estimate the point-wise mean-squared-error of NW estimator, whose expres- sions given by
MSENW (z) == E (Fsw(t)) {X
(E [Nw(2=) f(x) {x)==)
(fow(e) - -B (Fww(()) | (x));=) 1) 1x
+
Binsyw (rr)²
Varww(#)
When estimating at x E D, the kernel KH maps the translated domain the scaling H-1/2 The mapped set differs depending on whether lies in the interior of D or near 0D. This geometric difference drives larger boundary enror for NW, which will underlie the performance gap between NW and I.I., We formalize the mapped kernel domain and the boundary layer.
x) into supp(K)
Definition A.1 (Exact kernel domain and boundary layer.). For any E D and bandwidth H, the exact kernel domain
D2, : 111 x) Ssspp(K) {u E IBd x + 112 D).
Define the boundary layer by
B(H) (xE:: Da,H # supp(K))
We use the following kernel- shorthands. Definition A.2 (Exact kernel moments). For all x E D and bandwidth H, we define
(z(2, H) [ J K(u) du, (1 H) = [ D
15(x; H) = f
(()) du,
wu T K(u) du.
Dx.
Define normalized moments (z, H) = (2 H)/H((z, H).
Then we are ready to estimate the bias and variance using H and n.
Lemma A.2. Under Assumption A.2 we have
Biasy (I) Vx(x)) H 1/2 122 H) + H)()H)), Varxw(x) = =Op( n|H 1/2
uniformly for all H E Hn. where H, { h²B h in n B B. 1, K< k(B) <
Hi and < <.
Proof. Defining f (f(X)), f(Xx)) for any fixed point xo E D, we have Bias (Xo) = (1 T w1) -ww-i111 11--11T[ - fl)
Specifically,
n -1 (1) = Σ
Kn(X, - Io), 1 f(x))1) =n -1 Σ Kn (x)) - f(xo))
16

---

Published as a conference paper at ICLR 2026
By Chebyshev's inequality, for any fixed H,
Σ Kn xo) K zo)p(=)d + Op n Vm K2 zo)p(r)dx
By Theorem in (Fan al. 1996), this bound holds uniformly over H Hn, after multiplying the stochastic term by v /log n. Hence,
KH(X, -20) x0)= = S,, Kn(x 20)p(()d + Op n n nnoon / K2 zo)p(z)dr
n
f
K(u)p(zo + u) + op((1) = p(zo)uo(( H) opp
Dzg.H
(18)
uniformly for H Hn.
Similarly,
K f(xo))
1=1
f(ro)) J. 200()(a)) - +Op n nlog / k2 -
=fo K(u)p(zo + 112 (J + 11) f(xo))du + op(1H11
DDo..
=f D
(((o) 12 (7 H) Vp(2.) (|) (|
=p(zo) ((D)) T 1/2 1/ (mo, H) + ([]
(19)
uniformly for all H E Hn.
Combining Equations 18
19
we have
Bias (mn) (1(o)) /2 V2 H) + Op([
uniformly for all H e Hn.
Then calculate variance
VarNw( (1 (1" W1)-
where Σ =
(x))) Analogously to equation 18
( ΣX n Σ K (X: =X))
/ (2 - -x0) + Op (n-1 ( nlog
Inlong / K (z)((()))(2))
log
=1111//
K²(u)o?(an 11222)p(20 + H1 122 + Op
H
3/4
(20)
Dra.t
n
llg
11/2
K²()) (2uu))((a))p))))dd((1+ (z)(1) +Op
n H 13/4
Dra.
/ {((w)du(1+ogg11)
12 (20)p(20)
Da
uniformly for all H Hn.
17

---

Published as a conference paper at ICLR 2026
Combining Equations 18 20
we have that
(a)) = n|H
holds uniformly for all H Hn.
A.3 INTEGRAL ERROR ESTIMATION LOCAL CONSTANT REGRESSION
Here we calculate the integral error for NW estimator. From Lemma A.2 we have that for any a D, the pointwise bias of NW estimator is
Biassw (I) f(x)" H (z. + (p(
By symmetry of the kernel K, the first moment (z, H) = 0 if and only if & B(H). Roughly, (iassw(x) = O(I [1/21]) when x E B H)) and Biassw(r) O() HIl) when æ & B(H). ). Hence
important to bound Sa( MSFNw (a)) dx and SDABC) MSFxx (I) dx separately. We first define
a smooth substitute for the boundary layer ((H).
Definition A.3 (Uniform Boundary Layer). For any bandwidth H and (0,1), the uniform boundary layer of D is
C(H, a) = - y te(y) y E OD, E [0, ahn3
where for any y E 0D, we define ely) as the inward unit normal y, and define hn(y) Vely) H ely).
We now show that one can choose constant a E (0,1) (independent of H) that "sandwiches" (H) between C(H, a:) and C(H, )
Lemma A.3. Under Assumptions A. A.3 H, we have
there exists E (0, 1) such that for all sufficiently small
C(H,a) C B(H) C(H,
Proof. We prove the claim with a min( 2 1 1/2)
Step 1: a(H, C B(H). Let denote the signed distance to 0D (positive inside D). For - te(y) E EC(H, a) and any v E IRd, the standard expansion (using the shape operator Syy gives
Φ(r v) + Un SuUT 0(1oo1)),
where vn = Tely) and er = - (I-e(y)e(y)T) By Assumption A.1 ||S is uniformly bounded.
Define
u u: Ki 1/2 k,,(y) H 1/2e(y) ERC
Then
||u]|2 = KI Nn, hn.(y)) ((y)))()))T)))))))))))))))))))) H H- = Ki (e(v) He(w)) (e
< KT x(H) S 1,
since x(H) < Ki by assumption. Hence E supp(K) (Assumption A.3 Using the expansion
with 1122u and bounded curvature,
6(x + H 112 - H /1 + O(I (3/2) = = 1/2 h (v) + +( -1/2 Since S ahn (y) and a S | 1/2 for sufficiently small H we have a(x + H2/2 > 0, i.e. + 22222u & D. Thus there exists E supp(K) with & Dr.H. i.c., Dz.H + supp(K), (K), so B(H).
Step 2: B(H) C C(H,a- Let = y- --((e(y) E B(H). By definition, there exists a E ()) with r + 11222 & D.i.e,
o(x 0(x+H112)) = - TT H 1/2 + 00
Hence
0(111) < hn(()) + O(IHH)) t 11 ( + M < ||u|| Tel He(y) +
For sufficiently small H, this yields/ <h-((( (since a 2 2). so c(H,
18

---

Published as a conference paper at ICLR 2026
With C(H, a ") in hand, bound the integrated variance and squared bias of the NW estimator.
Lemma A.4. Under Assumptions A.7
A.2 and A.3 we have, uniformly for all H Hm.
S. BinsSw(()) dx Op (p(1) 3/2)
and J. Varnw(x) dx = A(n H 1/2))
Proof.
J Wias (a))dr / (B(H)) Bias (z)dx
Bias (z)dr
(DBB(H)
= = JE(I) dx + 1p(111121)
< 0,,(HHI SeCH.a 1)
da + +,((
(1((1 /
didS(y) + + (1
= 00((1HH32211
since (v) = Vely) He(y)) = 01122113/3 uniformly under bounded condition number.
For the variance, Lemma A.2 gives Varnw(1) (x) = Op | - 1122) pointwise (uniformly over H e H.,). Integrating D (whose volume is constant) preserves the order:
J VarNw(I)dr= = / O(n (12-)))) 1
A.4 PROOF OF PROPOSITION 2.1
Combining Lemmas A.I A.4
obtain the final conclusion.
Theorem A.1 (Precise statement of Proposition 2.Ip. Under Assumptions A.7 A.2 tion f to be estimated is not within the the global affine class G,then
A.3
the func-
E MSEcL (z) dz ((1) and E Sf MSENw(2) dx On 3/(13)))
where fNw uses the optimal bandwidth H e Hn.
Proof. The lower bound for MSEc1, follows directly from Lemma A.I
For NW, Lemma A.4 and the definition of H, imply
S Var (2) dr = 6(n-th-9
Bias (x) dx = Op(h)),
uniformly for E [n n (since ||l|| = (h?) and H11/2 =(h"") under bounded condition number). Hence
J. MSENW (x) dr 0 hi 1
n nhd On(n -3/(433))
the minimizer h -n 11/(43) E [n_",,n of h (nh
B APPENDIX: PROOF OF PROPOSITION 2.2
The local linear estimator at x E is fu(()) (X Wx) -
where
X
X1
T0
X,
2o
Here inherit Assumptions
A.I
A.2
A.3
and continue set dl 1 as in Appendix A
19

---

Published as a conference paper at ICLR 2026
B.1 POINT-WISH ERROR ESTIMATION
We have already derived the point-wise error of the local constant estimator in Lemma A.4 do the same for the local linear estimator.
We now
Lemma B.1. Under Assumption A.2 we have, uniformly for all H E Hn.
(2)) = O,(1HH),, Vari..(x) = Op( 12 1
where Ha = (H- B²B: : h in",, ". B > O, |B] = 1, x(B) S k1} with constants
Proof. For any fixed point xn D, we have
(i))),1( = e (X
W X
X((o)
1
= 2 (X (XTWXX xI ( WQQ + +x)(t((H))
where Q = 1=(( H,(XI - z).... ..... (Xn - x) x) Hy - x)!T.
VarLL( = ei (X WX) 'x T EX)( T (X)
where E = diag (K - (x,))
n (xXxXXXX =
nー -ー m KH (X, =00)
Li=l
n Σ i= KH (X, 20) (X, =20)T
n- - K (X; - -zo) (X, - -xo) n n-
K (X; - xo) (X, - zo) xo) TT
1x ΣX)
n -ー> T1
1= -
n " -00) =TO)T
n n
K2 11
- (X,
-xo) z0) n '- K 2x,) - zo)(X, x00) T
X W(f - (f(ox)
n- E -x) - -Vz)))(((zz) (X, =X)) xo)
- (x)) f(zo) Vf(xo) zo))
By the same uniform LLN and Vlog arguments used in the proof of Lemma A.2 have uni- formly over H E Hn:
n
n
-x, x0)00 -XO)T
n
log 11 (w/10r
( 1/2 (K)H1/2 +0,
(k),
3/4 1)/11111
n
n
KH - 20) = H) H)
Σ
- - 2x)0 - xo)(X, - =00
n
1=1
= (600)(( (2 + op (H111
20

---

Published as a conference paper at ICLR 2026
Kn(X; (onXX-30)(x xo)" =p(zo))i H) 1//2 (H1)
KH 200 zo)(X; =00 = ==)T==ppp 0
1=1
KH - xo)(f(Xi) - f(xo) - Vf(2o) (XX, -x)) - (K) + (x(()(H)
72
1m1
10
-20) f(xo) (=00) (x, zo)) xo) = (1(11//11),
where
Ro (Zo. H) = I (²())du, (I H) [- Dzo.
(u)uuudu, R2(ro, H) So Dzo.
(²(u)un' du.
Dao.
Then we have the expression of every element all matrices. A standard blockwise inversion therefore yields
(iAsLI( O,(HH VaLLL()) = Opl n//11//2
uniformly over H H.s proving the claim.
B.2 INTEGRAL ERROR ESTIMATION
From Lemma A.2 have that for any E D, the pointwise bias of local constant estimator
Bias (x) v H (x, (z,
(p(
By symmetry the kernel K, the first moment satisfies (x, H) = if and only x & B(H). We will show that if f has sufficiently large normal gradient measurable subset of OD, then SBL Bias (a:))w() will have dominant order ( (3321))....
We begin by lower-bounding Vf(x) 1 (x H)), to do which we first lower-bound e(y) p/ - H).
Lemma B.2. Under Assumption A./ A.3 there exists a E (0, 1), c, > 0 and C, such that for all sufficiently small H and all OD, all E [0, (b)l
(e(u) T 11 te(y), H) 2 (v)
|V te(y). HIll s C..
Proof. Recall Dz,H {u Bd + uu E D}. For x y te(y) with y E aD, choose an orthogonal Qyy) such that QQy) /2 h, (v) ca. Define the rotated domain
Dy u E _d y te(y) H 13 e D).
Define the corresponding moments on Dy,,
(0 = - K(u) du, (y, H): - J u k((u) du, (,, Pi. H.) (y 1, H)
Then
Mo, (y. t. H) = nā(x, H), Mi, H. = pī H),
(. = H),
and hence
e(y) H /2 (1, H) = (v) eJ (y, H).
(21)
By the standard signed-distance expansion with shape operator S, (Assumption A.I for Rd
( - te(y) H2/2 - hh)())
1= 3/13))
21

---

Published as a conference paper at ICLR 2026
where 27 (I - -(y))cly)
( " (z))) and ||S is uniformly bounded. Therefore I va hn) + 0 (11
y y- tely) + H (())
80
Dys. = { v BBd va (n(
t 0(1111111)
Consequently,
/ By..
(U,, =
K(v) de
{ dui:d-1
(t/bnly)
K (v va) dva + 0111/1
-ge-1
2 dui:d-1
--VI-I-- (V11dd1, ua) dua + H1
Using positivity and boundedness of K on Bd, there exist < < < oo such that
C < (U,, H) < D.
(22)
Similarly,
e (y,
VaK(U) du t//in(u)
Byn..
duida-1
Us (v doa + 0(1
+ 0 H1122)
Note g()) and g(T) decreases as T 4 Choose a = E (0,1) with g(an) < Then for all E [0, an hn (v)).
(i, S g(an) 0(1HH) < 0.
(23)
Combining equation
22 equatiin 23 for small H,
(u, H)
Ig(an)]
[e] (),
(v,t, Mô (v. H)
2D
Using equation 21 yields the first bound with c. = (g((a))/(2D) boundedness of K and supp(K) C Bdd
The second bound follows from
Since the leading term of (iassw(*) is V/() HJ), the lower bound on HT/22 te(y), H)] alone does not guarantee order H(22]]): also require suffi- ciently large normal derivative of f along OD.
Definition B.1 (Extreme boundary gradient class). For any domain and constants m and M, we define class of functions EDD, m, M), where f E(D,m. m. M) iff there exist measurable m and [|Vr[(y)l < M where
C AD with S(T) O and constants m, M such that [0)] 2 0. 0zf(2) v f() e(y) and Vrf(() = e(y)e(y) fV)()) (V)(())
Then we can prove that if the function to estimated is within E(D, m, M) where m and M are specifically chosen constants that independent of H, the NW has integral squared bias with high order.
Lemma B.3. Under Assumptions A.1 A.2 and A.3 if E E(D, m. M) with
0, K 1
2c, Ks
1/2C. mM 2 C > 0,
then uniformly over H E Hm,
/ Biash dr = H1 3322)
and
Varnw(x) de 9p(n 17 11/2)
22

---

Published as a conference paper at ICLR 2026
Remark.
It is easy construct f and D satisfying Assumptions A.1
A.2
and E(D, m, M)).
For example, on (211, : of + 22 S 11. (21 m, M.
.CIRI ( + 23) works with suitable
2n.
Proof.
We first lower-bound bias. Choosing the minimum given by Lemmas A.3
B.2
have
{ Biask 2 SSLn (iaskw(z)dd 2
Bias (x)da
JB(H))
((HH.))
= JmD (ahn(v)
( (Syy)
-tty dtdS(y)
(24)
2 C2
H p7 ? te(u), H))2 + 02(1119//111) atds(v).
where the last inequality is derived from Lemma A.2 and the boundedness of Sg from Assumption A.1
For all E ah,, (y) and e (0.
- te(y)) 12 - Vrfly ( 1/
@ - 11
te(y),H))
& (1 v) - (1
7 (VrS( - ( 12 te(y), H1)
d 11
= (1 110:0000000000011111000011111111 1
- 2 ((1 m)m -
Here (a) uses the fact Vj = ( (b) uses the fact (A+ 2 (1-1)) -1
for all 1j E (0, 1), (c) uses Lemma B.2
Setting 7 = MC.N1/2
we have
mc,
Vi(y te(u)) H HT te(v), H))
((1 - (1-m)m²cch p 12 || + 0(1113//111
(25)
( -2. -2 Cmmm) ||||| 0(11H3/111)
2 C1 Cill + +1(1
Combining Equations
24 25 we have that
I (2)dew 2 aC2 (GillHI 0/(1112//11)) / ())d((())) ( (1/41/1)3311)
The last equation holds because hn (v) 2 xi x1 Using Lemma A.2 it is straightforward to show that
/ VaINw(I)dr= = 2p 12/2111).
We then show that local linear regression enjoys strictly lower bias than NW on all functions.
Lemma B.4. Under Assumptions A.1
we have So (i)s Op(1|H211) and uniformly hold for all H Hn.
A.2
A.3
So VarLa(z)dx= O,(n-1|1111
23

---

Published as a conference paper at ICLR 2026
Proof. It is straightforward from Lemma B.1
Then it is easy to prove the final conclusion.
Theorem B.1 (Precise statement of Proposition 2.2 Under Assumptions A.7 tion] I to be estimated is within E(D, m, M) where
the func-
2xi'mm² 20,61 1/2 C.mM C1 0,
have E S. MSENw (z)dzz 12nn (3((3))) and I Sn (T) O (n==++))))), where both fow and fLA. are at their optimal bandwidth H Hn..
Proof. According
Lemma B.3 and the definition of Ha. we have that we have So Biassw (x))d (h²) and Jo Varnw(z))dz 2p(m uniformly hold for all h
[n a n b1. So we have
S. (ISENw(F)de (1² + nh. = 2pp(n -onn-3//m 3/(d+3))
even for optimal E Hy. The last equality holds because = n n-11/d E [n="n -1 is the minimizer of h + (nh*)-1
According to Lemma B.4, we have that
MSELL(z)de = (, + nhad On(nn
for the optimal H E Ha. According to the definition of S2p and Op. it is straightforward to deduce the conclusion.
C APPENDIX: CONJUGATE GRADIENT SOLVER
The Conjugate Gradient (CG) method (Hestenes & Stiefel 1952 is an iterative algorithm for solving systems of linear equations with symmetric positive-definite matrices. As described in Section we solve the linear systems with matrix Y, for blocks of queries in parallel:
D = Mis for i E {(r
+ +L..,,,,
(26)
In the FlashLLA forward algorithm [T CG is applied within each row block, with Qr. M,, m, Wr being the block quantities, X, being the solution to be computed and Y, being the right-hand side. We use the simplest initialization X, 0 for CG. Hence the initial residual r; is set to be ys. i.e..
(())) + Y, in Algorithm 2
The core computation is the matrix-vector product Σ,p, for the search vectors p; computed in the lines 10 The result is stored in matrix Ep. This operation has high I/O intensity due to the quirement to stream through the entire matrix HBM during each CG iteration. Consequently, controlling the number of iterations is crucial for both the efficiency and convergence. While maxi- mal iteration number < d can be manually set, further considerations are necessary to ensure the numerical stability and performance.
First, the convergence and convergent rate of CG are greatly influenced by its spectral condition. However, the conditioning varies significantly across positions. For example, for early tokens, the matrix Σ; is low-rank and requires relatively large À to maintain positive definiteness. To address this, we make the regularization learnable and data -ata-dependent:
A sigmoid(WAaz).
(27)
The dimension of the weight WA E ddxds controls the granularity of the regularization. Setting d, = d enables per-dimensional regularization, though empirically set d, d, suffices, where d, denotes the head dimension.
Additionally, since the CG solves multiple systems in parallel, different system converge different iterations. To prevent the numerical issues from affecting early converged system, we employ an active mask that disables iterations for systems whose residual norm fall below the tolerance
24

---

Published as a conference paper at ICLR 2026
Algorithm 2 FlashLLA CG Solver
Require: Variables K in HBM, X,. Yr, Qr, My, mrr in SRAM, block sizes Be, regularization X, tolerance e, max iterations bandwidth
1: Initialize on-chip: R( < Y, E B.xd xd p(U) + Y ER..d.
2: for porl=1 to 7 do
3: Initialize on-chip: Σp 0 R2.x
4: c to nn/B.) B. do
5: Load Kg from HBM to SRAM.
6: Compute W = K /h beast(m.)).
Compute Σ/ Ep
((e-1) 1) (W K Ke.
7:
8:
end for
9:
Compute Pg brsum y, Q+) and PM brsum( o My).
Compute 2 ast) = 1 (1(ast)) - -o @ M PM PM @ Q + boast @ PooQ, +P((-)
10:
11:
Compute rsum (1 @ R((11)) and check convergence.
Compute a n/rsum( ((t-1) @ Ep )
13:
Compute (²) = ((--1) + boast (a) @ P(t-1)
14:
Compute R(O) R(t-1) bcast(a)) o 2p
y-(last)
15:
Compute 3 rsum ((() //n)
16: Compute p(t) R(t) @ P(t-1).
17:
end for
return X, the solution of Σ;X- Y, for i the block.
D APPENDIX: BACKWARD DERIVATION
This section provides the detailed derivation of the backward pass. Defining the following variables:
b.
Uz
Wij
Thi
Ujj
9:
OL
E
IRd
(28)
8,
ab,
j=1
j=1
0C 0sig
Wi
M/ = 9, 03
3
Sigs
Cjj
(29)
0 j=1
8:
The gradient of the loss with respect by is given by:
OC 0o,
n Σ
sig9e
Av=STG
(30)
The gradient of qi and ky is related to sij. which can be broken down into wy and 25 path in the computation graph. The partial gradients of the loss with respect to Wijs 2ij are given by:
Oc
OL Osij
+ OC Ows
0p: Dwij
OC T 0Σ; J.
Tr
(31)
Owi
Jsij Owij
Ow; Owij
Twiy
Mip
T OL
0L
B, + 24 0pm
+ 2/j 0Σ, 2ij
(32)
8,
ac
a 0sis
0c 0μ:
0L 0E.
Ozs
=
asú 0zi
+ 0hf
02g
Tr
Dwsy
(33)
ac + 2w'sj
aL
= -cgipt + Wij 0ps
E 2ij
(34)
25

---

Published as a conference paper at ICLR 2026
Then the partial gradients of the loss with respect to ky and qi given by:
dc
dL owij
0C Ozij
0C wij Jwii h qb +
dc
+
(35)
0k.
Owij Jkg
Ozij 0k:
Ōzij
0C
dL owij
ac Ozij
0C Droi h
OL wij kj kzis
+
(36)
dq:
Swij
dq.
Ozig dq:
In order to compute the partial gradients of μ4 and E,, denote
Σ
Cig _n
Σ Caski
Σ Czj 2i
(37)
j=1
which can be computed with existing Conjugate Gradient solver
Then
0C =
Oc dmg
OC 06,
= -14 + 23.pps
(38)
0p.
==1 Jna 0p:
06, 8p:
OL Onij
dC a5;
1 OL T
1
06, 0Σ; 2 8ps
+ usp
(39)
aΣ,
Jniy OΣ;
2
We denote the following variables:
4 = U + 2Bcast(B) R
(40)
We can materialize the gradient of wey for every i, pair and then perform the reduction.
dc
OC
OC
Wij Jwij
8i + Way 0p:
+ Wiy 24 0Σ; Zsj
(41)
We omit the Q, K in the relmm for simplicity, then the gradient of wij can be computed as:
1
Aw T W (-bcast(B) relmm(Aw) 2 geelm(AA) C relmm(A)
(42)
1
2 relmm(U) o relmm(1)) (R))
(43)
Then compute the gradient of k, and q, through wy branch as follows:
A K L TrQ. A
1 h
Aw K
(44)
For the zij path, we avoid materializing the third-order tensor by performing the reduction internally,
OC
ac
Cyp: +
Wif 0p:
+ 2
Wid
(45)
Ozig
i=j
i=j
OL
OC
ac
Σ
Cipps +
Wii 0pt
+ 2
Wij
(46)
Oziz
/=1
j=1
1=1
Then we can compute the gradient of the loss with respect to 2ij as follows:
A/4 C R R A. - (W (elmm(A,,)) R + (W relnm(R))
TR
(47)
4 = brruu(u))))) (C) @ R+ brsum(W)) 3 A.
(48)
brsum(W W relmm(A,)) O R+ brsum(W @ re)n
(49)
Hence the gradient of the loss with respect ky and q can be computed
AK AK + K =0= +42
(50)
26

---

Published as a conference paper at ICLR 2026
Paranente Dimetammee Segmen
Stgzest Rea
Sumn See to
Samr Sectd
Suneet 153
Samer Sax
Samat Str
Sarast Ser v0l
Figure 6: Test-time regression with across all input dimension and segment size S.
E APPENDIX: COMPLEXITY ANALYSIS
In this section, provide detailed discussion of the computation and IO complexity of FlashLLA and its sparse variants, together with empirical profiling results and numerical precision analysis of FLASHLLA kernel. This complements the main paper by offering concrete guidance how computation and memory can be reduced in future work.
E.1 COMPUTATION AND IO COMPLEXITY
LLA extends the Softmax attention operator by performing regression-style computation across both sequence and feature dimensions. The additional cost arises primarily from solving a d x d linear system per query. Therefore, the computation LLA can be amortized over sequential dimensions and feature dimensions. Specifically, we can consider two strategies for reducing the computational and memory cost of LLA:
1. Sliding- Window LLA (SW-LLA): Applying sliding window reduces both computa- tion and IO, particularly for CG-based solvers. Theoretically, local linear regressors has improved sample efficiency compared to local constant methods. Hence, local linear atten- tion has the potential to achieve the same error level with restricted window size.
2. Block-Diagonal LLA (BD-LLA): The full LLA formulation captures all feature correla- tions, which is the dominant source of computational and memory cost and often requires ridge regularization ensure positive definiteness. Imposing block-diagonal structure re- duces both compute and memory while relaxing conditioning requirements. Suppose the block-diagonal structure consists of g diagonal blocks, each of size d/g. We consider the following two specifications:
g d: only per-feature scaling is modeled. This introduces only O(nd) memory and element-wise multiplications beyond Softmax attention. CG is unnecessary.
= d/2: pairwise feature correlations are modeled. This requires inverting 2 x matrices per block, which is computationally inexpensive and well-conditioned in practice. CG is unnecessary in this regime.
These dimensional sparsity strategies often eliminate the need for CG entirely while retain- ing much of the modeling flexibility of full LLA.
We summarize the per-query complexity of representative methods in Table | Here, n denotes the sequence length, d the head dimension, T the number of CG iterations (if used), g the number of diagonal blocks in block-diagonal LLA (BD-LLA), and B the block size used in the kernel.
E.2 LATENCY PROFILING
27

---

Published as a conference paper at ICLR 2026
Method
Computation
10
SW Support
Softmax Attention
(nd) O(Td)+ (dP) O (nd) O + O(nd)
2nd/B ((d)
MesaNet
X
LLA
(T + 31+33dd/1 B +T mndd B
BD-LLA (CG)
BD-LLA (direct)
O + O(nd)
3nd/ B
Table I: Per-query computation and IO complexity.
We profile latency (ms) and peak memory (GB) during prefill with head dimension 128 and batch size 32. Results are summarized Figure 7 FLASILLA reduces working memory and improves latency by 50x-200x over the naive implementation depending on sequence length and CG iterations. Notably, naive LLA runs out- of-memory beyond length 4096, while FlashLLA implementation scales to 8192 tokens with moderate memory growth.
Nuave LLA
10.
FlviLLA
LisHTAA (Tii1))
Lakurn (ne))
(T-12)
Atrention
We also evaluate numerical accuracy by comparing our BF16 CG kernel against a naive FP32 implementation that uses torch. 1 meiinalg.... solve. Table 2 reports absolute and relative error versus CG iteration count.
Figure 7: Prefill latency scal- ing (log-log) with dimension 128 and batch size 32.
Relative error decreases rapidly and reaches proximately 16 iteration. Additional iterations yield negligible improvements, indicating that a small constant iteration budget is sufficient in practice.
floor ap-
T
1
2
4
8
16
32
64
Absolute Error
1.49e-3
4.51e-4 0.651
9.13e-5
1.98e-5
1.72e-5
1.72e-5
1.72e-5 0.011
Relative Error
1.898
0.129
0.016
0.011
0.011
Table Numerical precision of BF16 CG kernel versus FP32 solve.
For comparison, if the naive implementation uses BF16: (1) Casting all computations to FP32 yields relative error 0.0173; (2) Casting to FP32 only before computing the final solution yields relative error 0.0174. Our BF16 CG kernel therefore achieves significantly better numerical stability than naive mixed-precision implementations. also suggests that the CG iterations could reduced further without sacrificing much accuracy, which would yield additional efficiency gains.
F APPENDIX: EXPERIMENTS
F.1 PIECEWISE LINEAR DATA GENERATION.
Let 2m be the number of segment with n log, n S For each section index c E {1...., n). define a sign pattern S. (st,1, Se.2...,, se,m) +1, by reading the liint-significant bits of c. For each segment in each sample, draw Z ~ NN0, Ta) and construct the data X E RR by flipping the first m coordinate of Z:
X, Sec Zyl Zz
j S m j > m
(51)
As the result, the constructed segment is a truncated Gaussian conditioned to lie the cone C.
{o E Idd Sajiy 2 0. = 1,...,m, m. where C, nC = b. Denote Tz(Z) = (S. @ 77 I 7 and segment distribution P., we have
X X Td(7)
(52)
28

---

Published as a conference paper at ICLR 2026
F.2 TRAINING CONFIGURATION
Time Regression. This experiment evaluates test-time adaptation without training any model parameters. We sweep over input dimension d E {8. 16, 32, 64, }28} and segment size S e {8, 16, 32, 64, 128, 256, 512, ]024] with fixed sequence length L = 1024. Performance is evalu- ated by averaging the mean squared error over 10,000 independently generated sequences.
In-Context Regression. We fix the input and output dimensions dg = d, 32. For each ran-
dom seed, we generate 100,000 training examples with noise level = 0.1 and 1,000 test examples with 5 0 (noiseless evaluation). For each sequence length L E (64, 128, 256, 512), over segment size S E [L/8, L/4, L/2, L} and learning rate (5 x 10 10-5, 10 10 x 10 10 *). All models are trained with the Adam optimizer (21,5g) (0.9, 0.999), weight decay 0.1, batch size 256, for a maximum of 100 epochs.
sweep
In-Context Associative Recall. We fix the vocabulary size U | = 8,192 For each sequence length L. E {64, 128, 256, 512), sweep over the number of key-value pairs {L/16, T,88, 1./4) and learning rate {10-4 4.5 x
10. We generate 20,000, 40,000, and 60,000 training exam- ples and 1,000 test examples each for the respective key-value pair counts. Training uses AdamW with (31, 82) = (0.9, 0.999), weight decay 0.1, batch size 256, for maximum of 32 epochs. Short convolution and feature map are disabled for all models.
Permutation State Tracking. We fix the vocabulary size |A| = 8,192. For each random seed, we generate 100,000 training examples and 1,000 test examples. For each position count N E {16, 24, 48, 96), sample the number of instructions S ~ Uniform(N/66 N/3) and use queries per example. We sweep over learning rate {10-a, 5 10-4, 10-3), "). Models are trained with Adam (81, 82) = (0.9, 0.999), weight decay 0.1, batch size 256, for a maximum of 64 epochs.
F.3 PRETRAINING ANALYSIS
This section supplements the main paper with additional pretraining results discussed in our rebuttal. Our goal is two-fold: (i) report evidence that LLA can yield consistent optimization benefits under settings that recover the kernel-regression perspective. and (ii) characterize qualitative differences in learned attention behavior when training with practical configurations.
Perplexity with Normalized Queries and Keys. When queries and keys are normalized, LLA recovers the kemeel-regression perspective most directly and empirically shows clear improvement over Softmax Attention in pretraining. In particular, our pretraining curves indicate that smaller ridge parameter À yields better learning curves and lower final perplexity. Concretely, we observe final losses of 3.287 (A=0.01), 3.421 (A=0.1), 3.509 (A=1), and 3.848 for standard Softmax At- tention (equivalently, LLA with X - oo). The results are shown Figure Sal While strong QK normalization appears to provide the most stable training in this regime, it can be overly restrictive and may hurt overall performance. Nonetheless, this normalization commonly adopted by linear- attention architectures (e.g., DeltaNet and MesaNet) as well as some Softmax Attention variants.
Attention Behavior Analysis. Under more practical training configurations without strong QK
normalization, we observe consistent and interpretable differences in attention behavior that suggest LLA induces a sharper and richer token interaction structure.
To quantify how broadly each token attends, we define an attention density metric as the number of positions whose absolute attention score exceeds a threshold (we use |siyl >> 0.1). We then average this density across tokens, heads, and layers. We find that Softmax Attention typically concentrates heavily 1-2 tokens average, while LLA more often attends to 2-4 tokens across most layers, indicating denser contextual aggregation. The results are shown in Figure Sb
Although LLA enforces a normalization constraint analogous to attention weights summing to one, it does not constrain the scores to be nonnegative. In practice, LLA exhibits both negative attention scores (as low -0.5) and larger positive scores (up to ss 1.8), whereas Softmax Attention score is restricted to be within (0, 1]. Signed contributions allow the model to subtract irrelevant components and amplify useful directions. The results are shown in Figure Sc
29

---

Published as a conference paper at ICLR 2026
Training Perplexity
Deusity of Attention Matrix
Range of Attention Score
LLA
Altn bilx
Altu
Perplott
LLA (A (100)
Max Min
Arto kilk
1.8
000
10000
Training Step
J00000
150000
303030
2.5
Lager
Position
(a) LLA Pretraining Perplexity.
(b) Attention Map Density,
(c) Attention Score Range.
F.4 ADDITIONAL EXPERIMENT RESULTS
In Figure 6 we provide the full test-time regression results across all input dimension d and seg- ment size S. The advantages of LLA scales with the dimension d and nonstationarity. In practical settings, as the dimensionality increases, is less likely do exact query q k2 as in this synthetic experiment. Consequently, kernel selectivity becomes less pronounced when noise is present in query points, limiting the potential advantages of both softmax attention and LLA compared the ideal conditions of this synthetic experiment. Nevertheless, the overall performance trends remain consistent with our main findings.
Figure 9 shows complete associative recall results across all sequence lengths L. For visual clarity, average results across different numbers of key-value pairs for each sequence length and plot the trajectory of the best score for each model. Gated DeltaNet exhibits distinctive training dynamics characterized by an extended plateau phase with minimal loss improvement, followed by an abrupt transition to significantly lower test loss and corresponding rapid accuracy improvement. The timing of this transition highly sensitive hyperparameters such as learning rate and dataset size.
In contrast, LLA demonstrates consistent, gradual improvement in test accuracy with correspond- ing smooth decrease in test loss throughout training. mirroring the behavior of Softmax Attention but more powerful. This stable convergence pattern remains robust across a wide range of learn- ing rates and dataset sizes. The marked difference in optimization dynamics suggests fundamental differences in how these models navigate the loss landscape and converge to solutions.
Test Reaal. Accurary
Lotil Soalee42
Lod32 Sendwa2
0.5
0.5
0.9
0.5
0-
11.
0.2
0.2
661 -11
1.00 Saedl
Saed
Send
0.5
0.5
0.5
0.5
8.1
0-
114
Seeel-11 -11
1.5
Saedl 11
1.0
Stedi-11
1.0
Send--0
11..
155
0.5
0.8-
0.5
0-
0.4
Figure 9: Test accuracy curves for associative recall across all sequence lengths L (averaged over random seeds). Results for different numbers of key-value pairs are averaged within each sequence length for visual clarity
30


Reviews:
[{"review_id": "Sexzl0TtPr", "reviewer": "Reviewer_Hn3s", "summary": "The paper analyzes the characteristics and performance of softmax and linear attention under the Test-Time Training (TTT) setting from the perspective of statistical regression. From this viewpoint, the authors propose a novel mechanism named Local Linear Attention (LLA), which they claim demonstrates superior advantages. Furthermore, the paper introduces an I/O-aware operator for LLA, designed to achieve more efficient memory utilization. The effectiveness of the proposed model is validated through toy experiments, which showcase its capabilities.", "questions": "1.\tThe paper's second claimed contribution, which asserts to have resolved the quadratic memory complexity, is fundamentally flawed. This is not a novel contribution. The issue of quadratic memory cost in standard self-attention was already solved by seminal works 3 years ago, such as FlashAttention. Furthermore, the authors' comparison is irrelevant for linear attention, which inherently possesses O(Nd) memory complexity due to its operator properties and never had a quadratic memory bottleneck to begin with. Claiming this as a contribution demonstrates a misunderstanding or misrepresentation of the current state-of-the-art in attention mechanisms. In other words, this \"contribution\" merely solves a problem that the proposed LLA method creates for itself.\n2.\tThe paper's second claimed contribution is achieving strong results on artificial \"toy\" experiments. However, the complete absence of validation in a pre-training setting is a critical omission. In the field of attention architecture research, evaluating a new mechanism's performance on large-scale pre-training has been a standard practice for years; failing to do so is highly unusual, even by the standards of a year ago. This limitation severely undermines the claim of \"good results,\" as it is unclear if the performance gains can generalize beyond these contrived scenarios. It gives the impression that the proposed model is only capable of solving problems that are, in a sense, artificially constructed for it.\n3.\tThe computational complexity of LLA raises serious questions about the paper's contribution. The authors state that LLA's complexity surpasses that of softmax attention. This is deeply ironic. For nearly a decade (since circa 2017), the community has dedicated immense effort to making attention mechanisms, particularly linear attention, more efficient. This paper, however, manages to achieve the opposite. In a surprising turn of events, this work presents a \"linear attention\" variant that is less efficient than the very baseline (softmax attention) it sought to outperform. One could sardonically remark that after all these years, this paper finally allows linear attention to generate a larger carbon footprint than softmax attention.\n4.\tThe paper would benefit from a more comprehensive complexity analysis. The proposed operator is unique, and a detailed breakdown of its computational and memory complexity would significantly improve clarity for the reader. I would encourage the authors to include a more thorough analysis to help the audience better understand the operator's behavior and trade-offs.", "limitations": "", "rating": 2, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The paper provides a solid theoretical foundation for its proposed method. The authors conduct a detailed analysis of the properties of linear and softmax attention from a statistical regression perspective. This approach is well-motivated and naturally leads to the formulation of their Local Linear Attention (LLA) framework.", "weaknesses": "A major weakness of this work is the limited scope of its experimental validation. The authors didn’t demonstrate the viability of Local Linear Attention (LLA) for pre-training, a critical aspect for assessing new attention mechanisms. Moreover, the paper completely has no evaluations on any practical downstream tasks, confining its results to toy experiments. Consequently, the practical relevance and usefulness of this research are questionable, as it is unclear whether the proposed method offers any real-world benefits.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OIpnoD0ZEq", "reviewer": "Reviewer_HAFq", "summary": "This paper presents a novel and insightful theoretical framework for understanding attention mechanisms through the lens of non-parametric statistics. The authors convincingly argue that standard Softmax Attention is analogous to a local constant (Nadaraya-Watson) estimator, which is known to suffer from theoretical limitations such as bias at data boundaries.\n\nAs a principled alternative, the paper proposes Local Linear Attention (LLA), a new attention mechanism derived from local linear regression. This model is theoretically superior, offering faster asymptotic convergence rates and mitigating the boundary bias issues inherent in the local constant approach.", "questions": "Please refer to my weakness part.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper's greatest strength is the novel connection it establishes between attention and non-parametric regression. Framing Softmax Attention as a local constant estimator is an insightful conceptual leap that provides a new lens for the community to understand and analyze attention.\n-  LLA is not an ad-hoc modification of attention. It is a new mechanism derived directly from a statistically superior estimator (local linear regression). This principled design, which provably addresses issues like boundary bias (as shown in the appendix), is a significant strength.\n- The authors clearly recognized the computational hurdles of LLA and put significant effort into addressing them.", "weaknesses": "- The primary experimental evidence comes from a synthetic piecewise-linear regression task. While this task perfectly isolates the theoretical benefits of LLA, it feels somewhat circular: it proves the local linear model is good at solving a local linear problem. It is not immediately obvious if this specific capability (handling sharp, non-stationary linear segments) is a major bottleneck in more common, large-scale tasks like natural language processing. The paper would be much stronger if it could demonstrate this benefit on a more complex, real-world benchmark where this specific type of non-stationarity is known to be an issue.\n- The \"FlashLLA\" name and its benchmarking could be misleading. The name invites a direct comparison to FlashAttention, which is a one-pass I/O-aware algorithm. However, FlashLLA relies on a Conjugate Gradient (CG) solver, which is iterative. As detailed in Algorithm 2, each iteration of the CG solver requires streaming through the Key matrix, resulting in a much higher I/O cost (multiplied by the number of iterations, $T$). The paper's main hardware benchmark (Figure 1) only plots memory usage and omits wall-clock runtime or throughput comparisons against a standard FlashAttention baseline. This makes it impossible to assess the true computational overhead. The authors admit in the limitations that the constant factor is \"higher,\" but this seems to understate the potential I/O bottleneck from the iterative solver.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "mYUVMXz2bX", "reviewer": "Reviewer_tiaX", "summary": "This paper introduces Local Linear Attention (LLA) as an enhancement to softmax attention. Within the test-time regression framework, the local linear regression objective in LLA can be viewed as an optimal interpolation between linear and softmax attention, yielding improved associative recall capabilities in terms of convergence behavior and boundary bias.\n\nThe authors further optimize LLA’s implementation through targeted improvements to ensure linear memory scalability and propose FlashLLA, a blockwise parallel and hardware-efficient algorithm. Empirical results across four synthetic tasks, e.g., test-time adaptation and in-context learning, demonstrate the method’s efficacy compared to existing attention mechanisms.", "questions": "1. In the regression formulation of LLA, why is the local function instantiated as $W(x-q)+b$? Is this particular form necessary to derive a “well-behaved” solution? I wonder the constraints and considerations behind this choice, and whether alternative linear functions exist that could also work.\n\n2. Please correct me if my understanding is inaccurate: within the proposed regression paradigm, locality is reflected in both the loss function  $f(x)$ and the weighting factor $w_{ij}$. For the former, softmax attention employs a constant value function, whereas LLA uses a parameterized local linear function. Could you clarify the distinct roles of these two forms of locality, and are both components essential to the design of LLA?\n\n3. In the test-time regression experiment (Fig. 2), does MesaNet use decay? It appears that the Attn model fails to effectively utilize in-distribution information—does this imply that parameterization and locality in the loss function are necessary? Moreover, why does MesaNet outperform LLA in the first segment but deteriorate significantly in later segments, and could data-dependent decay influence this observation?\n\n4. In the MQAR experiment (Fig. 4b), why does the in-context recall of Attention perform substantially worse than GDN or Mamba in the short-sequence (<256) setting? Intuitively, for a token-wise matching task of this kind, Attention should be able to learn such dependencies quite easily.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "1. The paper is well-organized and clearly written, providing readers with a solid understanding of the test-time regression framework and the motivation behind LLA. The presentation of the LLA algorithm is also systematic, including its regression, matrix-parallel, and blockwise form.\n\n2. The authors demonstrate a deep understanding of the problem, offering thorough theoretical analysis and proofs. The resulting local linear regression design is conceptually novel and well-justified.\n\n3. The synthetic experiments indicate that LLA exhibits a distinct superiority over other attention mechanisms in certain abilities, including in-context regression and associative recall.", "weaknesses": "1. The experiments are conducted only on synthetic tasks. While I understand that LLA introduces additional complexity in both computation and memory I/O at the kernel level, it remains uncertain how well results from such small-scale, controlled tasks can transfer to real-world LM capabilities.\n\n2. The paper provides limited analysis and interpretation of the experimental results (except for the state-tracking task). For instance, it is confusing why MesaNet outperforms LLA in the first segment of the test-time regression task.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ND1WYpPANA", "reviewer": "Reviewer_Uxja", "summary": "This paper builds upon the test-time regression interpretation of attention mechanisms proposed by Wang et al. (2025) and introduces a new architecture called Local Linear Attention (LLA), which aims to improve upon softmax attention. Specifically, the authors formulate LLA as the solution to local linear regression under the test-time regression framework, claiming that it achieves superior performance compared to softmax attention (corresponding to local constant regression) and MesaNet (corresponding to global linear regression). Furthermore, they propose a practical GPU-optimized algorithm to alleviate the high computational cost of LLA and make it efficient in practice.", "questions": "1. How does the local polynomial regression mentioned in line 197 relate to the proposed method?\n2. In line 236 (and in the title), what exactly does “optimal” refer to? In what sense is it optimal?\n3. Regarding line 470: the paper states that the conjugate gradient method is sufficient with a constant number of iterations. Is there any theoretical justification or empirical evidence supporting this claim quantitatively?", "limitations": "", "rating": 2, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The paper presents a theoretically motivated architectural design.\n- It proposes a practical, GPU-aware implementation of the algorithm.", "weaknesses": "1. The validity of the proposed architecture heavily relies on theoretical assumptions specific to in-context learning and associative memory. Moreover, the theory depends on test-time regression framework, which provides only a partial understanding of one aspect of the attention mechanism. While theory-driven design is sometimes valuable, the current theoretical justification is not comprehensive enough to fully validate the effectiveness of the proposed architecture and superiority over softmax attention.\n2. The experimental validation, like the theory, is limited to settings where the test-time regression interpretation applies — namely, in-context learning and associative recall. To convincingly justify a new architecture, experiments in a practical scenarios is necessary.\n3. The paper does not include comparisons of computational cost (time and memory) with softmax attention, linear attention, or MesaNet. Beyond demonstrating efficiency improvements via FlashLLA, the paper should also show how competitive the proposed method becomes relative to existing approaches.\n4. The description “it suffers when predicting near the boundary of the data support, particularly with symmetric kernels like RBF” in line 194 is unclear. Since this point seems to be crucial to differentiating the proposed method from softmax attention, a more detailed explanation is needed.\n5. Many references are currently to arXiv preprints, even though several of them likely have published versions in conferences or journals. Please consider citing the published versions instead.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)