Local Linear Attention: An Optimal Interpolation of Linear and Softmax Attention For Test-Time Regression

## Abstract

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 Θ ​ ( n 2 ​ d ) \Theta(n^{2}d) and Θ ​ ( n ​ d 2 ) \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.
Code is available at https://github.com/Yifei-Zuo/Flash-LLA .

## 1 Introduction

Transformer-based architectures have dominated modern AI systems and powered breakthroughs across various fields.
The core computation primitive–Softmax Attention (Vaswani et al., 2023), adaptively processes and aggregates contextual information.
Recent research on architecture innovation has proposed numerous variants of attention mechanism such as Linear Attention (LA) (Yang et al., 2025b, a; Siems et al., 2025; Sun et al., 2023; Ma et al., 2024; Liu et al., 2024; Katharopoulos et al., 2020; Yang et al., 2024) and state space models (SSMs) (Gu et al., 2022a, b; Gu and Dao, 2024; Dao and Gu, 2024; Poli et al., 2023).
While these methods offer remarkable efficiency improvements on long sequences, they often incur a performance penalty compared to Softmax Attention (Bick et al., 2025; Jelassi et al., 2024).
Meanwhile, many of the design choices such as gating and forgetting factor (Yang et al., 2025a; Lin et al., 2025; Gers et al., 2000; Yang et al., 2023) are often guided by heuristics or empirical results, lacking a principled understanding.
Conversely, the search for more expressive attention mechanisms, even at an additional computational cost, has been relatively 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 a test-time optimizer for layer-specific regression problems.
A natural question arises: *can we systematically improve the Softmax Attention mechanism from the perspective of test-time regression?*

In this work, we propose *local linear attention (LLA)*, an upgrade to 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 $\Theta(n^{2}d)$ and $\Theta(nd^{2})$ memory complexity with two optimizations, where $n$ is the sequence length and $d$ is the dimension.
We then introduce FlashLLA, a blockwise and hardware-efficient algorithm to parallelize the computation on modern accelerators.
In addition, we implement and profile a 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 limitations of LLA compared to strong baselines.

Figure: Figure 1: A comparison of regression strategies: global linear models (e.g., SSMs, MesaNet, DeltaNet) employ query agnostic linear fits and suffer from irreducible approximation error due to model misspecification; local constant models (e.g., Softmax Attention) perform query-specific local averaging but exhibit boundary bias; local linear models (e.g., LLA) achieve a superior bias-variance trade-off combining locality and linear fitting.
Refer to caption: figures/front-lla.png

### 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., 2025b, a; Siems et al., 2025; Sun et al., 2023; Ma et al., 2024; Liu et al., 2024; Yang et al., 2024) and SSMs (Gu et al., 2022b, a; Gu and Dao, 2024; Dao and Gu, 2024; Poli et al., 2023) were proposed in search of more efficient alternatives for long-context sequence generation.
These methods maintain a constant sized hidden state (Katharopoulos et al., 2020) during decoding and update it like a linear RNN.
This state update behavior is also referred as fast-weight programming (Schmidhuber, 1992; Schlag et al., 2021).
Essentially, MesaNet (von Oswald et al., 2025) is a LA variant that preconditions the hidden state and achieves optimal regression objectives among linear models.

#### In-Context Learning by Optimization.

A growing body of work suggests that attention mechanism implicitly performs optimization algorithms to achieve in-context learning (Garg et al., 2023; Akyürek et al., 2022; von Oswald et al., 2023; Kirsch et al., 2024; Zhang et al., 2024; Mahankali et al., 2023; 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 a 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, 2023) performs a block-wise online softmax to reduce I/O latency.
Several other approaches, including NSA (Yuan et al., 2025), SeerAttention (Gao et al., 2024), MoBA (Lu et al., 2025), and Block Sparse Attention (Xiao, 2025), leverage sparsity to lower the effective computational cost while preserving GPU utilization.
Flash Linear Attention (Yang et al., 2023) provides a hardware-friendly formulation of linear attention through chunk-wise computation.

### 1.2 Notation.

We use upper-case letters to denote matrices and lower-case letters to denote vectors.
For a matrix $X$, we denote its Frobenius norm as $\|X\|_{F}$ and the Hadamard product as $\odot$.
For a vector $x$, we denote its Euclidean norm as $\|x\|_{2}$.
Furthermore, define $\texttt{rsum}(X)=X\mathbf{1}$ for matrix $X$ and $\texttt{bcast}(x)=x\mathbf{1}^{\top}$ for vector $x$, where $\mathbf{1}$ is a vector of ones.
We use the abbreviation $\texttt{brsum}(x)=\texttt{bcast}(\texttt{rsum}(x))$.

## 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 limitations of LA and Softmax Attention.
Lastly, we introduce the formulation for LLA.

### 2.1 Attention as Test-Time Regression

In test-time regression framework, the attention mechanism is interpreted as a layer-specific regression solver.
The goal is to approximate an unknown regression function $f:\mathbb{R}^{d}\mapsto\mathbb{R}^{d}$ using historical key-value pairs.
To be specific, given a hypothesis space $\mathcal{F}$ and a position $1\leq i\leq n$, an estimator $\hat{f}_{i}\in\mathcal{F}$ is fitted on the dataset $\mathcal{D}_{i}=\{(k_{j},v_{j})\in\mathbb{R}^{d}\times\mathbb{R}^{d}\}_{j=1}^{i}$, where the attention keys $k_{j}$ serve as the features and attention values $v_{j}$ as the labels.
The prediction is made at a query $q_{i}\in\mathbb{R}^{d}$, which is treated as a test data point.

#### Linear Attention as Parametric Regression.

Parametric model constrains the function class $\mathcal{F}$ to a set of functions defined by a finite-dimensional parameter $\theta\in\Theta$.
The most fundamental instantiation is the linear regression, which sets $\mathcal{F}=\{f_{\theta}(x)=Wx+b\mid\theta=(W,b),W\in\mathbb{R}^{d\times d},b\in\mathbb{R}^{d}\}$.
We omit the intercept $b$ for simplicity. At each position $i$, the parameter $W_{i}$ is estimated by solving the following least square problem on the training dataset $\mathcal{D}_{i}$,

$$ $\displaystyle\min_{W}\mathcal{L}(W;\mathcal{D}_{i})=\frac{1}{2}\sum_{j=1}^{i}\gamma_{ij}\|v_{j}-Wk_{j}\|^{2}_{2}+\lambda\|W\|^{2}_{F},$ (1) $$

where $\gamma_{ij}\in\mathbb{R}$ is a weighting factor and $\lambda\geq 0$ is the ridge regularization penalty.
For appropriate regularization, objective equation [1](https://arxiv.org/html/2510.01450v1#S2.E1) admits a closed-form optimal solution. MesaNet (von Oswald et al., 2025, 2024) hardcodes this solution for the case $\gamma_{ij}=1$, where the prediction is given by,

$$ $\displaystyle\hat{f}_{\texttt{Mesa}}(q_{i})=\hat{W}_{i}^{\texttt{Mesa}}q_{i}=\biggl(\underbrace{\sum_{j=1}^{i}v_{j}k_{j}^{\top}}_{S_{i}}\biggr)\biggl(\underbrace{\sum_{j=1}^{i}k_{j}k_{j}^{\top}+\lambda I}_{H_{i}}\biggr)^{-1}q_{i}.$ (2) $$

Across different position $i$, the weight $W_{i}^{\texttt{Mesa}}$ can be updated recurrently by maintaining two statistics $S_{i}$ and $H_{i}$ with $\Theta(d^{2})$ 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_{i}\approx I$ in equation [2](https://arxiv.org/html/2510.01450v1#S2.E2), leading to a suboptimal solution of the least square problem.
It can also be shown that LA variants and SSMs such as GLA (Yang et al., 2024), RetNet (Sun et al., 2023), RWKV (Peng et al., 2024, 2025) and Mamba (Gu and Dao, 2024; Dao and Gu, 2024) can be derived by the approximation $H_{i}\approx I$ with different weighting schemes $\gamma_{ij}$.
Besides exact solutions, vanilla LA and variants such as DeltaNet (Yang et al., 2025b) and Gated DeltaNet (Yang et al., 2025a) can be interpreted as performing one step of first order stochastic gradient descent on weight $W$.
We refer readers to (Wang et al., 2025) and tables in (Peng et al., 2025; Yang et al., 2025a) for detailed derivations and how each model is implemented.

#### Softmax Attention is Non-Parametric.

Non-parametric regression makes minimal structural assumptions on the function class $\mathcal{F}$. A canonical example is the kernel regression, where the estimator $\hat{f}$ is defined directly from the data in a way that depends on the query point.
Specifically, for a query vector $q_{i}$, let $\mathcal{F}(q_{i})$ denote the local function class around $q_{i}$. The local regression objective is defined as:

$$ $\displaystyle\min_{f\in\mathcal{F}(q_{i})}\mathcal{L}(f;\mathcal{D}_{i})=\frac{1}{2}\sum_{j=1}^{i}w_{ij}\|v_{j}-f(k_{j})\|_{2}^{2},$ (3) $$

where $w_{ij}=K_{h}(q_{i},k_{j})\in\mathbb{R}$ is a query-dependent weight that measures the locality of the training point $k_{j}$ to the query $q_{i}$.
The simplest instantiation is the constant model, where $\mathcal{F}(q_{i})=\{f_{\theta}(x)=\theta\in\mathbb{R}^{d},\forall x\in\mathbb{R}^{d}\}$.
Solving objective equation [3](https://arxiv.org/html/2510.01450v1#S2.E3) with this function class yields:

$$ $\displaystyle\hat{f}(q_{i})=\sum_{j=1}^{i}s_{ij}v_{j},\quad s_{ij}=\frac{w_{ij}}{\sum_{j^{\prime}=1}^{i}w_{ij^{\prime}}}$ (4) $$

Consider an RBF kernel $w_{ij}=\exp(-\|k_{j}-q_{i}\|^{2}/h)$ with bandwidth $h=2\sqrt{d}$, the estimator equation [4](https://arxiv.org/html/2510.01450v1#S2.E4) exactly recovers the softmax attention when QK normalization (Dehghani et al., 2023; Wortsman et al., 2023; Team, 2024) is applied,
since in this case $w_{ij}\propto\exp(q_{i}^{\top}k_{j}/\sqrt{d})$ and the common constant factor cancels in the division.
It is also known as the Nadaraya-Watson (NW) kernel regression (Nadaraya, 1964; Watson, 1964; Bierens, 1988) in statistics literature.
We note that 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 $h$ in a data-dependent manner.

### 2.2 Learning Behavior of Associative Recall

Attention mechanisms are often evaluated by the associative memory capacity (Zhong et al., 2025; Behrouz et al., 2025; Ramsauer et al., 2021).
Specifically, given a training set of key-value pairs $\{(k_{j},v_{j})\}_{j=1}^{i}$, the model is expected to retrieve the value $v_{j}$ associated with $k_{j}$ when queried at $q=k_{j}$.
This objective can be exactly captured by the Mean Square Error (MSE). For example, the retrieval error in vanilla LA is given by,

$$ $\displaystyle\mathrm{MSE}_{i}^{\texttt{LA}}=\frac{1}{i}\sum_{j=1}^{i}\|S_{i}k_{j}-v_{j}\|^{2}=\frac{1}{i}\sum_{j=1}^{i}\|(k_{j}^{\top}k_{j}-1)v_{j}+\sum_{j^{\prime}\neq j}v_{j^{\prime}}k_{j^{\prime}}^{\top}k_{j}\|^{2},$ (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 et al., 2025) quantitatively analyze this error by the inverse of signal-to-noise ratio $\text{SNR}^{-1}$, which is essentially a 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-variance trade-off.

#### Irreducible Approximation Error of Global Linear Model.

Recall that global linear models correspond to solving a least square problem equation [1](https://arxiv.org/html/2510.01450v1#S2.E1) over the hypothesis space $\mathcal{F}=\{f_{\theta}(x)=Wx+b\}$.
When the ground truth function $f$ is not global linear, any estimator $\hat{f}\in\mathcal{F}$ will suffer from a non-vanishing approximation error due to model misspecification.
In contrast, the local constant model is a nonparametric estimator that does not impose structural assumptions on the function class except for smoothness or regularity.
Consequently, the approximation error vanishes asymptotically with proper assumptions. In fact, we have the following separation result between global linear (GL) and local constant (NW) estimators:

###### Proposition 2.1 .

Let $(X_{i},Y_{i})_{i=1}^{n}$ be i.i.d., $X_{i}\in\mathbb{R}^{d}$ supported on a bounded set $D\subset\mathbb{R}^{d}$, and
$Y_{i}=f(X_{i})+\varepsilon_{i}\in\mathbb{R}^{d_{y}}$ with $\mathbb{E}[\varepsilon_{i}\mid X_{i}]=0$ and $\mathbb{E}[\varepsilon_{i}^{2}\mid X_{i}]=\sigma^{2}(X_{i})$.
Let $\widehat{f}_{\mathrm{GL}}$ denote a global-linear estimator and $\widehat{f}_{\mathrm{NW}}$ the local-constant (NW) estimator with optimal bandwidth.
Under mild assumptions, if $f$ is not globally linear, then

$$ $\mathbb{E}\int_{D}||\widehat{f}_{\mathrm{GL}}(x)-f(x)||^{2}dx=\Omega(1)\;,\;\mathbb{E}\int_{D}||\widehat{f}_{\mathrm{NW}}(x)-f(x)||^{2}dx=O(n^{-3/(d+3)}).$ $$

###### Proposition 2.1 .

#### Irreducible Boundary Bias of Local Constant Model.

Despite the appealing convergent property of local constant model, it suffers when predicting near the boundary of the data support, particularly with symmetric kernels like RBF.
This phenomenon becomes more pronounced in high-dimension and likely to occur more frequently in autoregressive prediction.
Local polynomial regression is a standard remedy in nonparametric statistics to address this issue. In Section [2.3](https://arxiv.org/html/2510.01450v1#S2.SS3), we will introduce LLA as a natural adaptation of local linear regression.
In fact, we have the following separation result between local constant (NW) and local linear (LL) estimators:

###### Proposition 2.2 .

Under the setting of Proposition [2.1](https://arxiv.org/html/2510.01450v1#S2.Thmproposition1), let $\widehat{f}_{\mathrm{NW}}$ and $\widehat{f}_{\mathrm{LL}}$ denote local-constant and local-linear estimators with their respective optimal bandwidths. Under mild assumptions, if $f$ has sufficiently large normal gradient along the boundary of $D$, then

$$ $\mathbb{E}\int_{D}||\widehat{f}_{\mathrm{NW}}(x)-f(x)||^{2}dx=\Omega(n^{-3/(d+3)})\;,\;\mathbb{E}\int_{D}||\widehat{f}_{\mathrm{LL}}(x)-f(x)||^{2}dx=O(n^{-4/(d+4)}).$ $$

The proof of proposition [2.1](https://arxiv.org/html/2510.01450v1#S2.Thmproposition1) and [2.2](https://arxiv.org/html/2510.01450v1#S2.Thmproposition2) are provided in Appendix [A](https://arxiv.org/html/2510.01450v1#A1) and [B](https://arxiv.org/html/2510.01450v1#A2) correspondingly.

###### Proposition 2.2 .

### 2.3 Local Linear Attention

#### Formulation.

For a query $q_{i}$, instantiate the local function class $\mathcal{F}(q_{i})=\{f_{\theta}(x)=b+W(x-q_{i})\mid\theta=(W,b),W\in\mathbb{R}^{d\times d},b\in\mathbb{R}^{d}\}$.
The regularized local linear regression objective is,

$$ $\displaystyle\min_{f\in\mathcal{F}(q_{i})}\mathcal{L}(f;\mathcal{D}_{i})=\frac{1}{2}\sum_{j=1}^{i}w_{ij}\|v_{j}-b-W(k_{j}-q_{i})\|^{2}+\lambda\|W\|^{2}_{F},$ (6) $$

where $w_{ij}$ are query-dependent kernel weights and $\lambda\geq 0$ is a ridge penalty.
This objective also admits a closed-form solution for the intercept $b$ and weight $W$. Importantly, at test time, the prediction is only made at $\hat{f}(q_{i})=\hat{b}_{i}$. Thus it suffices to derive the formulation for the intercept.
Define $z_{ij}=k_{j}-q_{i}$ and the following query-specific statistics,

$$ $\displaystyle\omega_{i}=\sum_{j=1}^{i}w_{ij}\in\mathbb{R},\quad\mu_{i}=\sum_{j=1}^{i}w_{ij}z_{ij}\in\mathbb{R}^{d},\quad\Sigma_{i}=\sum_{j=1}^{i}w_{ij}z_{ij}z_{ij}^{\top}+\lambda I\in\mathbb{R}^{d\times d}.$ (7) $$

Also denote $\rho_{i}=\Sigma_{i}^{-1}\mu_{i}\in\mathbb{R}^{d}$.
The optimal intercept can be computed as follows,

$$ $\displaystyle\hat{f}(q_{i})=\hat{b}_{i}=\sum_{j=1}^{i}s_{ij}v_{j},\quad s_{ij}=w_{ij}\frac{1-z_{ij}^{\top}\rho_{i}}{\omega_{i}-\mu_{i}^{\top}\rho_{i}}.$ (8) $$

Similar to maintaining $H_{i}$ in MesaNet, LLA requires capturing the precondition matrix $\Sigma_{i}$.
However, a key difference is that $H_{i}$ is built on global statistics that are independent of the query, whereas $\Sigma_{i}$ is constructed from features centered around the specific query $q_{i}$ for each position $1\leq i\leq n$.
Consequently, LLA requires a KV cache of size $\Theta(nd)$ similar to Softmax Attention, rather than constant-size recurrent states as in the LA family.

#### LLA Interpolates Linear and Softmax Attention.

A more interpretable form of equation [8](https://arxiv.org/html/2510.01450v1#S2.E8) can be obtained by decomposing the prediction into two components.
Suppose that the weight matrix $\hat{W}_{i}$ is given prior to solving equation [6](https://arxiv.org/html/2510.01450v1#S2.E6), then the optimal intercept can be expressed as,

$$ $\displaystyle\hat{b}_{i}=\sum_{j=1}^{i}s_{ij}(v_{j}-\hat{W}_{i}k_{j})+\hat{W}_{i}q_{i},\quad s_{ij}=\frac{w_{ij}}{\sum_{j^{\prime}=1}^{i}w_{ij^{\prime}}}.$ (9) $$

The first term is a local constant regression to predict the residuals $v_{j}-\hat{W}_{i}k_{j}$, while the second term is a linear prediction based on $\hat{W}_{i}$.
The formulation recovers LLA if $\hat{W}_{i}$ is obtained by optimally solving equation [6](https://arxiv.org/html/2510.01450v1#S2.E6).
However, by allowing suboptimal estimation, one can construct $\hat{W}_{i}$ as a recurrent state similar to LA.
This decomposition reveals how LLA interpolates between Linear and Softmax Attention and provides a template for designing new algorithms.

## 3 Practical Algorithm

In this section, we provide a detailed discussion of the computation involved in LLA.
We highlight two major challenges in naïve implementations and develop a practical block-wise algorithm FlashLLA that scales efficiently on modern accelerators such as GPUs.

### 3.1 Memory Efficient Primitives

#### Avoid Pairwise Materialization.

The first bottleneck is the evaluation of vectors $z_{ij}=k_{j}-q_{i}$ for every $1\leq j\leq i\leq n$, which requires $\Theta(n^{2}d)$ memory to materialize.
This pairwise difference is later used in the formulation of $\mu_{i}$ and $\Sigma_{i}$ defined in equation [7](https://arxiv.org/html/2510.01450v1#S2.E7) as well as the inner product $z_{ij}^{\top}\rho_{i}$ in equation [8](https://arxiv.org/html/2510.01450v1#S2.E8).
For both cases, explicit materialization of $z_{ij}$ can be avoid by algebraically separating the contributions of $k_{j}$ and $q_{i}$ to the final result.
Specifically, the statistics $\mu_{i}$ and $\Sigma_{i}$ can be reformulated in terms of intermediate quantities that are independent of the query and can then be transformed to recover the original centered statistics:

$$ $\displaystyle\tilde{\mu}_{i}=\sum_{j=1}^{i}w_{ij}k_{j}\in\mathbb{R}^{d},\quad\tilde{\Sigma}_{i}=\sum_{j=1}^{i}w_{ij}k_{j}k_{j}^{\top}+\lambda I\in\mathbb{R}^{d\times d}$ (10) $\displaystyle\mu_{i}=\tilde{\mu}_{i}-\omega_{i}q_{i},\quad\Sigma_{i}=\tilde{\Sigma}_{i}-\tilde{\mu}_{i}q_{i}^{\top}-q_{i}\tilde{\mu}_{i}^{\top}+\omega_{i}q_{i}q_{i}^{\top}$ (11) $$

The computation in equation [10](https://arxiv.org/html/2510.01450v1#S3.E10) and equation [11](https://arxiv.org/html/2510.01450v1#S3.E11) only requires vectors $k_{j}$ and $q_{i}$ individually, reducing the memory cost to $\Theta(nd)$.
The same principle applies to computing inner products of the form $z_{ij}^{\top}x_{i}=k_{j}^{\top}x_{i}-q_{i}x_{i}$ for any vector $x_{i}\in\mathbb{R}^{d}$.
The matrix operator relative matrix multiplication (relmm) for this optimization as follows,

$$ $\displaystyle\texttt{relmm}(X,Q,K):=XK^{\top}-\texttt{brsum}(X\odot Q).$ (12) $$

This operator is invoked once in the forward computation with $x_{i}=\rho_{i}$, but appears multiple times in the backward with other variables.
Further details on the backward are provided in Appendix [D](https://arxiv.org/html/2510.01450v1#A4).

#### Matrix-Free Inversion via Conjuagte Gradients.

The second bottleneck arises in solving linear systems of the form $\Sigma_{i}^{-1}x_{i}$ for some vector $x_{i}\in\mathbb{R}^{d}$. Directly inverting $\Sigma_{i}$ for every $1\leq i\leq n$ incurs a prohibitive $\Theta(nd^{2})$ memory.
Following the approach in MesaNet (von Oswald et al., 2025), we exploit the sum-of-rank-one structure of $\Sigma_{i}$ and solve the linear system iteratively using the conjugate gradient (CG) method (Hestenes and Stiefel, 1952).
The key insight is that CG only evaluate the matrix-vector product $\Sigma_{i}p$ for a search direction $p\in\mathbb{R}^{d}$ without explicit matrix materialization:

$$ $\displaystyle\Sigma_{i}p$ $\displaystyle=\sum_{j=1}^{i}w_{ij}(k_{j}^{\top}p)k_{j}-(q_{i}^{\top}p)\tilde{\mu}_{i}-(\tilde{\mu}_{i}^{\top}p)q_{i}+(\omega_{i}q_{i}^{\top}p)q_{i}+\lambda p.$ (13) $$

Each term only involves inner products or weighted sums over keys and the query, both of which can be computed efficiently using batched matrix multiplication with $\Theta(nd)$ memory.
This CG operation of $\Sigma_{i}$ is invoked once in the forward computation with $x_{i}=\mu_{i}$ and twice in the backward computation with other variables.
Further details on the CG algorithm are provided in [C](https://arxiv.org/html/2510.01450v1#A3).

### 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\in\mathbb{R}^{n\times d}$ 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 tril operator that preserves the lower-triangular matrix.
Then the output $O\in\mathbb{R}^{n\times d}$ can be computed as follows,

$$ $\displaystyle W=\texttt{tril}(\exp(QK^{\top}/h)),$ $\displaystyle\quad M=WK-\texttt{brsum}(W)\odot Q$ (14) $\displaystyle R=\texttt{CGSolve}(M,Q,K,\lambda),$ $\displaystyle\quad\delta=\texttt{rsum}(W)-\texttt{rsum}(M\odot R)$ (15) $$

$$ $\displaystyle O=$ $\displaystyle\left(\frac{1-\texttt{relmm}(R,Q,K)}{\texttt{bcast}(\delta)}\odot W\right)V,$ (16) $$

where $W\in\mathbb{R}^{n\times n}$ is the matrix of kernel weight $w_{ij}$, $M\in\mathbb{R}^{n\times d}$ stores the first-order statistics $\mu_{i}$,
$R\in\mathbb{R}^{n\times d}$ contains the solution to the linear systems $\Sigma_{i}^{-1}\mu_{i}$ for every $1\leq i\leq n$, and $\texttt{CGSolve}(\cdot)$ invokes the CG algorithm that construct $\Sigma_{i}$ implicitly and solve these systems in parallel.
The division and subtraction are performed element-wise.
The single-head computation can be naturally extended to multi-head the same way as in standard multi-head attention mechanisms.

Figure: Algorithm 1 FlashLLA Forward Pass

#### Blockwise Algorithm.

Denote $B_{r},B_{c}$ as the block size for queries and keys/values along the sequence length dimension and $r,c$ as the block index.
Denote $Q_{r}\in\mathbb{R}^{B_{r}\times d}$ and $K_{c},V_{c}\in\mathbb{R}^{B_{c}\times d}$ as the block-wise representations.
The forward pass of FlashLLA is summarized in Algorithm [1](https://arxiv.org/html/2510.01450v1#alg1).

Figure: Figure 2: FlashLLA reduces the working set memory to $\Theta(nd)$. The figure shows the profiling result for $d=128$, OOM points are omitted.
Refer to caption: x9.png

Since the statistics equation [7](https://arxiv.org/html/2510.01450v1#S2.E7) for each query are computed independently, the forward pass of LLA can be naturally made parallel for batched queries.
Therefore, the algorithm proceeds by iterating over query blocks $r$.
Within each query block, the iteration over key/value blocks has three passes.
(i) The first pass (line [6](https://arxiv.org/html/2510.01450v1#alg1.l6)-[12](https://arxiv.org/html/2510.01450v1#alg1.l12)) corresponds to accumulating the statistics $M_{r},\omega_{r}$ in an online fashion. Similar to online softmax (Ye, 2023; Milakov and Gimelshein, 2018), we maintain a running maximum $m_{r}$ to ensure numerical stability when computing the kernel weights.
This trick is valid as the computation equation [8](https://arxiv.org/html/2510.01450v1#S2.E8) is homogeneous in $w_{ij}$.
(ii) The second pass (line [15](https://arxiv.org/html/2510.01450v1#alg1.l15)) is encapsulated in the $\texttt{CGSolve}(\cdot)$ operator (see Appendix [C](https://arxiv.org/html/2510.01450v1#A3) for details).
(iii) The third pass (line [17](https://arxiv.org/html/2510.01450v1#alg1.l17)-[22](https://arxiv.org/html/2510.01450v1#alg1.l22)) computes the final output $O_{r}$ using the pre-computed results and the values $V_{c}$. To save computation in backward pass, we also store the intermediate $R_{r}$ and denominator $\delta_{r}$ alongside the output into HBM.

We implement and benchmark the algorithm [1](https://arxiv.org/html/2510.01450v1#alg1) in a custom Triton kernel ($\sim$500 lines of Python) across a range of dimensions and batch sizes on a single NVIDIA H200 GPU.
Figure [2](https://arxiv.org/html/2510.01450v1#S3.F2) demonstrate the quadratic dependency and quickly runs out of memory for naïve method.
In contrast, the blockwise Triton kernel significantly reduces the working set and scales linearly with sequence length, making it hardware-efficient and feasible for long-context and large batch training or inference.

## 4 Empirical Results

#### Test-Time Regression on Non-Stationary Data.

Figure: Figure 3: Test-time regression performance on a piecewise-linear task. The figures demonstrate position-wise MSE for $d=64$ with $S\in\{64,256,512,1024\}$. Results are averaged over $10{,}000$ independently sampled sequences; LLA outperforms other baselines and benefits from more in-segment data; MesaNet excels only before the first shift. The y-axis uses a logarithmic scale.
Refer to caption: x10.png

We first devise a synthetic piecewise-linear regression task to isolate the test-time adaptation capabilities of different attention mechanisms directly without training the query, key and value projections.
Each sample is a length-$L$ sequence partitioned into $L/S$ contiguous segments with $S$ being the segment size.
For each $c\in\{1,\ldots,L/S\}$, keys $k_{i}\in\mathbb{R}^{d}$ are drawn from a segment-specific distribution $P_{c}$ supported on a distinct cone in the input space (see Appendix [E](https://arxiv.org/html/2510.01450v1#A5) for construction details).
The corresponding values $v_{i}$ are generated by a segment-specific linear function $v_{i}=A_{c}k_{i}+\epsilon_{i}$ for $i\in\{(c-1)S+1,\ldots,cS\}$ where $A_{c}\sim\mathcal{N}(0,I)$ and $\epsilon_{i}\sim\mathcal{N}(0,\delta^{2}I)$.
This design ensures the generated data $\{(k_{i},v_{i})\}_{i=1}^{L}$ has non-stationary input distribution and conditional mapping $f_{c}(k)=A_{c}k$.

Figure: Figure 4: The advantage of LLA scales with the data dimension. Axes are in logarithmic scale.
Refer to caption: x11.png

We evaluate a 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 is a 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 is measured by the position-wise MSE $\ell_{i}=\|\hat{f}(k_{i})-v_{i}\|_{2}^{2}$ for $i\in\{1,\ldots,L\}$ to capture the adaptation capability along the sequence.
We also investigate the scaling behavior by evaluating the MSE ratio $\sum_{j=1}^{L}\ell_{j}^{\texttt{Model}}/\sum_{j=1}^{L}\ell_{j}^{\texttt{LLA}}$ for each model compared to LLA across different values of $d$ and $S$.
The results are summarized in Figure [3](https://arxiv.org/html/2510.01450v1#S4.F3).
LLA consistently outperforms other mechanisms as more non-stationarity is observed, even though MesaNet achieves higher performance in the first segment.
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 [4](https://arxiv.org/html/2510.01450v1#S4.F4)), indicating the potential for adaptation to larger models and datasets.

#### In-Context Regression on Non-Stationary Data.

We next evaluate the models’ ability to perform in-context regression on non-stationary, piecewise-linear data.
The data generation process follows the same principle as in the test-time regression task.
The data points $\{x_{i}\in\mathbb{R}^{d_{x}}\}$ are generated from segment-specific distributions $P_{c}$ and the target is given by $y_{i}=A_{c}x_{i}+\epsilon_{i}\in\mathbb{R}^{d_{y}}$ for $i\in\{(c-1)S+1,\ldots,cS\}$, where $A_{c}\sim\mathcal{N}(0,I_{d_{y}\times d_{x}}/d_{x})$ and $\epsilon_{i}\sim\mathcal{N}(0,\delta^{2}I_{d_{y}})$.
Query $x^{\prime}\in\mathbb{R}^{d_{x}}$ is randomly sampled from the segment distributions $P_{c}$. Each in-context regression prompt $X$ is constructed by concatenating $L$ shuffled input-target pairs with $L^{\prime}$ queries:

$$ $\displaystyle X=\begin{pmatrix}x_{1}&x_{2}&\cdots&x_{L}&x_{1}^{\prime}&\cdots&x^{\prime}_{L^{\prime}}\\ y_{1}&y_{2}&\cdots&y_{L}&0&\cdots&0\end{pmatrix}\in\mathbb{R}^{(d_{x}+d_{y})\times(L+L^{\prime})}.$ (17) $$

In contrast to the test-time regression setting, the query, key and value projections are parameterized.
The model $f_{\theta}:\mathbb{R}^{d_{x}+d_{y}}\to\mathbb{R}^{d_{y}}$ is trained to predict the target $Y$, where the label to each query is generated by $y^{\prime}_{i}=A_{c}x_{i}^{\prime}$.
Specifically, for a dataset $\mathcal{D}_{\texttt{train}}=\{(X^{(b)},Y^{(b)})\}_{b=1}^{B}$, we minimize the MSE loss $\mathcal{L}(\theta;\mathcal{D}_{\texttt{train}})$ on the query tokens and report the test error $\mathcal{L}(\theta^{\ast};\mathcal{D}_{\texttt{test}})$ after training.

We compare LLA against several strong baselines, including Softmax Attention, Mamba (Gu and Dao, 2024; Dao and Gu, 2024), GLA (Yang et al., 2024), Hyena (Poli et al., 2023) and Gated DeltaNet (Yang et al., 2025b, a).
We fix the dimension $d_{x}=d_{y}=32$ and query number $L^{\prime}=16$ for all sweeps over segment sizes and evaluate two-layer models without MLPs.
The results in Figure [5](https://arxiv.org/html/2510.01450v1#S4.F5)[5(a)](https://arxiv.org/html/2510.01450v1#S4.F5.sf1) 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 of language models which requires the model to 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 $\mathcal{A}_{k},\mathcal{A}_{v}$ and a set of key-value pairs $(k_{i},v_{i})\in\mathcal{A}_{k}\times\mathcal{A}_{v}$, the set $\{k_{i}\mapsto v_{i}\}$ defines a many-to-one key-value association.
The model is prompted with a sequence of key-value pairs and then queried with keys sampled from the context to predict the corresponding value:

$$ $\underbrace{k_{1},v_{1},k_{2},v_{2},\ldots,k_{N},v_{N}}_{\text{context}}\ \underbrace{k_{i_{1}},o_{i_{1}},\ldots,k_{i_{2}},o_{i_{2}},\ldots,k_{i_{N}},o_{i_{N}}}_{\text{query + answer}},$ $$

where $i_{1},\ldots,i_{N}$ is a permutation of $\{1,\ldots,N\}$ and $o_{i}$ is the model’s prediction.
The position of query follows a power-law to simulate the realistic language data distribution.

As indicated in (Wang et al., 2025), a single short convolution layer is sufficient to solve next-token recall tasks.
Therefore, we disable short convolution in all baselines to ensure fair comparison. We set $\|\mathcal{A}_{k}\cup\mathcal{A}_{v}\|=8k$ and sweep over different sequence lengths and number of KV pairs.
The test recall accuracy results in Figure [5](https://arxiv.org/html/2510.01450v1#S4.F5)[5(b)](https://arxiv.org/html/2510.01450v1#S4.F5.sf2) indicate that the advantages of LLA can be effectively transferred 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 [E.3](https://arxiv.org/html/2510.01450v1#A5.SS3).

Figure: (a) Test Error of In-Context Regression.
Refer to caption: x12.png

#### Permutation State-Tracking.

We then test models’ state-tracking ability by permutation state-tracking task.
Given an initial assignment of items to positions, a 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

$$ $\underbrace{p_{1}{=}\,a_{1},\,p_{2}{=}\,a_{2},\,\ldots,\,p_{N}{=}\,a_{N}}_{\text{initial state}}\ \ \texttt{\#}\ \ \underbrace{i_{1}\,j_{1},\,i_{2}\,j_{2},\,\ldots,\,i_{S}\,j_{S}}_{\text{swap instruction}}\ \ \texttt{\#}\ \ \underbrace{q_{1}{=}\,o_{1},\,\ldots q_{N^{\prime}}{=}\,o_{N^{\prime}}}_{\text{query + answer}}.$ $$

Figure: Figure 6: Test Accuracy of State Tracking. Results show the best accuracy averaged over 3 random seeds.
Refer to caption: x14.png

Here $p_{n}\in\{1,\ldots,N\}$ denotes position $n$; $a_{n}\in\mathcal{A}$ is the item initially assigned to $n$; each $(i_{s},j_{s})\in\{1,\ldots,N\}^{2}$ is a swap instruction exchanging the items at positions $i_{s}$ and $j_{s}$; and $q$ is the queried position.
The target is the item at position $q$ after applying all $S$ swaps.
We include explicit delimiter tokens #, =, and , for structure.

We draw the number of swaps as $S\sim\text{Uniform}(N/6,N/3)$ and set $|\mathcal{A}|=8k$ for each example.
The results in Figure [6](https://arxiv.org/html/2510.01450v1#S4.F6) show that LLA achieves test accuracy on par with Softmax Attention across $N$.
This outcome is expected from a complexity-theoretic perspective: constant-depth Softmax Attention is no more expressive than constant-depth threshold circuits $\mathsf{TC}^{0}$ and has limited ability to realize unbounded-depth state-tracking as $N$ grows (Hahn, 2020; Merrill and Sabharwal, 2023).
By Eqs. [7](https://arxiv.org/html/2510.01450v1#S2.E7) and [8](https://arxiv.org/html/2510.01450v1#S2.E8), 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 a constant extra circuit layer.
Its performance therefore matches the theoretical limits of Softmax Attention, explaining the results in Figure [6](https://arxiv.org/html/2510.01450v1#S4.F6).

## 5 Limitations and Future Directions

#### High Computation and I/O Intensity.

Despite the significant reduction in memory consumption, LLA’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 is an important direction for future work.
Furthermore, while FlashLLA achieves the same $\Theta(nd+n^{2})$ 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-aware optimizations, such as sliding windows or sparsity, could further reduce I/O complexity.

#### Kernel Development and Evaluation on LLMs.

This work evaluates LLA on synthetic and moderate-scale tasks; its efficacy on large language models remains an ongoing question.
Training LLMs with LLA using PyTorch implementation is 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 in equation [9](https://arxiv.org/html/2510.01450v1#S2.E9), LLA provides an optimal interpolation between Linear and Softmax Attention in solving the regression objective.
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 circuit complexity of Softmax Attention.
For instance, future work could explore the integration of state-of-the-art Linear Attention architectures such as DeltaNet and Mamba using this template.

## Acknowledgements

Yifei Zuo performed the research project partly during his internship at Snowfake AI Research. Authors would also like to thank Boyi Liu for meaningful discussions and valuable feedback.

## References

- Ahn et al. (2023)
Ahn, K., Cheng, X., Daneshmand, H. and Sra, S. (2023).
Transformers learn to implement preconditioned gradient descent for in-context learning.
[https://arxiv.org/abs/2306.00297](https://arxiv.org/abs/2306.00297)
- Akyürek et al. (2022)
Akyürek, E., Schuurmans, D., Andreas, J., Ma, T. and Zhou, D. (2022).
What learning algorithm is in-context learning? investigations with linear models.
arXiv preprint arXiv:2211.15661.
- Arora et al. (2023)
Arora, S., Eyuboglu, S., Timalsina, A., Johnson, I., Poli, M., Zou, J., Rudra, A. and Ré, C. (2023).
Zoology: Measuring and improving recall in efficient language models.
[https://arxiv.org/abs/2312.04927](https://arxiv.org/abs/2312.04927)
- Behrouz et al. (2025)
Behrouz, A., Razaviyayn, M., Zhong, P. and Mirrokni, V. (2025).
It’s all connected: A journey through test-time memorization, attentional bias, retention, and online optimization.
[https://arxiv.org/abs/2504.13173](https://arxiv.org/abs/2504.13173)
- Bick et al. (2025)
Bick, A., Xing, E. and Gu, A. (2025).
Understanding the skill gap in recurrent language models: The role of the gather-and-aggregate mechanism.
[https://arxiv.org/abs/2504.18574](https://arxiv.org/abs/2504.18574)
- Bierens (1988)
Bierens, H. J. (1988).
The nadaraya–watson kernel regression function estimator.
Working Paper 1988-58, Faculty of Economics and Business Administration, Vrije Universiteit Amsterdam, Amsterdam.
- Dai et al. (2023)
Dai, D., Sun, Y., Dong, L., Hao, Y., Ma, S., Sui, Z. and Wei, F. (2023).
Why can gpt learn in-context? language models implicitly perform gradient descent as meta-optimizers.
[https://arxiv.org/abs/2212.10559](https://arxiv.org/abs/2212.10559)
- Dao (2023)
Dao, T. (2023).
Flashattention-2: Faster attention with better parallelism and work partitioning.
[https://arxiv.org/abs/2307.08691](https://arxiv.org/abs/2307.08691)
- Dao et al. (2022)
Dao, T., Fu, D. Y., Ermon, S., Rudra, A. and Ré, C. (2022).
Flashattention: Fast and memory-efficient exact attention with io-awareness.
[https://arxiv.org/abs/2205.14135](https://arxiv.org/abs/2205.14135)
- Dao and Gu (2024)
Dao, T. and Gu, A. (2024).
Transformers are ssms: Generalized models and efficient algorithms through structured state space duality.
[https://arxiv.org/abs/2405.21060](https://arxiv.org/abs/2405.21060)
- Dehghani et al. (2023)
Dehghani, M., Djolonga, J., Mustafa, B., Padlewski, P., Heek, J., Gilmer, J., Steiner, A. P., Caron, M., Geirhos, R., Alabdulmohsin, I. et al. (2023).
Scaling vision transformers to 22 billion parameters.
In International conference on machine learning. PMLR.
- Fan et al. (1996)
Fan, J., Gijbels, I., Hu, T.-C. and Huang, L.-S. (1996).
A study of variable bandwidth selection for local polynomial regression.
Statistica Sinica 113–127.
- Gao et al. (2024)
Gao, Y., Zeng, Z., Du, D., Cao, S., Zhou, P., Qi, J., Lai, J., So, H. K.-H., Cao, T., Yang, F. et al. (2024).
Seerattention: Learning intrinsic sparse attention in your llms.
arXiv preprint arXiv:2410.13276.
- Garg et al. (2023)
Garg, S., Tsipras, D., Liang, P. and Valiant, G. (2023).
What can transformers learn in-context? a case study of simple function classes.
[https://arxiv.org/abs/2208.01066](https://arxiv.org/abs/2208.01066)
- Gers et al. (2000)
Gers, F. A., Schmidhuber, J. and Cummins, F. (2000).
Learning to forget: Continual prediction with lstm.
Neural computation, 12 2451–2471.
- Gu and Dao (2024)
Gu, A. and Dao, T. (2024).
Mamba: Linear-time sequence modeling with selective state spaces.
[https://arxiv.org/abs/2312.00752](https://arxiv.org/abs/2312.00752)
- Gu et al. (2022a)
Gu, A., Goel, K. and Ré, C. (2022a).
Efficiently modeling long sequences with structured state spaces.
[https://arxiv.org/abs/2111.00396](https://arxiv.org/abs/2111.00396)
- Gu et al. (2022b)
Gu, A., Johnson, I., Timalsina, A., Rudra, A. and Ré, C. (2022b).
How to train your hippo: State space models with generalized orthogonal basis projections.
[https://arxiv.org/abs/2206.12037](https://arxiv.org/abs/2206.12037)
- Hahn (2020)
Hahn, M. (2020).
Theoretical limitations of self-attention in neural sequence models.
Transactions of the Association for Computational Linguistics, 8 156–171.
- Hestenes and Stiefel (1952)
Hestenes, M. R. and Stiefel, E. (1952).
Methods of conjugate gradients for solving linear systems.
Journal of Research of the National Bureau of Standards, 49 409–436.
- Jelassi et al. (2024)
Jelassi, S., Brandfonbrener, D., Kakade, S. M. and Malach, E. (2024).
Repeat after me: Transformers are better than state space models at copying.
[https://arxiv.org/abs/2402.01032](https://arxiv.org/abs/2402.01032)
- Katharopoulos et al. (2020)
Katharopoulos, A., Vyas, A., Pappas, N. and Fleuret, F. (2020).
Transformers are rnns: Fast autoregressive transformers with linear attention.
[https://arxiv.org/abs/2006.16236](https://arxiv.org/abs/2006.16236)
- Kirsch et al. (2024)
Kirsch, L., Harrison, J., Sohl-Dickstein, J. and Metz, L. (2024).
General-purpose in-context learning by meta-learning transformers.
[https://arxiv.org/abs/2212.04458](https://arxiv.org/abs/2212.04458)
- Lin et al. (2025)
Lin, Z., Nikishin, E., He, X. O. and Courville, A. (2025).
Forgetting transformer: Softmax attention with a forget gate.
[https://arxiv.org/abs/2503.02130](https://arxiv.org/abs/2503.02130)
- Liu et al. (2024)
Liu, B., Wang, R., Wu, L., Feng, Y., Stone, P. and Liu, Q. (2024).
Longhorn: State space models are amortized online learners.
[https://arxiv.org/abs/2407.14207](https://arxiv.org/abs/2407.14207)
- Lu et al. (2025)
Lu, E., Jiang, Z., Liu, J., Du, Y., Jiang, T., Hong, C., Liu, S., He, W., Yuan, E., Wang, Y. et al. (2025).
Moba: Mixture of block attention for long-context llms.
arXiv preprint arXiv:2502.13189.
- Ma et al. (2024)
Ma, X., Yang, X., Xiong, W., Chen, B., Yu, L., Zhang, H., May, J., Zettlemoyer, L., Levy, O. and Zhou, C. (2024).
Megalodon: Efficient llm pretraining and inference with unlimited context length.
[https://arxiv.org/abs/2404.08801](https://arxiv.org/abs/2404.08801)
- Mahankali et al. (2023)
Mahankali, A., Hashimoto, T. B. and Ma, T. (2023).
One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention.
[https://arxiv.org/abs/2307.03576](https://arxiv.org/abs/2307.03576)
- Merrill and Sabharwal (2023)
Merrill, W. and Sabharwal, A. (2023).
The parallelism tradeoff: Limitations of log-precision transformers.
Transactions of the Association for Computational Linguistics, 11 531–545.
- Milakov and Gimelshein (2018)
Milakov, M. and Gimelshein, N. (2018).
Online normalizer calculation for softmax.
[https://arxiv.org/abs/1805.02867](https://arxiv.org/abs/1805.02867)
- Nadaraya (1964)
Nadaraya, E. A. (1964).
On estimating regression.
Theory of Probability and Its Applications, 9 141–142.
- Peng et al. (2024)
Peng, B., Goldstein, D., Anthony, Q., Albalak, A., Alcaide, E., Biderman, S., Cheah, E., Du, X., Ferdinan, T., Hou, H., Kazienko, P., GV, K. K., Kocoń, J., Koptyra, B., Krishna, S., Jr., R. M., Lin, J., Muennighoff, N., Obeid, F., Saito, A., Song, G., Tu, H., Wirawan, C., Woźniak, S., Zhang, R., Zhao, B., Zhao, Q., Zhou, P., Zhu, J. and Zhu, R.-J. (2024).
Eagle and finch: Rwkv with matrix-valued states and dynamic recurrence.
[https://arxiv.org/abs/2404.05892](https://arxiv.org/abs/2404.05892)
- Peng et al. (2025)
Peng, B., Zhang, R., Goldstein, D., Alcaide, E., Du, X., Hou, H., Lin, J., Liu, J., Lu, J., Merrill, W., Song, G., Tan, K., Utpala, S., Wilce, N., Wind, J. S., Wu, T., Wuttke, D. and Zhou-Zheng, C. (2025).
Rwkv-7 ”goose” with expressive dynamic state evolution.
[https://arxiv.org/abs/2503.14456](https://arxiv.org/abs/2503.14456)
- Poli et al. (2023)
Poli, M., Massaroli, S., Nguyen, E., Fu, D. Y., Dao, T., Baccus, S., Bengio, Y., Ermon, S. and Ré, C. (2023).
Hyena hierarchy: Towards larger convolutional language models.
[https://arxiv.org/abs/2302.10866](https://arxiv.org/abs/2302.10866)
- Press et al. (2022)
Press, O., Smith, N. A. and Lewis, M. (2022).
Train short, test long: Attention with linear biases enables input length extrapolation.
[https://arxiv.org/abs/2108.12409](https://arxiv.org/abs/2108.12409)
- Ramsauer et al. (2021)
Ramsauer, H., Schäfl, B., Lehner, J., Seidl, P., Widrich, M., Adler, T., Gruber, L., Holzleitner, M., Pavlović, M., Sandve, G. K., Greiff, V., Kreil, D., Kopp, M., Klambauer, G., Brandstetter, J. and Hochreiter, S. (2021).
Hopfield networks is all you need.
[https://arxiv.org/abs/2008.02217](https://arxiv.org/abs/2008.02217)
- Schlag et al. (2021)
Schlag, I., Irie, K. and Schmidhuber, J. (2021).
Linear transformers are secretly fast weight programmers.
[https://arxiv.org/abs/2102.11174](https://arxiv.org/abs/2102.11174)
- Schmidhuber (1992)
Schmidhuber, J. (1992).
Learning to control fast-weight memories: An alternative to dynamic recurrent networks.
Neural Computation, 4 131–139.
- Siems et al. (2025)
Siems, J., Carstensen, T., Zela, A., Hutter, F., Pontil, M. and Grazzi, R. (2025).
Deltaproduct: Improving state-tracking in linear rnns via householder products.
[https://arxiv.org/abs/2502.10297](https://arxiv.org/abs/2502.10297)
- Sun et al. (2023)
Sun, Y., Dong, L., Huang, S., Ma, S., Xia, Y., Xue, J., Wang, J. and Wei, F. (2023).
Retentive network: A successor to transformer for large language models.
[https://arxiv.org/abs/2307.08621](https://arxiv.org/abs/2307.08621)
- Team (2024)
Team, C. (2024).
Chameleon: Mixed-modal early-fusion foundation models.
arXiv preprint arXiv:2405.09818.
- Vaswani et al. (2023)
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L. and Polosukhin, I. (2023).
Attention is all you need.
[https://arxiv.org/abs/1706.03762](https://arxiv.org/abs/1706.03762)
- von Oswald et al. (2023)
von Oswald, J., Niklasson, E., Randazzo, E., Sacramento, J., Mordvintsev, A., Zhmoginov, A. and Vladymyrov, M. (2023).
Transformers learn in-context by gradient descent.
[https://arxiv.org/abs/2212.07677](https://arxiv.org/abs/2212.07677)
- von Oswald et al. (2025)
von Oswald, J., Scherrer, N., Kobayashi, S., Versari, L., Yang, S., Schlegel, M., Maile, K., Schimpf, Y., Sieberling, O., Meulemans, A., Saurous, R. A., Lajoie, G., Frenkel, C., Pascanu, R., y Arcas, B. A. and Sacramento, J. (2025).
Mesanet: Sequence modeling by locally optimal test-time training.
[https://arxiv.org/abs/2506.05233](https://arxiv.org/abs/2506.05233)
- von Oswald et al. (2024)
von Oswald, J., Schlegel, M., Meulemans, A., Kobayashi, S., Niklasson, E., Zucchet, N., Scherrer, N., Miller, N., Sandler, M., y Arcas, B. A., Vladymyrov, M., Pascanu, R. and Sacramento, J. (2024).
Uncovering mesa-optimization algorithms in transformers.
[https://arxiv.org/abs/2309.05858](https://arxiv.org/abs/2309.05858)
- Wang et al. (2025)
Wang, K. A., Shi, J. and Fox, E. B. (2025).
Test-time regression: a unifying framework for designing sequence models with associative memory.
[https://arxiv.org/abs/2501.12352](https://arxiv.org/abs/2501.12352)
- Watson (1964)
Watson, G. S. (1964).
Smooth regression analysis.
Sankhyā: The Indian Journal of Statistics, Series A, 26 359–372.
- Wortsman et al. (2023)
Wortsman, M., Liu, P. J., Xiao, L., Everett, K., Alemi, A., Adlam, B., Co-Reyes, J. D., Gur, I., Kumar, A., Novak, R. et al. (2023).
Small-scale proxies for large-scale transformer training instabilities.
arXiv preprint arXiv:2309.14322.
- Xiao (2025)
Xiao, G. (2025).
Statistics behind block sparse attention.
[https://guangxuanx.com/blog/block-sparse-attn-stats.html](https://guangxuanx.com/blog/block-sparse-attn-stats.html).
- Yang et al. (2025a)
Yang, S., Kautz, J. and Hatamizadeh, A. (2025a).
Gated delta networks: Improving mamba2 with delta rule.
[https://arxiv.org/abs/2412.06464](https://arxiv.org/abs/2412.06464)
- Yang et al. (2023)
Yang, S., Wang, B., Shen, Y., Panda, R. and Kim, Y. (2023).
Gated linear attention transformers with hardware-efficient training.
arXiv preprint arXiv:2312.06635.
- Yang et al. (2024)
Yang, S., Wang, B., Shen, Y., Panda, R. and Kim, Y. (2024).
Gated linear attention transformers with hardware-efficient training.
[https://arxiv.org/abs/2312.06635](https://arxiv.org/abs/2312.06635)
- Yang et al. (2025b)
Yang, S., Wang, B., Zhang, Y., Shen, Y. and Kim, Y. (2025b).
Parallelizing linear transformers with the delta rule over sequence length.
[https://arxiv.org/abs/2406.06484](https://arxiv.org/abs/2406.06484)
- Ye (2023)
Ye, Z. (2023).
From online softmax to flashattention.
Course notes for ”ML for ML Systems” (CSE 599M), University of Washington, Spring 2023.
Online; University of Washington.
[https://courses.cs.washington.edu/courses/cse599m/23sp/notes/flashattn.pdf](https://courses.cs.washington.edu/courses/cse599m/23sp/notes/flashattn.pdf)
- Yuan et al. (2025)
Yuan, J., Gao, H., Dai, D., Luo, J., Zhao, L., Zhang, Z., Xie, Z., Wei, Y., Wang, L., Xiao, Z. et al. (2025).
Native sparse attention: Hardware-aligned and natively trainable sparse attention.
arXiv preprint arXiv:2502.11089.
- Zhang et al. (2024)
Zhang, R., Frei, S. and Bartlett, P. L. (2024).
Trained transformers learn linear models in-context.
Journal of Machine Learning Research, 25 1–55.
- Zhong et al. (2025)
Zhong, S., Xu, M., Ao, T. and Shi, G. (2025).
Understanding transformer from the perspective of associative memory.
[https://arxiv.org/abs/2505.19488](https://arxiv.org/abs/2505.19488)

## Appendix A Appendix: Proof of Proposition 2.1

Let $(X_{i},Y_{i})_{i=1}^{n}$ be i.i.d. with $X_{i}\in\mathbb{R}^{d}$ supported on a bounded domain $D\subset\mathbb{R}^{d}$ with density $p$, and

$$ $Y_{i}\;=\;f(X_{i})+\varepsilon_{i}\in\mathbb{R}^{d_{y}},\qquad\mathbb{E}[\varepsilon_{i}\mid X_{i}]=0,\ \ \mathbb{E}[\varepsilon_{i}^{2}\mid X_{i}]=\sigma^{2}(X_{i}).$ $$

The global linear estimator at $x\in D$ is

$$ $\widehat{f}_{\mathrm{GL}}(x)=\widehat{\theta}^{\top}(1,x^{\top})^{\top},$ $$

where $\widehat{\theta}\in\mathbb{R}^{d+1}$ minimizes the empirical squared loss over the global affine class.

Let $K:\mathbb{R}^{d}\to[0,\infty)$ be a bounded, compactly supported, radially symmetric kernel with $\int K(u)\,du=1$. For a symmetric p.d. bandwidth matrix $H=H_{n}\succ 0$ with a constant condition-number upper bound $\kappa_{1}$, write

$$ $K_{H}(u)\;:=\;|H|^{-1/2}\,K(H^{-1/2}u),\qquad\|H\|\to 0,\qquad n|H|^{1/2}\to\infty.$ $$

The NW estimator at $x\in D$ is

$$ $\widehat{f}_{\mathrm{NW}}(x)\;=\;\frac{\bm{1}^{\top}W\bm{Y}}{\bm{1}^{\top}W\bm{1}},$ $$

where $W:=\mathrm{Diag}\!\big(K_{H}(X_{i}-x)\big)$, $\bm{1}:=(1,\ldots,1)^{\top}\!\in\mathbb{R}^{n}$ and $\bm{Y}:=(Y_{1},\ldots,Y_{n})\in\mathbb{R}^{n\times d_{y}}$.

To make the analysis easier, we assume the following smoothness requirements.

###### Assumption A.1 .

The domain $D$ has $C^{2}$ boundary (defined as $\partial D$) with principal curvatures uniformly bounded by $\kappa_{2}$.

###### Assumption A.2 .

The density function $p$ of $X$ satisfies $p\in C^{1}(D)$. For all dimensions $j\in\{1,\ldots,d_{y}\}$, the function $f$ satisfies $f_{j}\in C^{2}(D)$, and the variance function $\sigma^{2}$ satisfies $\sigma^{2}_{j}\in C(D)$.

We also assume that the kernel $K$ has easy-to-handle support.

###### Assumption A.3 .

$K$ is radial, compactly supported in the unit ball $\mathbb{B}^{d}:=\{u\in\mathbb{R}^{d}:||u||\leq 1\}=:\mathrm{supp}(K)$.

Throughout, $\mathbb{E}[\cdot]$ denotes expectation with respect to the randomness of the training sample $\mathcal{S}_{n}:=\{(X_{i},Y_{i})\}_{i=1}^{n}$, while $\int_{D}(\cdot)\,dx$ denotes the Lebesgue integral over the spatial domain $D$. Because

$$ $\mathbb{E}\!\int_{D}\|\widehat{f}(x)-f(x)\|^{2}\,dx\;=\;\sum_{j=1}^{d_{y}}\mathbb{E}\!\int_{D}\bigl(\widehat{f}_{j}(x)-f_{j}(x)\bigr)^{2}\,dx,$ $$

the output dimension $d_{y}$ only induces a summation across components and does not affect the order; hence, without loss of generality, we take $d_{y}=1$ below.

###### Assumption A.1 .

###### Assumption A.2 .

###### Assumption A.3 .

### A.1 Integral Error Estimation of Global Linear Regression

We consider the global affine class

$$ $\mathcal{G}\;:=\;\bigl\{\,g_{\theta}(x)=\beta_{0}+\beta^{\top}x\;:\;\theta=(\beta_{0},\beta^{\top})^{\top}\in\mathbb{R}^{1+d}\,\bigr\}.$ $$

###### Lemma A.1 .

If $f\notin\mathcal{G}$, there exists a constant $A_{D}^{*}>0$ such that for every $n$,

$$ $\mathbb{E}\!\left[\;\int_{D}\bigl(\widehat{f}_{\mathrm{GL}}(x)-f(x)\bigr)^{2}\,dx\;\right]\;\geq\;A_{D}^{\star}.$ $$

###### Proof.

Let $L^{2}(D,dx):=\{h:D\to\mathbb{R}\ \text{with}\ \int_{D}h(x)^{2}\,dx<\infty\}$ endowed with inner product $\langle h_{1},h_{2}\rangle=\int_{D}h_{1}(x)h_{2}(x)\,dx$. The set $\mathcal{G}$ is a finite-dimensional linear subspace and hence closed in $L^{2}(D,dx)$. By the Projection Theorem, the $L^{2}(D,dx)$-orthogonal projection of $f$ onto $\mathcal{G}$ exists and is unique:

$$ $g^{\dagger}\in\arg\min_{g\in\mathcal{G}}\;\|f-g\|_{L^{2}(D)}^{2}\;=\;\arg\min_{g\in\mathcal{G}}\;\int_{D}\bigl(f(x)-g(x)\bigr)^{2}\,dx.$ $$

Set $A_{D}^{\star}:=\|f-g^{\dagger}\|_{L^{2}(D)}^{2}$. If $f\notin\mathcal{G}$, then $f-g^{\dagger}\neq 0$ in $L^{2}(D)$ and thus $A_{D}^{\star}>0$.

Fix an arbitrary realization $\mathcal{S}_{n}:=(X_{i},Y_{i})_{i=1}^{n}$. Since $\widehat{f}_{\mathrm{GL}}(\cdot;\mathcal{S}_{n})\in\mathcal{G}$, the optimality of $g^{\dagger}$ yields

$$ $\int_{D}\bigl(\widehat{f}_{\mathrm{GL}}(x;\mathcal{S}_{n})-f(x)\bigr)^{2}\,dx\;\geq\;\inf_{g\in\mathcal{G}}\int_{D}\bigl(g(x)-f(x)\bigr)^{2}\,dx\;=\;A_{D}^{\star}.$ $$

This inequality holds for every sample $\mathcal{S}_{n}$. Taking expectation over the training data proves the claim.
∎

###### Lemma A.1 .

###### Proof.

### 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 expressions are given by

$$ $\displaystyle\mathrm{MSE}_{\mathrm{NW}}(x)$ $\displaystyle=\mathbb{E}\left[\left(\widehat{f}_{\mathrm{NW}}(x)-f(x)\right)^{2}\mid\{X_{i}\}_{i=1}^{n}\right]$ $\displaystyle=\underbrace{\left(\mathbb{E}\left[\widehat{f}_{\mathrm{NW}}(x)-f(x)\mid\{X_{i}\}_{i=1}^{n}\right]\right)^{2}}_{\mathrm{Bias}_{\mathrm{NW}}(x)^{2}}+\underbrace{\mathbb{E}\left[\left(\widehat{f}_{\mathrm{NW}}(x)-\mathbb{E}\left[\widehat{f}_{\mathrm{NW}}(x)\mid\{X_{i}\}_{i=1}^{n}\right]\right)^{2}\mid\{X_{i}\}_{i=1}^{n}\right]}_{\mathrm{Var}_{\mathrm{NW}}(x)}.$ $$

When estimating at $x\in D$, the kernel $K_{H}$ maps the translated domain $(D-x)$ into $\mathrm{supp}(K)$ via the scaling $H^{-1/2}$. The mapped set differs depending on whether $x$ lies in the interior of $D$ or near $\partial D$. This geometric difference drives a larger boundary error for $\mathrm{NW}$, which will underlie the performance gap between $\mathrm{NW}$ and $\mathrm{LL}$. We formalize the mapped kernel domain and the boundary layer.

###### Definition A.1 (Exact kernel domain and boundary layer.) .

For any $x\in D$ and bandwidth $H$, the exact kernel domain is

$$ $D_{x,H}:=H^{-1/2}(D-x)\cap\mathrm{supp}(K)=\{u\in\mathbb{B}^{d}:x+H^{1/2}u\in D\}.$ $$

Define the boundary layer by

$$ $\mathcal{B}(H):=\{\,x\in D:\ D_{x,H}\neq\mathrm{supp}(K)\,\}.$ $$

We use the following kernel-moment shorthands.

###### Definition A.2 (Exact kernel moments) .

For all $x\in D$ and bandwidth $H$, we define

$$ $\mu_{0}^{\star}(x,H):=\int_{D_{x,H}}K(u)\,du,\qquad\mu_{1}^{\star}(x,H):=\int_{D_{x,H}}u\,K(u)\,du,\qquad\mu_{2}^{\star}(x,H):=\int_{D_{x,H}}uu^{\top}K(u)\,du.$ $$

Define normalized moments $\bar{\mu}_{r}^{\star}(x,H):=\mu^{\star}_{r}(x,H)/\mu^{\star}_{0}(x,H)$.

Then we are ready to estimate the bias and variance using $H$ and $n$.

###### Lemma A.2 .

Under Assumption [A.2](https://arxiv.org/html/2510.01450v1#A1.Thmassumption2), we have

$$ $\mathrm{Bias}_{\mathrm{NW}}(x)=\nabla f(x)^{\top}H^{1/2}\bar{\mu}_{1}^{\star}(x,H)\;+\;O_{p}(\|H\|),\quad\mathrm{Var}_{\mathrm{NW}}(x)=\Theta_{p}(\frac{1}{n|H|^{1/2}})$ $$

uniformly for all $H\in\mathcal{H}_{n}$, where $\mathcal{H}_{n}:=\{H=h^{2}B:h\in[n^{-a},n^{-b}],B\succ 0,|B|=1,\kappa(B)\leq\kappa_{1}\}$ and $0<b<a<1$.

###### Proof.

Defining $\bm{f}:=(f(X_{1}),\ldots f(X_{n}))^{\top}$, for any fixed point $x_{0}\in D$, we have

$$ $\displaystyle\mathrm{Bias}_{\mathrm{NW}}(x_{0})$ $\displaystyle=(\bm{1}^{\top}W\bm{1})^{-1}\bm{1}^{\top}W(\bm{f}-f(x_{0})\bm{1}).$ $$

Specifically,

$$ $\displaystyle n^{-1}(\bm{1}^{\top}W\bm{1})=n^{-1}\sum_{i=1}^{n}K_{H}(X_{i}-x_{0}),\quad n^{-1}\bm{1}^{\top}W(\bm{f}-f(x_{0})\bm{1})=n^{-1}\sum_{i=1}^{n}K_{H}(X_{i}-x_{0})(f(X_{i})-f(x_{0})).$ $$

By Chebyshev’s inequality, for any fixed $H$,

$$ $\displaystyle n^{-1}\sum_{i=1}^{n}K_{H}(X_{i}-x_{0})=\int_{D}K_{H}(x-x_{0})p(x)dx+O_{p}\left(n^{-1}\sqrt{n\int_{D}K_{H}^{2}(x-x_{0})p(x)dx}\right).$ $$

By Theorem 1 in (Fan et al., 1996), this bound holds *uniformly* over $H\in\mathcal{H}_{n}$ after multiplying the stochastic term by $\sqrt{\log n}$. Hence,

$$ $\displaystyle n^{-1}\sum_{i=1}^{n}K_{H}(X_{i}-x_{0})=\int_{D}K_{H}(x-x_{0})p(x)dx+O_{p}\left(n^{-1}\sqrt{n\log n\int_{D}K_{H}^{2}(x-x_{0})p(x)dx}\right)$ (18) $\displaystyle=\int_{D_{x_{0},H}}K(u)p(x_{0}+H^{1/2}u)du+o_{p}(1)=p(x_{0})\mu_{0}^{*}(x,H)+o_{p}(1)$ $$

uniformly for $H\in\mathcal{H}_{n}$.

Similarly,

$$ $\displaystyle n^{-1}\sum_{i=1}^{n}K_{H}(X_{i}-x_{0})(f(X_{i})-f(x_{0}))$ (19) $\displaystyle=\int_{D}K_{H}(x-x_{0})p(x)(f(x)-f(x_{0}))dx+O_{p}\left(n^{-1}\sqrt{n\log n\int_{D}K_{H}^{2}(x-x_{0})(f(x)-f(x_{0}))^{2}p(x)dx}\right)$ $\displaystyle=\int_{D_{x_{0},H}}K(u)p(x_{0}+H^{1/2}u)(f(x_{0}+H^{1/2}u)-f(x_{0}))du+o_{p}(||H||)$ $\displaystyle=\int_{D_{x_{0},H}}K(u)(p(x_{0})+\nabla p(x_{0})H^{1/2}u+O(||H||))(\nabla f(x_{0})^{\top}H^{1/2}u+O(||H||))du+o_{p}(||H||)$ $\displaystyle=p(x_{0})\nabla f(x_{0})^{\top}H^{1/2}\mu_{1}^{*}(x_{0},H)+\nabla p(x_{0})H^{1/2}\mu_{2}^{*}(x_{0},H)H^{1/2}\nabla f(x_{0})+o_{p}(||H||)$ $\displaystyle=p(x_{0})\nabla f(x_{0})^{\top}H^{1/2}\mu_{1}^{*}(x_{0},H)+O_{p}(||H||)$ $$

uniformly for all $H\in\mathcal{H}_{n}$.

Combining Equations [18](https://arxiv.org/html/2510.01450v1#A1.E18), [19](https://arxiv.org/html/2510.01450v1#A1.E19), we have

$$ $\mathrm{Bias}_{\mathrm{NW}}(x_{0})=\nabla f(x_{0})^{\top}H^{1/2}\bar{\mu}_{1}^{\star}(x_{0},H)\;+\;O_{p}(\|H\|)$ $$

uniformly for all $H\in\mathcal{H}_{n}$.

Then we calculate variance

$$ $\mathrm{Var}_{\mathrm{NW}}(x_{0})=(\bm{1}^{\top}W\bm{1})^{-1}(\bm{1}^{\top}\Sigma\bm{1})(\bm{1}^{\top}W\bm{1})^{-1}$ $$

where $\Sigma:=\text{diag}(K_{H}^{2}(X_{i}-x_{0})\sigma^{2}(X_{i}))$. Analogously to equation [18](https://arxiv.org/html/2510.01450v1#A1.E18),

$$ $\displaystyle n^{-1}(\bm{X}^{\top}\Sigma\bm{X})=n^{-1}\sum_{i=1}^{n}K_{H}^{2}(X_{i}-x_{0})\sigma^{2}(X_{i})$ (20) $\displaystyle=\int_{D}K_{H}^{2}(x-x_{0})\sigma^{2}(x)p(x)dx+O_{p}\left(n^{-1}\sqrt{n\log n\int_{D}K_{H}^{4}(x-x_{0})\sigma^{4}(x)p(x)dx}\right)$ $\displaystyle=|H|^{-1/2}\int_{D_{x_{0},H}}K^{2}(u)\sigma^{2}(x_{0}+H^{1/2}u)p(x_{0}+H^{1/2}u)du+O_{p}\left(\sqrt{\frac{\log n}{n}}|H|^{-3/4}\right)$ $\displaystyle=|H|^{-1/2}\int_{D_{x_{0},H}}K^{2}(u)\sigma^{2}(x_{0})p(x_{0})du(1+o_{p}(1))+O_{p}\left(\sqrt{\frac{\log n}{n}}|H|^{-3/4}\right)$ $\displaystyle=|H|^{-1/2}\sigma^{2}(x_{0})p(x_{0})\int_{D_{x_{0}},H}K^{2}(u)du(1+o_{p}(1))$ $$

uniformly for all $H\in\mathcal{H}_{n}$.

Combining Equations [18](https://arxiv.org/html/2510.01450v1#A1.E18), [20](https://arxiv.org/html/2510.01450v1#A1.E20), we have that

$$ $\mathrm{Var}_{\mathrm{NW}}(x)=\Theta_{p}(\frac{1}{n|H|^{1/2}})$ $$

holds uniformly for all $H\in\mathcal{H}_{n}$.
∎

###### Definition A.1 (Exact kernel domain and boundary layer.) .

###### Definition A.2 (Exact kernel moments) .

###### Lemma A.2 .

###### Proof.

### A.3 Integral Error Estimation of Local Constant Regression

Here we calculate the integral error for NW estimator.
From Lemma [A.2](https://arxiv.org/html/2510.01450v1#A1.Thmlemma2) we have that for any $x\in D$, the pointwise bias of NW estimator is

$$ $\mathrm{Bias}_{\mathrm{NW}}(x)=\nabla f(x)^{\top}H^{1/2}\bar{\mu}_{1}^{\star}(x,H)\;+\;O_{p}(\|H\|).$ $$

By symmetry of the kernel $K$, the first moment $\bar{\mu}_{1}^{\star}(x,H)=0$ if and only if $x\notin\mathcal{B}(H)$. Roughly, $\mathrm{Bias}_{\mathrm{NW}}(x)=O(\|H^{1/2}\|)$ when $x\in\mathcal{B}(H)$ and $\mathrm{Bias}_{\mathrm{NW}}(x)=O(\|H\|)$ when $x\notin\mathcal{B}(H)$. Hence it is important to bound $\int_{\mathcal{B}(H)}\mathrm{MSE}_{\mathrm{NW}}(x)\,dx$ and $\int_{D\setminus\mathcal{B}(H)}\mathrm{MSE}_{\mathrm{NW}}(x)\,dx$ separately. We first define a smooth substitute for the boundary layer $\mathcal{B}(H)$.

###### Definition A.3 (Uniform Boundary Layer) .

For any bandwidth $H$ and $\alpha\in(0,1)$, the uniform boundary layer of $D$ is

$$ $\mathcal{C}(H,\alpha):=\{x=y-te(y):y\in\partial D,t\in[0,\alpha h_{n}(y)]\}$ $$

where for any $y\in\partial D$, we define $e(y)$
as the inward unit normal at $y$, and
define
$h_{n}(y):=\sqrt{\,e(y)^{\top}H\,e(y)\,}$.

We now show that one can choose a constant $\alpha\in(0,1)$ (independent of $H$) that “sandwiches” $\mathcal{B}(H)$ between $\mathcal{C}(H,\alpha)$ and $\mathcal{C}(H,\alpha^{-1})$.

###### Lemma A.3 .

Under Assumptions [A.1](https://arxiv.org/html/2510.01450v1#A1.Thmassumption1), [A.3](https://arxiv.org/html/2510.01450v1#A1.Thmassumption3), there exists $\alpha\in(0,1)$ such that for all sufficiently small $H$, we have

$$ $\mathcal{C}(H,\alpha)\subset\ \mathcal{B}(H)\subset\mathcal{C}(H,\alpha^{-1}).$ $$

###### Proof.

We prove the claim with $\alpha=\min\!\bigl(\tfrac{1}{2},\tfrac{1}{2}\kappa_{1}^{-1/2}\bigr)$.

*Step 1: $\mathcal{C}(H,\alpha)\subset\mathcal{B}(H)$.*
Let $\phi$ denote the signed distance to $\partial D$ (positive inside $D$). For $x=y-te(y)\in\mathcal{C}(H,\alpha)$ and any $v\in\mathbb{R}^{d}$, the standard expansion (using the shape operator $S_{y}$) gives

$$ $\phi(x+v)=-t+v_{n}-\tfrac{1}{2}v_{T}^{\top}S_{y}v_{T}+O(\|v\|^{3}),$ $$

where $v_{n}:=v^{\top}e(y)$ and $v_{T}:=(I-e(y)e(y)^{\top})v$. By Assumption [A.1](https://arxiv.org/html/2510.01450v1#A1.Thmassumption1), $\|S_{y}\|$ is uniformly bounded.

Define

$$ $u:=\kappa_{1}^{-1/2}\,h_{n}(y)\,H^{-1/2}e(y)\in\mathbb{R}^{d}.$ $$

Then

$$ $\|u\|^{2}=\kappa_{1}^{-1}\,h_{n}(y)^{2}\,e(y)^{\top}H^{-1}e(y)=\kappa_{1}^{-1}\,\frac{\bigl(e(y)^{\top}He(y)\bigr)\bigl(e(y)^{\top}H^{-1}e(y)\bigr)}{1}\;\leq\;\kappa_{1}^{-1}\,\kappa(H)\;\leq\;1,$ $$

since $\kappa(H)\leq\kappa_{1}$ by assumption. Hence $u\in\mathrm{supp}(K)$ (Assumption [A.3](https://arxiv.org/html/2510.01450v1#A1.Thmassumption3)). Using the expansion with $v=H^{1/2}u$ and bounded curvature,

$$ $\phi\bigl(x+H^{1/2}u\bigr)=-t+u^{\top}H^{1/2}e(y)+O(\|H\|^{3/2})=-t+\kappa_{1}^{-1/2}h_{n}(y)+O(\|H\|^{3/2}).$ $$

Since $t\leq\alpha h_{n}(y)$ and $\alpha\leq\tfrac{1}{2}\kappa_{1}^{-1/2}$, for sufficiently small $H$ we have $\phi(x+H^{1/2}u)>0$, i.e., $x+H^{1/2}u\notin D$. Thus there exists $u\in\mathrm{supp}(K)$ with $u\notin D_{x,H}$, i.e., $D_{x,H}\neq\mathrm{supp}(K)$, so $x\in\mathcal{B}(H)$.

*Step 2: $\mathcal{B}(H)\subset\mathcal{C}(H,\alpha^{-1})$.*
Let $x=y-te(y)\in\mathcal{B}(H)$. By definition, there exists $u\in\mathrm{supp}(K)$ with $x+H^{1/2}u\notin D$, i.e.,

$$ $\phi(x+H^{1/2}u)=-t+u^{\top}H^{1/2}e(y)+O(\|H\|)>0.$ $$

Hence

$$ $t<u^{\top}H^{1/2}e(y)+O(\|H\|)\leq\|u\|\,\sqrt{e(y)^{\top}He(y)}+O(\|H\|)\leq h_{n}(y)+O(\|H\|).$ $$

For sufficiently small $H$, this yields $t<\alpha^{-1}h_{n}(y)$ (since $\alpha^{-1}\geq 2$), so $x\in\mathcal{C}(H,\alpha^{-1})$.
∎

With $\mathcal{C}(H,\alpha^{-1})$ in hand, we bound the integrated variance and squared bias of the $\mathrm{NW}$ estimator.

###### Lemma A.4 .

Under Assumptions [A.1](https://arxiv.org/html/2510.01450v1#A1.Thmassumption1), [A.2](https://arxiv.org/html/2510.01450v1#A1.Thmassumption2), and [A.3](https://arxiv.org/html/2510.01450v1#A1.Thmassumption3), we have, uniformly for all $H\in\mathcal{H}_{n}$,

$$ $\int_{D}\mathrm{Bias}_{\mathrm{NW}}^{2}(x)\,dx=O_{p}\bigl(\|H\|^{3/2}\bigr)\quad\text{and}\quad\int_{D}\mathrm{Var}_{\mathrm{NW}}(x)\,dx=\Theta\!\bigl(n^{-1}|H|^{-1/2}\bigr).$ $$

###### Proof.

$$ $\displaystyle\int_{D}\mathrm{Bias}_{\mathrm{NW}}^{2}(x)dx=\int_{\mathcal{B}(H)}\mathrm{Bias}_{\mathrm{NW}}^{2}(x)dx+\int_{D\backslash\mathcal{B}(H)}\mathrm{Bias}_{\mathrm{NW}}^{2}(x)dx$ $\displaystyle=O_{p}(||H||)\int_{\mathcal{B}(H)}dx+O_{p}(||H^{2}||)$ $\displaystyle\leq O_{p}(||H||)\int_{\mathcal{C}(H,\alpha^{-1})}dx+O_{p}(||H^{2}||)$ $\displaystyle=O_{p}(||H||)\int_{\partial D}\int_{0}^{\frac{h_{n}(y)}{\alpha}}dtdS(y)+O_{p}(||H^{2}||)$ $\displaystyle=O_{p}(||H^{3/2}||),$ $$

since $h_{n}(y)=\sqrt{e(y)^{\top}He(y)}=\Theta(\|H\|^{1/2})$ uniformly under bounded condition number.

For the variance, Lemma [A.2](https://arxiv.org/html/2510.01450v1#A1.Thmlemma2) gives $\mathrm{Var}_{\mathrm{NW}}(x)=\Theta_{p}(n^{-1}|H|^{-1/2})$ pointwise (uniformly over $H\in\mathcal{H}_{n}$). Integrating over $D$ (whose volume is constant) preserves the order:

$$ $\int_{D}\mathrm{Var}_{\mathrm{NW}}(x)dx=\int_{D}\Theta(n^{-1}|H^{1/2}|^{-1})dx=\Theta(n^{-1}|H^{1/2}|^{-1}).$ $$

∎

###### Definition A.3 (Uniform Boundary Layer) .

###### Lemma A.3 .

###### Proof.

###### Lemma A.4 .

###### Proof.

### A.4 Proof of Proposition 2.1

Combining Lemmas [A.1](https://arxiv.org/html/2510.01450v1#A1.Thmlemma1), [A.4](https://arxiv.org/html/2510.01450v1#A1.Thmlemma4), we obtain the final conclusion.

###### Theorem A.1 (Precise statement of Proposition 2.1 ) .

Under Assumptions [A.1](https://arxiv.org/html/2510.01450v1#A1.Thmassumption1), [A.2](https://arxiv.org/html/2510.01450v1#A1.Thmassumption2), [A.3](https://arxiv.org/html/2510.01450v1#A1.Thmassumption3), if the function $f$ to be estimated is not within the the global affine class $\mathcal{G}$,then

$$ $\mathbb{E}\!\int_{D}\mathrm{MSE}_{\mathrm{GL}}(x)\,dx=\Omega(1)\quad\text{and}\quad\mathbb{E}\!\int_{D}\mathrm{MSE}_{\mathrm{NW}}(x)\,dx=O\!\bigl(n^{-3/(d+3)}\bigr),$ $$

where $\widehat{f}_{\mathrm{NW}}$ uses the optimal bandwidth $H\in\mathcal{H}_{n}$.

###### Proof.

The lower bound for $\mathrm{MSE}_{\mathrm{GL}}$ follows directly from Lemma [A.1](https://arxiv.org/html/2510.01450v1#A1.Thmlemma1).

For $\mathrm{NW}$, Lemma [A.4](https://arxiv.org/html/2510.01450v1#A1.Thmlemma4) and the definition of $\mathcal{H}_{n}$ imply

$$ $\int_{D}\mathrm{Bias}_{\mathrm{NW}}^{2}(x)\,dx=O_{p}(h^{3}),\qquad\int_{D}\mathrm{Var}_{\mathrm{NW}}(x)\,dx=\Theta\!\bigl(n^{-1}h^{-d}\bigr)$ $$

uniformly for $h\in[n^{-a},n^{-b}]$ (since $\|H\|=\Theta(h^{2})$ and $|H|^{1/2}=\Theta(h^{d})$ under bounded condition number). Hence

$$ $\int_{D}\mathrm{MSE}_{\mathrm{NW}}(x)\,dx=O_{p}\!\Bigl(h^{3}+\frac{1}{n\,h^{d}}\Bigr)=O_{p}\!\bigl(n^{-3/(d+3)}\bigr),$ $$

at the minimizer $h=n^{-1/(d+3)}\in[n^{-a},n^{-b}]$ of $h^{3}+(nh^{d})^{-1}$.
∎

###### Theorem A.1 (Precise statement of Proposition 2.1 ) .

###### Proof.

## Appendix B Appendix: Proof of Proposition 2.2

The local linear estimator at $x\in D$ is

$$ $\widehat{f}_{\mathrm{LL}}(x)=(\bm{X}^{\top}W\bm{X})^{-1}\bm{X}^{\top}W\bm{Y},$ $$

where

$$ $\bm{X}:=\begin{bmatrix}1&\ldots&1\\ X_{1}-x_{0}&\ldots&X_{n}-x_{0}\end{bmatrix}^{\top}.$ $$

Here we inherit Assumptions [A.1](https://arxiv.org/html/2510.01450v1#A1.Thmassumption1), [A.2](https://arxiv.org/html/2510.01450v1#A1.Thmassumption2), [A.3](https://arxiv.org/html/2510.01450v1#A1.Thmassumption3) and continue to set $d_{y}=1$ as in Appendix [A](https://arxiv.org/html/2510.01450v1#A1).

### B.1 Point-wise Error Estimation

We have already derived the point-wise error of the local constant estimator in Lemma [A.4](https://arxiv.org/html/2510.01450v1#A1.Thmlemma4). We now do the same for the local linear estimator.

###### Lemma B.1 .

Under Assumption [A.2](https://arxiv.org/html/2510.01450v1#A1.Thmassumption2), we have, uniformly for all $H\in\mathcal{H}_{n}$,

$$ $\mathrm{Bias}_{\mathrm{LL}}(x)=O_{p}(\|H\|),\qquad\mathrm{Var}_{\mathrm{LL}}(x)=O_{p}\!\Bigl(\frac{1}{n|H|^{1/2}}\Bigr),$ $$

where $\mathcal{H}_{n}:=\{H=h^{2}B:\ h\in[n^{-a},n^{-b}],\ B\succ 0,\ |B|=1,\ \kappa(B)\leq\kappa_{1}\}$ with constants $0<b<a<1$.

###### Proof.

For any fixed point $x_{0}\in D$, we have

$$ $\displaystyle\mathrm{Bias}_{\mathrm{LL}}(x_{0})$ $\displaystyle=e_{1}(\bm{X}^{\top}W\bm{X})^{-1}\bm{X}^{\top}W(\bm{f}-\bm{X}(f(x_{0}),\nabla f(x_{0})^{\top})^{\top})$ $\displaystyle=\frac{1}{2}e_{1}(\bm{X}^{\top}W\bm{X})^{-1}\bm{X}^{\top}W(Q+o_{p}(\text{tr}(H)))$ $$

where $Q=[(X_{1}-x)^{\top}\mathcal{H}_{f}(X_{1}-x),\ldots,(X_{n}-x)^{\top}\mathcal{H}_{f}(X_{n}-x)]^{\top}$.

$$ $\mathrm{Var}_{\mathrm{LL}}(x_{0})=e_{1}(\bm{X}^{\top}W\bm{X})^{-1}(\bm{X}^{\top}\Sigma\bm{X})(\bm{X}^{\top}W\bm{X})^{-1}e_{1}^{\top}$ $$

where $\Sigma:=\text{diag}(K_{H}^{2}(X_{i}-x_{0})\sigma^{2}(X_{i}))$.

$$ $n^{-1}(\bm{X}^{\top}W\bm{X})=\begin{bmatrix}n^{-1}\sum_{i=1}^{n}K_{H}(X_{i}-x_{0})&n^{-1}\sum_{i=1}^{n}K_{H}(X_{i}-x_{0})\,(X_{i}-x_{0})^{\top}\\ n^{-1}\sum_{i=1}^{n}K_{H}(X_{i}-x_{0})\,(X_{i}-x_{0})&n^{-1}\sum_{i=1}^{n}K_{H}(X_{i}-x_{0})\,(X_{i}-x_{0})(X_{i}-x_{0})^{\top}\end{bmatrix}$ $$

$$ $n^{-1}(\bm{X}^{\top}\Sigma\bm{X})=\begin{bmatrix}n^{-1}\sum_{i=1}^{n}K_{H}^{2}(X_{i}-x_{0})\sigma^{2}(X_{i})&n^{-1}\sum_{i=1}^{n}K_{H}^{2}(X_{i}-x_{0})\sigma^{2}(X_{i})(X_{i}-x_{0})^{\top}\\ n^{-1}\sum_{i=1}^{n}K_{H}^{2}(X_{i}-x_{0})\sigma^{2}(X_{i})(X_{i}-x_{0})&n^{-1}\sum_{i=1}^{n}K_{H}^{2}(X_{i}-x_{0})\sigma^{2}(X_{i})(X_{i}-x_{0})(X_{i}-x_{0})^{\top}\end{bmatrix}$ $$

$$ $\displaystyle n^{-1}\bm{X}^{\top}W(\bm{f}-\bm{X}(f(x_{0}),\nabla f(x_{0})^{\top})^{\top})$ $\displaystyle=\begin{bmatrix}n^{-1}\sum_{i=1}^{n}K_{H}(X_{i}-x_{0})(f(X_{i})-f(x_{0})-\nabla f(x_{0})^{\top}(X_{i}-x_{0}))\\ n^{-1}\sum_{i=1}^{n}K_{H}(X_{i}-x_{0})(f(X_{i})-f(x_{0})-\nabla f(x_{0})^{\top}(X_{i}-x_{0}))(X_{i}-x_{0})\end{bmatrix}$ $$

By the same uniform LLN and $\sqrt{\log n}$ arguments used in the proof of Lemma [A.2](https://arxiv.org/html/2510.01450v1#A1.Thmlemma2), we have uniformly over $H\in\mathcal{H}_{n}$:

$$ $\displaystyle n^{-1}\sum_{i=1}^{n}K_{H}^{2}(X_{i}-x_{0})\sigma^{2}(X_{i})=|H|^{-1/2}\sigma^{2}(x_{0})p(x_{0})R_{0}(x_{0},H)(1+o_{p}(1))$ $\displaystyle n^{-1}\sum_{i=1}^{n}K_{H}^{2}(X_{i}-x_{0})\sigma^{2}(X_{i})(X_{i}-x_{0})^{\top}$ $\displaystyle\quad=\sigma^{2}(x_{0})p(x_{0})|H|^{-1/2}R_{1}(K)H^{1/2}+O_{p}\left((\sqrt{\frac{\log n}{n}}|H|^{-3/4}+1)H^{1/2}\bm{1}\right)$ $\displaystyle n^{-1}\sum_{i=1}^{n}K_{H}(X_{i}-x_{0})=p(x_{0})\mu_{0}^{*}(x_{0},H)+o_{p}(1)$ $\displaystyle n^{-1}\sum_{i=1}^{n}K_{H}^{2}(X_{i}-x_{0})\sigma^{2}(X_{i})(X_{i}-x_{0})(X_{i}-x_{0})^{\top}$ $\displaystyle\quad=\sigma^{2}(x_{0})p(x_{0})|H|^{-1/2}H^{1/2}R_{2}(K)H^{1/2}+o_{p}\left(|H|^{-1/2}H\right)$ $$

$$ $\displaystyle n^{-1}\sum_{i=1}^{n}K_{H}(X_{i}-x_{0})(X_{i}-x_{0})^{\top}=p(x_{0})\mu_{1}^{*}(x_{0},H)^{\top}H^{1/2}+O_{p}(H\bm{1})$ $\displaystyle n^{-1}\sum_{i=1}^{n}K_{H}(X_{i}-x_{0})(X_{i}-x_{0})(X_{i}-x_{0})^{\top}=O_{p}(H)$ $\displaystyle n^{-1}\sum_{i=1}^{n}K_{H}(X_{i}-x_{0})(f(X_{i})-f(x_{0})-\nabla f(x_{0})^{\top}(X_{i}-x_{0}))=p(x_{0})\mu_{2}(K)\text{tr}\left(H\mathcal{H}_{f}(x_{0})\right)+o_{p}(\text{tr}(H))$ $\displaystyle n^{-1}\sum_{i=1}^{n}K_{H}(X_{i}-x_{0})(f(X_{i})-f(x_{0})-\nabla f(x_{0})^{\top}(X_{i}-x_{0}))(X_{i}-x_{0})=O_{p}(H^{3/2}\bm{1}),$ $$

where

$$ $R_{0}(x_{0},H)=\int_{D_{x_{0},H}}K^{2}(u)du,\quad R_{1}(x_{0},H)=\int_{D_{x_{0},H}}K^{2}(u)udu,\quad R_{2}(x_{0},H)=\int_{D_{x_{0},H}}K^{2}(u)uu^{\top}du.$ $$

Then we have the expression of every element in all matrices. A standard blockwise inversion therefore yields

$$ $\mathrm{Bias}_{\mathrm{LL}}(x)=O_{p}(\|H\|),\quad\mathrm{Var}_{\mathrm{LL}}(x)=O_{p}(\frac{1}{n|H|^{1/2}}).$ $$

uniformly over $H\in\mathcal{H}_{n}$, proving the claim.

∎

###### Lemma B.1 .

###### Proof.

### B.2 Integral Error Estimation

From Lemma [A.2](https://arxiv.org/html/2510.01450v1#A1.Thmlemma2) we have that for any $x\in D$, the pointwise bias of local constant estimator is

$$ $\mathrm{Bias}_{\mathrm{NW}}(x)=\nabla f(x)^{\top}H^{1/2}\bar{\mu}_{1}^{\star}(x,H)\;+\;O_{p}(\|H\|).$ $$

By symmetry of the kernel $K$, the first moment satisfies $\bar{\mu}_{1}^{\star}(x,H)=0$ if and only if $x\notin\mathcal{B}(H)$. We will show that if $f$ has a sufficiently large normal gradient on a measurable subset of $\partial D$, then $\int_{\mathcal{B}(H)}\mathrm{Bias}^{2}_{\mathrm{NW}}(x)dx$ will have a dominant order $\Theta(\|H^{3/2}\|)$.

We begin by lower-bounding $|\nabla f(x)^{\top}H^{1/2}\bar{\mu}_{1}^{\star}(x,H)|$, to do which we first lower-bound $e(y)^{\top}H^{1/2}\bar{\mu}_{1}^{*}(y-te(y),H)$.

###### Lemma B.2 .

Under Assumption [A.1](https://arxiv.org/html/2510.01450v1#A1.Thmassumption1), [A.3](https://arxiv.org/html/2510.01450v1#A1.Thmassumption3), there exists $\alpha\in(0,1)$, $c_{*}>0$ and $C_{*}>0$ such that for all sufficiently small $H$ and all $y\in\partial D$, all $t\in[0,\alpha h_{n}(y)]$,

$$ $|e(y)^{\top}H^{1/2}\bar{\mu}_{1}^{*}(y-te(y),H)|\geq c_{*}h_{n}(y)\qquad||\bar{\mu}_{1}^{*}(y-te(y),H)||\leq C_{*}.$ $$

###### Proof.

Recall $D_{x,H}=\{u\in\mathbb{B}^{d}:\ x+H^{1/2}u\in D\}$. For $x=y-te(y)$ with $y\in\partial D$, choose an orthogonal $Q(y)$ such that
$Q(y)H^{1/2}e(y)=h_{n}(y)\,e_{d}$. Define the rotated domain

$$ $\widetilde{D}_{y,t,H}:=\{\,v\in\mathbb{B}^{d}:\ y-te(y)+H^{1/2}Q(y)v\in D\,\}.$ $$

Define the corresponding moments on $\widetilde{D}_{y,t,H}$:

$$ $\mu_{0,v}^{\star}(y,t,H):=\int_{\widetilde{D}_{y,t,H}}K(u)\,du,\quad\mu_{1,v}^{\star}(y,t,H):=\int_{\widetilde{D}_{y,t,H}}u\,K(u)\,du,\quad\bar{\mu}_{1,v}^{\star}(y,t,H):=\frac{\mu_{1,v}^{\star}(y,t,H)}{\mu_{0,v}^{\star}(y,t,H)}.$ $$

Then

$$ $\mu_{0,v}^{\star}(y,t,H)=\mu_{0}^{\star}(x,H),\quad Q(y)\,\mu_{1,v}^{\star}(y,t,H)=\mu_{1}^{\star}(x,H),\quad Q(y)\,\bar{\mu}_{1,v}^{\star}(y,t,H)=\bar{\mu}_{1}^{\star}(x,H),$ $$

and hence

$$ $e(y)^{\top}H^{1/2}\bar{\mu}_{1}^{\star}(x,H)=h_{n}(y)\,e_{d}^{\top}\bar{\mu}_{1,v}^{\star}(y,t,H).$ (21) $$

By the standard signed-distance expansion with shape operator $S_{y}$ (Assumption [A.1](https://arxiv.org/html/2510.01450v1#A1.Thmassumption1)), for $v\in\mathbb{R}^{d}$,

$$ $\phi\bigl(y-te(y)+H^{1/2}Q(y)v\bigr)=-t+h_{n}(y)\,v_{d}-\tfrac{1}{2}z_{T}^{\top}S_{y}z_{T}+O(\|H\|^{3/2}),$ $$

where $z_{T}:=(I-e(y)e(y)^{\top})H^{1/2}Q(y)v$ and $\|S_{y}\|$ is uniformly bounded. Therefore

$$ $y-te(y)+H^{1/2}Q(y)v\in D\ \iff\ v_{d}<\frac{t}{h_{n}(y)}+O(\|H\|^{1/2}),$ $$

so

$$ $\widetilde{D}_{y,t,H}=\Bigl\{\,v\in\mathbb{B}^{d}:\ v_{d}<\frac{t}{h_{n}(y)}+O(\|H\|^{1/2})\,\Bigr\}.$ $$

Consequently,

$$ $\displaystyle\mu_{0,v}^{\star}(y,t,H)$ $\displaystyle=\int_{\widetilde{D}_{y,t,H}}K(v)\,dv$ $\displaystyle=\int_{\mathbb{R}^{d-1}}\!\!dv_{1:d-1}\int_{-\sqrt{1-\sum_{i=1}^{d-1}v_{i}^{2}}}^{t/h_{n}(y)}K(v_{1:d-1},v_{d})\,dv_{d}\;+\;O(\|H\|^{1/2})$ $\displaystyle\geq\int_{\mathbb{R}^{d-1}}\!\!dv_{1:d-1}\int_{-\sqrt{1-\sum_{i=1}^{d-1}v_{i}^{2}}}^{0}K(v_{1:d-1},v_{d})\,dv_{d}\;+\;O(\|H\|^{1/2}).$ $$

Using positivity and boundedness of $K$ on $\mathbb{B}^{d}$, there exist $0<C<D<\infty$ such that

$$ $C<\mu_{0,v}^{\star}(y,t,H)<D.$ (22) $$

Similarly,

$$ $\displaystyle e_{d}^{\top}\mu_{1,v}^{\star}(y,t,H)$ $\displaystyle=\int_{\widetilde{D}_{y,t,H}}v_{d}K(v)\,dv$ $\displaystyle=\int_{\mathbb{R}^{d-1}}\!\!dv_{1:d-1}\int_{-\sqrt{1-\sum_{i=1}^{d-1}v_{i}^{2}}}^{t/h_{n}(y)}v_{d}\,K(v_{1:d-1},v_{d})\,dv_{d}\;+\;O(\|H\|^{1/2})$ $\displaystyle=:g\!\left(\tfrac{t}{h_{n}(y)}\right)+O(\|H\|^{1/2}).$ $$

Note $g(0)<0$ and $g(\tau)$ decreases as $\tau\downarrow 0$. Choose $\alpha=\alpha_{1}\in(0,1)$ with $g(\alpha_{1})<0$. Then for all $t\in[0,\alpha_{1}h_{n}(y)]$,

$$ $e_{d}^{\top}\mu_{1,v}^{\star}(y,t,H)\leq g(\alpha_{1})+O(\|H\|^{1/2})<0.$ (23) $$

Combining equation [22](https://arxiv.org/html/2510.01450v1#A2.E22)–equation [23](https://arxiv.org/html/2510.01450v1#A2.E23), for small $H$,

$$ $\bigl|e_{d}^{\top}\bar{\mu}_{1,v}^{\star}(y,t,H)\bigr|=\frac{|e_{d}^{\top}\mu_{1,v}^{\star}(y,t,H)|}{\mu_{0,v}^{\star}(y,t,H)}\;\geq\;\frac{|g(\alpha_{1})|}{2D}.$ $$

Using equation [21](https://arxiv.org/html/2510.01450v1#A2.E21) yields the first bound with $c_{*}:=|g(\alpha_{1})|/(2D)$. The second bound follows from boundedness of $K$ and $\mathrm{supp}(K)\subset\mathbb{B}^{d}$.
∎

Since the leading term of $\mathrm{Bias}_{\mathrm{NW}}(x)$ is $\nabla f(x)^{\top}H^{1/2}\bar{\mu}_{1}^{\star}(x,H)$, the lower bound on $|e(y)^{\top}H^{1/2}\bar{\mu}_{1}^{\star}(y-te(y),H)|$ alone does not guarantee order $\Theta(\|H^{1/2}\|)$; we also require a sufficiently large normal derivative of $f$ along $\partial D$.

###### Definition B.1 (Extreme boundary gradient class) .

For any domain $D$ and constants $m$ and $M$, we define a class of functions $\mathcal{E}(D,m,M)$, where $f\in\mathcal{E}(D,m,M)$ iff there exist a measurable $\Gamma\subset\partial D$ with $S(\Gamma)>0$ and constants $m$, $M$ such that $|\partial_{e}f(y)|\geq m$ and $||\nabla_{T}f(y)||<M$ where $\partial_{e}f(y)=\nabla f(y)^{\top}e(y)$ and $\nabla_{T}f(y)=(I-e(y)e(y)^{\top})\nabla f(y)$.

Then we can prove that if the function to be estimated is within $\mathcal{E}(D,m,M)$ where $m$ and $M$ are specifically chosen constants that are independent of $H$, the NW has an integral squared bias with high order.

###### Lemma B.3 .

Under Assumptions [A.1](https://arxiv.org/html/2510.01450v1#A1.Thmassumption1), [A.2](https://arxiv.org/html/2510.01450v1#A1.Thmassumption2), and [A.3](https://arxiv.org/html/2510.01450v1#A1.Thmassumption3), if $f\in\mathcal{E}(D,m,M)$ with

$$ $c_{*}^{2}\,\kappa_{1}^{-1}m^{2}-2c_{*}\,\kappa_{1}^{-1/2}C_{*}\,mM\;\geq\;C_{1}>0,$ $$

then uniformly over $H\in\mathcal{H}_{n}$,

$$ $\int_{D}\mathrm{Bias}_{\mathrm{NW}}^{2}(x)\,dx=\Omega_{p}(\|H\|^{3/2})\qquad\text{and}\qquad\int_{D}\mathrm{Var}_{\mathrm{NW}}(x)\,dx=\Omega_{p}\!\bigl(n^{-1}|H|^{-1/2}\bigr).$ $$

###### Remark .

It is easy to construct $f$ and $D$ satisfying Assumptions [A.1](https://arxiv.org/html/2510.01450v1#A1.Thmassumption1), [A.2](https://arxiv.org/html/2510.01450v1#A1.Thmassumption2) and $f\in\mathcal{E}(D,m,M)$. For example, on $D=\{(x_{1},x_{2}):x_{1}^{2}+x_{2}^{2}\leq 1\}$, $f(x_{1},x_{2})=\frac{\sqrt{c_{1}\kappa_{1}}}{2c_{*}}(x_{1}^{2}+x_{2}^{2})$ works with suitable $m,M$.

###### Proof.

We first lower-bound bias. Choosing $\alpha$ as the minimum $\alpha$ given by Lemmas [A.3](https://arxiv.org/html/2510.01450v1#A1.Thmlemma3), [B.2](https://arxiv.org/html/2510.01450v1#A2.Thmlemma2), we have

$$ $\displaystyle\int\mathrm{Bias}_{\mathrm{NW}}^{2}(x)dx\geq\int_{\mathcal{B}(H)}\mathrm{Bias}_{\mathrm{NW}}^{2}(x)dx\geq\int_{\mathcal{C}(H,\alpha)}\mathrm{Bias}_{\mathrm{NW}}^{2}(x)dx$ (24) $\displaystyle=\int_{\partial D}\int_{0}^{\alpha h_{n}(y)}\mathrm{Bias}_{\mathrm{NW}}^{2}(y-te(y))\text{det}\left(I-tS_{y}\right)dtdS(y)$ $\displaystyle\geq C_{2}\int_{\Gamma}\int_{0}^{\alpha h_{n}(y)}\left((\nabla f(y-te(y))^{\top}H^{1/2}\bar{\mu}_{1}^{\star}(y-te(y),H))^{2}\;+\;O_{p}(||H^{3/2}||)\right)dtdS(y),$ $$

where the last inequality is derived from Lemma [A.2](https://arxiv.org/html/2510.01450v1#A1.Thmlemma2) and the boundedness of $S_{y}$ from Assumption [A.1](https://arxiv.org/html/2510.01450v1#A1.Thmassumption1).

For all $t\in\alpha h_{n}(y)$ and all $\eta\in(0,1)$,

$$ $\displaystyle\left(\nabla f(y-te(y))^{\top}H^{1/2}\bar{\mu}_{1}^{\star}(y-te(y),H)\right)^{2}$ $\displaystyle\overset{(a)}{=}\left(\partial_{e}f(y-te(y))e(y)^{\top}H^{1/2}\bar{\mu}_{1}^{\star}(y-te(y),H)+\nabla_{T}f(y-te(y))^{\top}H^{1/2}\bar{\mu}_{1}^{\star}(y-te(y),H)\right)^{2}$ $\displaystyle\overset{(b)}{\geq}(1-\eta)\left(\partial_{e}f(y-te(y))e(y)^{\top}H^{1/2}\bar{\mu}_{1}^{\star}(y-te(y),H)\right)^{2}-\eta^{-1}\left(\nabla_{T}f(y-te(y))^{\top}H^{1/2}\bar{\mu}_{1}^{\star}(y-te(y),H)\right)^{2}$ $\displaystyle\overset{(c)}{\geq}(1-\eta)(\partial_{e}f(y-te(y)))^{2}h_{n}^{2}(y)c_{*}^{2}-\eta^{-1}||\nabla_{T}f(y-te(y)||^{2}C_{*}^{2}||H||$ $\displaystyle=(1-\eta)(\partial_{e}f(y))^{2}h_{n}^{2}(y)c_{*}^{2}-\eta^{-1}||\nabla_{T}f(y)||^{2}C_{*}^{2}||H||+O(||H^{3/2}||)$ $\displaystyle\geq\left((1-\eta)m^{2}c_{*}^{2}\kappa_{1}^{-1}-\eta^{-1}M^{2}C_{*}^{2}\right)||H||+O(||H^{3/2}||).$ $$

Here $(a)$ uses the fact $\nabla f=\partial_{e}f(y)e(y)+\nabla_{T}f(y)$, $(b)$ uses the fact $(A+B)^{2}\geq(1-\eta)A^{2}-\eta^{-1}B^{2}$ for all $\eta\in(0,1)$, $(c)$ uses Lemma [B.2](https://arxiv.org/html/2510.01450v1#A2.Thmlemma2).

Setting $\eta=\frac{MC_{*}\kappa_{1}^{1/2}}{mc_{*}}$, we have

$$ $\displaystyle\left(\nabla f(y-te(y))^{\top}H^{1/2}\bar{\mu}_{1}^{\star}(y-te(y),H)\right)^{2}$ (25) $\displaystyle\geq\left((1-\eta)m^{2}c_{*}^{2}\kappa_{1}^{-1}-\eta^{-1}M^{2}C_{*}^{2}\right)||H||+O(||H^{3/2}||)$ $\displaystyle=\left(c_{*}^{2}\kappa_{1}^{-1}m^{2}-2c_{*}\kappa_{1}^{-1/2}C_{*}mM\right)||H||+O(||H^{3/2}||)$ $\displaystyle\geq C_{1}||H||++O(||H^{3/2}||).$ $$

Combining Equations [24](https://arxiv.org/html/2510.01450v1#A2.E24), [25](https://arxiv.org/html/2510.01450v1#A2.E25), we have that

$$ $\displaystyle\int\mathrm{Bias}_{\mathrm{NW}}^{2}(x)dx$ $\displaystyle\geq\alpha C_{2}\left(C_{1}||H||\;+\;O_{p}(||H^{3/2}||)\right)\int_{\Gamma}h_{n}(y)dS(y)=\Omega_{p}(||H^{3/2}||).$ $$

The last equation holds because $h_{n}(y)\geq\kappa_{1}^{-1}||H^{1/2}||$.
Using Lemma [A.2](https://arxiv.org/html/2510.01450v1#A1.Thmlemma2), it is straightforward to show that

$$ $\int_{\Omega}\mathrm{Var}_{\mathrm{NW}}(x)dx=\Omega_{p}(n^{-1}|H^{1/2}|^{-1}).$ $$

∎

We then show that local linear regression enjoys strictly lower bias than NW on all functions.

###### Lemma B.4 .

Under Assumptions [A.1](https://arxiv.org/html/2510.01450v1#A1.Thmassumption1), [A.2](https://arxiv.org/html/2510.01450v1#A1.Thmassumption2), [A.3](https://arxiv.org/html/2510.01450v1#A1.Thmassumption3), we have $\int_{D}\mathrm{Bias}_{\mathrm{LL}}^{2}(x)dx=O_{p}(||H^{2}||)$ and $\int_{D}\mathrm{Var}_{\mathrm{LL}}(x)dx=O_{p}(n^{-1}|H^{1/2}|^{-1})$ uniformly hold for all $H\in\mathcal{H}_{n}$.

###### Proof.

It is straightforward from Lemma [B.1](https://arxiv.org/html/2510.01450v1#A2.Thmlemma1).
∎

Then it is easy to prove the final conclusion.

###### Theorem B.1 (Precise statement of Proposition 2.2 ) .

Under Assumptions [A.1](https://arxiv.org/html/2510.01450v1#A1.Thmassumption1), [A.2](https://arxiv.org/html/2510.01450v1#A1.Thmassumption2), [A.3](https://arxiv.org/html/2510.01450v1#A1.Thmassumption3), if the function $f$ to be estimated is within $\mathcal{E}(D,m,M)$ where

$$ $c_{*}^{2}\kappa_{1}^{-1}m^{2}-2c_{*}\kappa_{1}^{-1/2}C_{*}mM\geq C_{1}>0,$ $$

we have $\mathbb{E}\int_{D}\mathrm{MSE}_{\mathrm{NW}}(x)dx=\Omega(n^{-3/(d+3)})$ and $\mathbb{E}\int_{D}\mathrm{MSE}_{\mathrm{LL}}(x)dx=O(n^{-4/(d+4)})$, where both $\widehat{f}_{\mathrm{NW}}$ and $\widehat{f}_{\mathrm{LL}}$ are at their optimal bandwidth $H\in\mathcal{H}_{n}$.

###### Proof.

According to Lemma [B.3](https://arxiv.org/html/2510.01450v1#A2.Thmlemma3) and the definition of $\mathcal{H}_{n}$, we have that we have $\int_{D}\mathrm{Bias}_{\mathrm{NW}}^{2}(x)dx=\Omega_{p}(h^{3})$ and $\int_{D}\mathrm{Var}_{\mathrm{NW}}(x)dx=\Omega_{p}(n^{-1}h^{-d})$ uniformly hold for all $h\in[n^{-a},n^{-b}]$. So we have

$$ $\int_{D}\mathrm{MSE}_{\mathrm{NW}}(x)dx=\Omega_{p}(h^{3}+\frac{1}{nh^{d}})=\Omega_{p}(n^{-3/(d+3)})$ $$

even for optimal $H\in\mathcal{H}_{n}$.
The last equality holds because $h=n^{-1/(d+3)}\in[n^{-a},n^{-b}]$ is the minimizer of $h^{3}+(nh^{d})^{-1}$.

According to Lemma [B.4](https://arxiv.org/html/2510.01450v1#A2.Thmlemma4), we have that

$$ $\int_{D}\mathrm{MSE}_{\mathrm{LL}}(x)dx=O_{p}(h^{4}+\frac{1}{nh^{d}})=O_{p}(n^{-4/(d+4)})$ $$

for the optimal $H\in\mathcal{H}_{n}$.
According to the definition of $\Omega_{p}$ and $O_{p}$, it is straightforward to deduce the conclusion.
∎

###### Lemma B.2 .

###### Proof.

###### Definition B.1 (Extreme boundary gradient class) .

###### Lemma B.3 .

###### Remark .

###### Proof.

###### Lemma B.4 .

###### Proof.

###### Theorem B.1 (Precise statement of Proposition 2.2 ) .

###### Proof.

## Appendix C Appendix: Conjugate Gradient Solver

The Conjugate Gradient (CG) method (Hestenes and Stiefel, 1952) is an iterative algorithm for solving systems of linear equations with symmetric positive-definite matrices.
As described in Section [3](https://arxiv.org/html/2510.01450v1#S3), we solve the linear systems with matrix $\Sigma_{i}$ for blocks of queries in parallel:

$$ $\displaystyle\Sigma_{i}x_{i}=y_{i},\quad\text{for }i\in\{(r-1)B_{r}+1,\ldots,rB_{r}\}$ (26) $$

In the FlashLLA forward algorithm [1](https://arxiv.org/html/2510.01450v1#alg1), CG is applied within each row block, with $Q_{r},M_{r},m_{r},\omega_{r}$ being the block quantities, $X_{r}$ being the solution to be computed and $Y_{r}$ being the right-hand side.
We use the simplest initialization $X_{r}\leftarrow 0$ for CG. Hence the initial residual $r_{i}$ is set to be $y_{i}$, i.e., $R^{(0)}\leftarrow Y_{r}$ in Algorithm [2](https://arxiv.org/html/2510.01450v1#alg2).

The core computation is the matrix-vector product $\Sigma_{i}p_{i}$ for the search vectors $p_{i}$ computed in the lines [4](https://arxiv.org/html/2510.01450v1#alg2.l4)-[10](https://arxiv.org/html/2510.01450v1#alg2.l10). The result is stored in matrix $\Sigma_{P}$.
This operation has high I/O intensity due to the requirement to stream through the entire $K$ matrix in HBM during each CG iteration.
Consequently, controlling the number of iterations is crucial for both the efficiency and convergence.
While maximal iteration number $T\leq 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 $\Sigma_{i}$ is low-rank and requires relatively large $\lambda$ to maintain positive definiteness.
To address this, we make the regularization $\lambda$ learnable and data-dependent:

$$ $\displaystyle\lambda_{i}=\mathrm{sigmoid}(W_{\lambda}x_{i}).$ (27) $$

The dimension of the weight $W_{\lambda}\in\mathbb{R}^{d\times d_{\lambda}}$ controls the granularity of the regularization. Setting $d_{\lambda}=d$ enables per-dimensional regularization,
though empirically set $d_{\lambda}=d_{h}$ suffices, where $d_{h}$ denotes the head dimension.

Additionally, since the CG solves multiple systems in parallel, different system converge at 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 $\epsilon$.

Figure: Algorithm 2 FlashLLA CG Solver

## Appendix D Appendix: Backward Derivation

This section provides the detailed derivation of the backward pass.
Defining the following variables:

$$ $\displaystyle\hat{b}_{i}$ $\displaystyle=\sum_{j=1}^{i}s_{ij}v_{j}=\sum_{j=1}^{i}w_{ij}\frac{n_{ij}}{\delta_{i}}v_{j},\quad g_{i}=\frac{\partial\mathcal{L}}{\partial\hat{b}_{i}}\in\mathbb{R}^{d}$ (28) $\displaystyle\gamma_{ij}$ $\displaystyle=g_{i}^{\top}v_{j}=\frac{\partial\mathcal{L}}{\partial s_{ij}},\quad\beta_{i}=\frac{1}{\delta_{i}}\sum_{j=1}^{i}\gamma_{ij}{s_{ij}},\quad c_{ij}=\frac{\gamma_{ij}w_{ij}}{\delta_{i}}$ (29) $$

The gradient of the loss with respect to $v_{j}$ is given by:

$$ $\displaystyle\frac{\partial\mathcal{L}}{\partial v_{j}}$ $\displaystyle=\sum_{i=j}^{n}s_{ij}g_{i},\quad\Delta_{V}=S^{\top}G$ (30) $$

The gradient of $q_{i}$ and $k_{j}$ is related to $s_{ij}$, which can be broken down into $w_{ij}$ and $z_{ij}$ path in the computation graph.
The partial gradients of the loss with respect to $w_{ij},z_{ij}$ are given by:

$$ $\displaystyle\frac{\partial\mathcal{L}}{\partial w_{ij}}$ $\displaystyle=\frac{\partial\mathcal{L}}{\partial s_{ij}}\frac{\partial s_{ij}}{\partial w_{ij}}+\frac{\partial\mathcal{L}}{\partial\omega_{i}}\frac{\partial\omega_{i}}{\partial w_{ij}}+\frac{\partial\mathcal{L}}{\partial\mu_{i}}^{\top}\frac{\partial\mu_{i}}{\partial w_{ij}}+\operatorname{Tr}\left(\frac{\partial\mathcal{L}}{\partial\Sigma_{i}}^{\top}\frac{\partial\Sigma_{i}}{\partial w_{ij}}\right)$ (31) $\displaystyle=\frac{\gamma_{ij}n_{ij}}{\delta_{i}}-\beta_{i}+z_{ij}^{\top}\frac{\partial\mathcal{L}}{\partial\mu_{i}}+z_{ij}^{\top}\frac{\partial\mathcal{L}}{\partial\Sigma_{i}}z_{ij}$ (32) $\displaystyle\frac{\partial\mathcal{L}}{\partial z_{ij}}$ $\displaystyle=\frac{\partial\mathcal{L}}{\partial s_{ij}}\frac{\partial s_{ij}}{\partial z_{ij}}+\frac{\partial\mathcal{L}}{\partial\mu_{i}}^{\top}\frac{\partial\mu_{i}}{\partial z_{ij}}+\operatorname{Tr}\left(\frac{\partial\mathcal{L}}{\partial\Sigma_{i}}^{\top}\frac{\partial\Sigma_{i}}{\partial w_{ij}}\right)$ (33) $\displaystyle=-c_{ij}\rho_{i}+w_{ij}\frac{\partial\mathcal{L}}{\partial\mu_{i}}+2w_{ij}\frac{\partial\mathcal{L}}{\partial\Sigma_{i}}z_{ij}$ (34) $$

Then the partial gradients of the loss with respect to $k_{j}$ and $q_{i}$ are given by:

$$ $\displaystyle\frac{\partial\mathcal{L}}{\partial k_{j}}$ $\displaystyle=\sum_{i=j}^{n}\frac{\partial\mathcal{L}}{\partial w_{ij}}\frac{\partial w_{ij}}{\partial k_{j}}+\frac{\partial\mathcal{L}}{\partial z_{ij}}\frac{\partial z_{ij}}{\partial k_{j}}=\sum_{i=j}^{n}\frac{\partial\mathcal{L}}{\partial w_{ij}}\frac{w_{ij}}{h}q_{i}+\frac{\partial\mathcal{L}}{\partial z_{ij}}$ (35) $\displaystyle\frac{\partial\mathcal{L}}{\partial q_{i}}$ $\displaystyle=\sum_{j=1}^{i}\frac{\partial\mathcal{L}}{\partial w_{ij}}\frac{\partial w_{ij}}{\partial q_{i}}+\frac{\partial\mathcal{L}}{\partial z_{ij}}\frac{\partial z_{ij}}{\partial q_{i}}=\sum_{j=1}^{i}\frac{\partial\mathcal{L}}{\partial w_{ij}}\frac{w_{ij}}{h}k_{j}-\frac{\partial\mathcal{L}}{\partial z_{ij}}$ (36) $$

In order to compute the partial gradients of $\mu_{i}$ and $\Sigma_{i}$, denote

$$ $\displaystyle u_{i}=\Sigma_{i}^{-1}\sum_{j=1}^{i}c_{ij}z_{ij}=\Sigma_{i}^{-1}\biggl[\sum_{j=1}^{i}c_{ij}k_{j}-\Bigl(\sum_{j=1}^{i}c_{ij}\Bigr)q_{i}\biggr]$ (37) $$

which can be computed with existing Conjugate Gradient solver [2](https://arxiv.org/html/2510.01450v1#alg2). Then

$$ $\displaystyle\frac{\partial\mathcal{L}}{\partial\mu_{i}}$ $\displaystyle=\sum_{j=1}^{i}\frac{\partial\mathcal{L}}{\partial n_{ij}}\frac{\partial n_{ij}}{\partial\mu_{i}}+\frac{\partial\mathcal{L}}{\partial\delta_{i}}\frac{\partial\delta_{i}}{\partial\mu_{i}}=-u_{i}+2\beta_{i}\rho_{i}$ (38) $\displaystyle\frac{\partial\mathcal{L}}{\partial\Sigma_{i}}$ $\displaystyle=\sum_{j=1}^{i}\frac{\partial\mathcal{L}}{\partial n_{ij}}\frac{\partial n_{ij}}{\partial\Sigma_{i}}+\frac{\partial\mathcal{L}}{\partial\delta_{i}}\frac{\partial\delta_{i}}{\partial\Sigma_{i}}=-\frac{1}{2}\rho_{i}\frac{\partial\mathcal{L}}{\partial\mu_{i}}^{\top}+\frac{1}{2}u_{i}\rho_{i}^{\top}$ (39) $$

We denote the following variables:

$$ $\displaystyle\Delta_{\mu}=-U+2\texttt{bcast}(\beta)\odot R$ (40) $$

We can materialize the gradient of $w_{ij}$ for every $i,j$ pair and then perform the reduction.

$$ $\displaystyle w_{ij}\frac{\partial\mathcal{L}}{\partial w_{ij}}=\gamma_{ij}s_{ij}+w_{ij}z_{ij}^{\top}\frac{\partial\mathcal{L}}{\partial\mu_{i}}+w_{ij}z_{ij}^{\top}\frac{\partial\mathcal{L}}{\partial\Sigma_{i}}z_{ij}$ (41) $$

We omit the $Q,K$ in the relmm for simplicity, then the gradient of $w_{ij}$ can be computed as:

$$ $\displaystyle\Delta_{W}$ $\displaystyle=\Gamma\odot S+W\odot(-\texttt{bcast}(\beta)+\texttt{relmm}(\Delta_{\mu})-\frac{1}{2}\texttt{relmm}(\Delta_{\mu})\odot\texttt{relmm}(R)$ (42) $\displaystyle+\frac{1}{2}\texttt{relmm}(U)\odot\texttt{relmm}(R))$ (43) $$

Then we can compute the gradient of $k_{j}$ and $q_{i}$ through $w_{ij}$ branch as follows:

$$ $\displaystyle\Delta_{W}^{K}$ $\displaystyle=\frac{1}{h}\Delta_{W}^{\top}Q,\quad\Delta_{W}^{Q}=\frac{1}{h}\Delta_{W}K$ (44) $$

For the $z_{ij}$ path, we avoid materializing the third-order tensor by performing the reduction internally,

$$ $\displaystyle\sum_{i=j}^{n}\frac{\partial\mathcal{L}}{\partial z_{ij}}$ $\displaystyle=\sum_{i=j}^{n}-c_{ij}\rho_{i}+\sum_{i=j}^{n}w_{ij}\frac{\partial\mathcal{L}}{\partial\mu_{i}}+2\sum_{i=j}^{n}w_{ij}\frac{\partial\mathcal{L}}{\partial\Sigma_{i}}z_{ij}$ (45) $\displaystyle\sum_{j=1}^{i}\frac{\partial\mathcal{L}}{\partial z_{ij}}$ $\displaystyle=\sum_{j=1}^{i}-c_{ij}\rho_{i}+\sum_{j=1}^{i}w_{ij}\frac{\partial\mathcal{L}}{\partial\mu_{i}}+2\sum_{j=1}^{i}w_{ij}\frac{\partial\mathcal{L}}{\partial\Sigma_{i}}z_{ij}$ (46) $$

Then we can compute the gradient of the loss with respect to $z_{ij}$ as follows:

$$ $\displaystyle\Delta_{Z}^{K}$ $\displaystyle=-C^{\top}R+W^{\top}\Delta_{\mu}-(W\odot\texttt{relmm}(\Delta_{\mu}))^{\top}R+(W\odot\texttt{relmm}(R))^{\top}U$ (47) $\displaystyle\Delta_{Z}^{Q}$ $\displaystyle=-\texttt{brsum}(C)\odot R+\texttt{brsum}(W)\odot\Delta_{\mu}$ (48) $\displaystyle\quad-\texttt{brsum}(W\odot\texttt{relmm}(\Delta_{\mu}))\odot R+\texttt{brsum}(W\odot\texttt{relmm}(R))\odot U$ (49) $$

Hence the gradient of the loss with respect to $k_{j}$ and $q_{i}$ can be computed as:

$$ $\displaystyle\Delta_{K}$ $\displaystyle=\Delta_{W}^{K}+\Delta_{Z}^{K},\quad\Delta_{Q}=\Delta_{W}^{Q}+\Delta_{Z}^{Q}$ (50) $$

Figure: Figure 7: Test-time regression with across all input dimension $d$ and segment size $S$.
Refer to caption: x15.png

## Appendix E Appendix: Experiments

### E.1 Piecewise Linear Data Generation.

Let $n=2^{m}$ be the number of segment with $n=\log_{2}n\leq d$. For each section index $c\in\{1,\ldots,n\}$, define a sign pattern
$S_{c}=(s_{c,1},s_{c,2},\ldots,s_{c,m})\in\{-1,+1\}^{m}$
by reading the least-significant bits of $c$.
For each segment in each sample, draw $Z\sim\mathcal{N}(0,I_{d})$ and construct the data $X\in\mathbb{R}^{d}$ by flipping the first $m$ coordinate of $Z$:

$$ $\displaystyle X_{j}=\begin{cases}S_{c,j}|Z_{j}|,&j\leq m\\ Z_{j},&j>m\end{cases}$ (51) $$

As the result, the constructed segment is a truncated Gaussian conditioned to lie in the cone $\mathcal{C}_{c}=\{x\in\mathbb{R}^{d}:s_{c,j}x_{j}\geq 0,j=1,\ldots,m\}$ where $\mathcal{C}_{i}\cap\mathcal{C}_{j}=\emptyset$.
Denote $T_{c}(Z)=(S_{c}\odot|Z_{1:m}|,|Z_{m+1:d}|)$ and segment distribution $P_{c}$, we have

$$ $\displaystyle X\sim P_{c}\iff X\overset{d}{=}T_{c}(Z)$ (52) $$

### E.2 Training Configuration

#### Test-Time Regression.

This experiment evaluates test-time adaptation without training any model parameters. We sweep over input dimension
$d\in\{8,16,32,64,128\}$ and segment size
$S\in\{8,16,32,64,128,256,512,1024\}$ with fixed sequence length
$L=1024$. Performance is evaluated by averaging the mean squared error over
$10{,}000$ independently generated sequences.

#### In-Context Regression.

We fix the input and output dimensions at
$d_{x}=d_{y}=32$. For each random seed, we generate
$100{,}000$ training examples with noise level $\delta=0.1$ and
$1{,}000$ test examples with $\delta=0$ (noiseless evaluation). For each sequence length
$L\in\{64,128,256,512\}$, we sweep over segment size
$S\in\{L/8,L/4,L/2,L\}$ and learning rate
$\{5\times 10^{-5},10^{-4},5\times 10^{-4},10^{-3}\}$.
All models are trained with the AdamW optimizer
$(\beta_{1},\beta_{2})=(0.9,0.999)$, weight decay $0.1$, batch size
$256$, for a maximum of $100$ epochs.

#### In-Context Associative Recall.

We fix the vocabulary size $|A_{k}\cup A_{v}|=8{,}192$.
For each sequence length $L\in\{64,128,256,512\}$,
we sweep over the number of key-value pairs
$\{L/16,L/8,L/4\}$ and learning rate
$\{10^{-4},5\times 10^{-4},10^{-3}\}$.
We generate $20{,}000$, $40{,}000$, and $60{,}000$ training examples
and $1{,}000$ test examples each for the respective key-value pair counts.
Training uses AdamW with $(\beta_{1},\beta_{2})=(0.9,0.999)$,
weight decay $0.1$, batch size $256$, for a 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\in\{16,24,48,96\}$, we sample the number of instructions
$S\sim\text{Uniform}(N/6,N/3)$ and use $8$ queries per example.
We sweep over learning rate $\{10^{-4},5\times 10^{-4},10^{-3}\}$.
Models are trained with AdamW $(\beta_{1},\beta_{2})=(0.9,0.999)$,
weight decay $0.1$, batch size $256$, for a maximum of $64$ epochs.

### E.3 Additional Experiment Results

Figure: Figure 8: Test accuracy curves for associative recall across all sequence lengths $L$ (averaged over 3 random seeds). Results for different numbers of key-value pairs are averaged within each sequence length for visual clarity
Refer to caption: x16.png

In Figure [7](https://arxiv.org/html/2510.01450v1#A4.F7), we provide the full test-time regression results across all input dimension $d$ and segment size $S$.
The advantages of LLA scales with the dimension $d$ and nonstationarity.
In practical settings, as the dimensionality increases, it is less likely to do exact query $q=k_{j}$ 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 to the ideal conditions of this synthetic experiment. Nevertheless, the overall performance trends remain consistent with our main findings.

Figure [8](https://arxiv.org/html/2510.01450v1#A5.F8) shows complete associative recall results across all sequence lengths $L$.
For visual clarity, we 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 is highly sensitive to hyperparameters such as learning rate and dataset size.

In contrast, LLA demonstrates consistent, gradual improvement in test accuracy with a corresponding 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 learning 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.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 2
presentation: 2
questions: 1.	The 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.
2.	The 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.
3.	The 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.
4.	The 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.
soundness: 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.
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.
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.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: Please refer to my weakness part.
soundness: 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.
-  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.
- The authors clearly recognized the computational hurdles of LLA and put significant effort into addressing them.
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.

As 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.
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.
- 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.

---

confidence: 3
contribution: 4
initial_rating: 6
presentation: 4
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.

2. 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?

3. 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?

4. 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.
soundness: 3
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.

2. 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.

3. The synthetic experiments indicate that LLA exhibits a distinct superiority over other attention mechanisms in certain abilities, including in-context regression and associative recall.
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.

The 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.
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.

2. 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.

---

confidence: 4
contribution: 2
initial_rating: 2
presentation: 2
questions: 1. How does the local polynomial regression mentioned in line 197 relate to the proposed method?
2. In line 236 (and in the title), what exactly does “optimal” refer to? In what sense is it optimal?
3. 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?
soundness: 2
strengths: - The paper presents a theoretically motivated architectural design.
- It proposes a practical, GPU-aware implementation of the algorithm.
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.
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.
2. 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.
3. 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.
4. 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.
5. 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.


Decision: Accept (Poster)