LoLA: Low-Rank Linear Attention With Sparse Caching

## Abstract

Abstract The per-token cost of transformer inference scales with context length, preventing its application to lifelong in-context learning. Linear attention is an efficient alternative that maintains a constant memory footprint, even on infinite context lengths. While this is a potential candidate for lifelong learning, it falls short in memory capacity. In this paper, we propose LoLA, a training-free augmentation to linear attention that boosts associative recall. LoLA distributes past key-value pairs from context into three memory systems: (i) recent pairs in a local sliding window cache; (ii) difficult-to-memorize pairs in a sparse, global cache; and (iii) generic pairs in the recurrent hidden state of linear attention. We show through ablations that our self-recall error metric is crucial to efficiently manage long-term associative memories. On pass-key retrieval tasks, LoLA improves the base model’s performance from 0.6% to 97.4% accuracy. This is achieved with a 4.6 × \times smaller cache than Llama-3.1 8B on 4K context length. LoLA also outperforms other 1B and 8B parameter subquadratic models on zero-shot commonsense reasoning tasks.

## 1 Introduction

Transformer-based large language models (LLMs) rely on storing all past tokens in an ever-growing key-value (KV) cache (Vaswani et al., 2017). This allows future query tokens to access past memories with associative recall, which enables in-context learning (Olsson et al., 2022). Since no previous information is discarded, the KV cache continues to grow with context length. This eventually leads to a memory bottleneck on long context tasks, such as lifelong in-context learning. As a result, transformers cannot condition next token predictions on arbitrarily long sequences.

Alternative architectures to transformers have been proposed—such as Mamba (Gu & Dao, 2024), DeltaNet (Schlag et al., 2021), linear attention (Katharopoulos et al., 2020), and others (Yang et al., 2024a; Behrouz et al., 2024; Sun et al., 2024)—to reduce the compute complexity from quadratic to linear. Additionally, these approaches reduce the memory cost from linear to constant.
In particular, linear attention removes the exponential dot product in softmax (Katharopoulos et al., 2020). This effectively collapses the unbounded KV-cache into a fixed-size matrix, which corresponds to a recurrently formed hidden state (i.e., a linear RNN). This constructs a linear associative memory map from keys to values. Past memories can be recalled through a vector-matrix product of an incoming query vector and the hidden state matrix. Linear attention enables constant-cost prediction per token when conditioned on arbitrarily long contexts.

While efficient and flexible, linear attention architectures lag behind transformers in terms of memory capacity. This is largely noticeable on tasks leveraging in-context learning (Paperno et al., 2016; Hsieh et al., 2024). The removal of the exponential dot product allows for non-orthogonal keys to interfere with the hidden state’s learned key-to-value map. This interference—denoted as a *memory collision*—impairs associative recall. Previous work used nonlinear query and key activations to improve the exponential dot product approximation (Choromanski et al., 2020; Zhang et al., 2024). However, these attempts are essentially performing a low-rank approximation of the infinite-rank exponential dot product kernel.

Additional use of sparse attention (Chen et al., 2021) can improve linear attention’s recall; however,
current hybrid approaches only focus on local information with sliding window attention (Arora et al., 2024; Zhang et al., 2025a; Lan et al., 2025; Van Nguyen et al., 2025). These approaches can struggle to recall critical, long-term facts that fall outside the window.

This raises our fundamental research question:

How can long term associative memory for subquadratic language models be improved?

### Contributions

We present LoLA: Low-rank Linear Attention with sparse caching. LoLA is a novel, training-free inference strategy that boosts the performance of hybrid linear attention layers. LoLA distributes historical tokens into three forms of memory: (i) recent KV pairs are stored in a sliding window cache, (ii) difficult-to-memorize pairs in a sparse global cache, and (iii) all other pairs are placed in a recurrent hidden-state matrix via linear attention. LoLA performs a self-recall check to see which KV pairs disagree with the current hidden state’s linear associative map. LoLA sparsely caches the interfering memories in full rank. The selection mechanism effectively mitigates memory collisions with a small, constant-sized cache. This inference strategy can be applied on top of previously trained linear attention + sliding window models (e.g., LoLCATs) to significantly improve associative recall. As a result, LoLA extracts stronger language modeling capabilities from the same base model weights.

Utilizes Self-Recall Error. We introduce an importance metric for key-value pairs to reduce memory collisions in the hidden state of linear attention. This is computed by determining if a key can recall its own value with linear attention. In our ablations, we show that this performance increase cannot be obtained from using a larger sliding window or other sparse attention metrics: self-recall is essential.

Enables Associative Recall. As a lightweight inference strategy, LoLA enables pass-key retrieval on up to 8K context lengths in needle-in-a-haystack tasks from the RULER benchmark (Hsieh et al., 2024). With a 4.6x smaller cache than Llama-3.1 8B (Grattafiori et al., 2024), our approach boosts accuracy from LoLCATs’ 0.6% to 97.4% at 4K context lengths with the same model weights.

Improves Language Modeling.
LoLA shows superior performance on zero-shot commonsense reasoning tasks among 1B and 8B parameter subquadratic architectures. This demonstrates that effective memory management can boost language modeling performance.

## 2 Preliminaries

In this section, we review softmax attention through the lens of associative memory. Then, we show how linear attention naturally forms a recurrent hidden state. We address practical implementations for training linear architectures and highlight unresolved drawbacks of previous approaches.

### 2.1 Softmax Attention as a Nonparametric, Online Learner

Transformers process a sequence of input tokens $\{\bm{x}_{t}\}_{t=1}^{n}$, for $\bm{x}_{t}\in\mathbb{R}^{d}$ (Vaswani et al., 2017). For each attention head, the input tokens are transformed into three distinct representations—queries, keys, and values—via trainable weight matrices $\mathbf{W}_{q},\mathbf{W}_{k}\in\mathbb{R}^{d\times d_{k}}$ and $\mathbf{W}_{v}\in\mathbb{R}^{d\times d_{v}}$. For a given token $\bm{x}_{t}$, define

$$ $\underbrace{\bm{q}_{t}=\mathbf{W}_{q}\bm{x}_{t}}_{\text{query}},\quad\underbrace{\bm{k}_{t}=\mathbf{W}_{k}\bm{x}_{t}}_{\text{key}},\quad\underbrace{\bm{v}_{t}=\mathbf{W}_{v}\bm{x}_{t}}_{\text{value}}.$ (1) $$

Causal attention uses the current query to recall past information from key-value pairs. The similarity between the query $\bm{q}_{t}$ and key $\bm{k}_{i}$ is denoted as $\alpha_{ti}\in(0,1)$. This similarity score determines how much value $\bm{v}_{i}$ is used for the current output token at time $t$. The output token $\bm{y}_{t}$ is defined by

$$ $\bm{y}_{t}=\sum_{i=1}^{t}\alpha_{ti}\,\bm{v}_{i}\in\mathbb{R}^{d_{v}},\quad\text{with}\quad\alpha_{ti}=\frac{\exp\left(\bm{q}_{t}^{\top}\bm{k}_{i}/\sqrt{d_{k}}\right)}{\sum_{j=1}^{t}\exp\left(\bm{q}_{t}^{\top}\bm{k}_{j}/\sqrt{d_{k}}\right)}.$ (2) $$

We view softmax attention as a nonparametric function $m_{t}:\mathbb{R}^{d_{k}}\to\mathbb{R}^{d_{v}}$ that fits to the past context at inference time. This online function $m_{t}$ learns to map keys to their associated value in context with $m_{t}(\bm{k}_{i})\approx\bm{v}_{i}$ for $(\bm{k}_{i},\bm{v}_{i})\in\{(\bm{k}_{i},\bm{v}_{i})\}_{i=1}^{t}$. Then, $m_{t}$ applies the learned transformation to the query, $\bm{y}_{t}=m_{t}(\bm{q}_{t})$. The set of past key-value pairs forms an online "training set" of input-output labels. The query acts as an unsupervised "test set".

Softmax attention caches all past key-value pairs to perform this non-parametric, or “look-up table”, operation. Since the function complexity scales with the context length, softmax attention can flexibly learn new context without forgetting past key-value associations. However, this process leads to an unbounded KV-cache size that scales linearly with sequence length $n$. Ultimately, this operation cannot be used for extremely long context scenarios, such as lifelong learning.

### 2.2 Linear Attention

To bound inference costs, linear attention methods replace the exponential dot product kernel (Katharopoulos et al., 2020) with a low-rank approximation. This enables models to maintain constant-size memory footprints even for infinite sequence lengths. With this replacement,

$$ $\exp\left(\frac{\bm{q}_{t}^{\top}\bm{k}_{j}}{\sqrt{d_{k}}}\right)\approx\phi(\bm{q}_{t})^{\top}\phi(\bm{k}_{j}),\quad\text{for}\quad\phi:\mathbb{R}^{d_{k}}\rightarrow\mathbb{R}^{D},$ (3) $$

the output token is approximated as

$$ $\displaystyle\bm{y}_{t}^{\top}=\sum_{i=1}^{t}\frac{\exp\left(\bm{q}_{t}^{\top}\bm{k}_{i}/\sqrt{d_{k}}\right)\;\bm{v}_{i}^{\top}}{\sum_{j=1}^{t}\exp\left(\bm{q}_{t}^{\top}\bm{k}_{j}/\sqrt{d_{k}}\right)}\approx\sum_{i=1}^{t}\frac{\phi(\bm{q}_{t})^{\top}\phi(\bm{k}_{i})\;\bm{v}_{i}^{\top}}{\sum_{j=1}^{t}\phi(\bm{q}_{t})^{\top}\phi(\bm{k}_{j})}$ $\displaystyle=\frac{\phi(\bm{q}_{t})^{\top}\left(\sum_{j=1}^{t}\phi(\bm{k}_{j})\;\bm{v}_{j}^{\top}\right)}{\phi(\bm{q}_{t})^{\top}\left(\sum_{j=1}^{t}\phi(\bm{k}_{j})\right)}$ $\displaystyle=\frac{\phi(\bm{q}_{t})^{\top}\mathbf{H}_{t}}{\phi(\bm{q}_{t})^{\top}\bm{s}_{t}}.$ (4) $$

This creates a hidden state matrix $\mathbf{H}_{t}\in\mathbb{R}^{D\times d_{v}}$ as the sum of key-value outer products. This effectively bounds the memory cost to $\mathcal{O}(Dd_{v})$, constant with respect to sequence length $n$. The hidden dimension size $D$ controls the approximation quality at the cost of computational efficiency. As a linear RNN, the hidden state $\mathbf{H}_{t}$ and normalization state $\bm{s}_{t}\in\mathbb{R}^{D}$ can be computed recurrently,

$$ $\mathbf{H}_{t}=\mathbf{H}_{t-1}+\phi(\bm{k}_{t})\,\bm{v}_{t}^{\top},\quad\bm{s}_{t}=\bm{s}_{t-1}+\phi(\bm{k}_{t}).$ (5) $$

In this formulation, linear attention stores each observation, or KV-pair, as a rank-one outer product. Rather than building a look-up table, linear attention parameterizes the key-to-value map as a linear function. While this approach is efficient, linear attention falls short in memory capacity as the number of orthogonal key-value pairs is bounded by the rank of $\mathbf{H}_{t}$. “Memory collisions” (Yang et al., 2024a) occur when new hidden state updates overwrite past key-value associations. This prevents the linear map from accurately modeling the context.

### 2.3 Efficient Training of Linear Attention

To reduce training costs, LoLCATs (Zhang et al., 2025a) and others (Bick et al., 2024; Wang et al., 2024; Bick et al., 2025; Goldstein et al., 2025) recycle large pretrained transformers into linear attention models with knowledge distillation (Hinton et al., 2015). These approaches minimize the difference between the pretrained transformer’s output $\bm{y}$ (i.e., the teacher) and linear attention’s output $\widehat{\bm{y}}$ (i.e., the student). In particular, LoLCATs uses a trainable nonlinear map for $\phi:\mathbb{R}^{d_{k}}\to\mathbb{R}^{D}$, constructed as

$$ $\phi(\bm{x})=\left[\exp(\bm{w}_{1}^{\top}\bm{x}),\ldots,\exp(\bm{w}_{D/2}^{\top}\,\bm{x}),\exp(-\bm{w}_{1}^{\top}\bm{x}),\ldots,\exp(-\bm{w}_{D/2}^{\top}\,\bm{x})\right]\in\mathbb{R}^{D},$ (6) $$

with learnable weights $\bm{w}_{i}\in\mathbb{R}^{d_{k}}$ (Zhang et al., 2024). This distillation approach freezes all other parameters, adjusting $\phi$ to minimize the loss

$$ $\mathcal{L}(\phi)=\sum_{\bm{q},\bm{k},\bm{v}}\|\bm{y}_{t}-\widehat{\bm{y}}_{t}\|,\quad\text{with}\quad\widehat{\bm{y}}_{t}=\frac{\phi(\bm{q}_{t})^{\top}\mathbf{H}_{t}}{\phi(\bm{q}_{t})^{\top}\bm{s}_{t}}.$ (7) $$

After attention distillation, the whole model is finetuned with LoRA (Hu et al., 2022). Overall, this procedure only requires 40 million training tokens from the Alpaca dataset (Taori et al., 2023), grouped in 1024-long sequences.

### 2.4 Disadvantages of Previous Approaches

Even with distillation, linear attention models struggle to accurately mimic the behavior of softmax attention. Initial work in linear attention proposed nonlinear query and key activations to improve the exponential dot product approximation (Choromanski et al., 2020; Zhang et al., 2024). These methods fall short as they are essentially performing a low-rank approximation of the infinite-rank exponential dot product kernel. In Appendix [E](https://arxiv.org/html/2505.23666v2#A5), we show that the exponential dot product kernel has slowly decaying singular values for simple data distributions. This implies that high-dimensional hidden states may be required for modest approximation errors.

Recent approaches attempt to address the poor performance of the low-rank approximation in linear attention by augmenting it with sliding window attention (Arora et al., 2024; Zhang et al., 2025a). These methods compute a *finite* number of recent tokens in a window with softmax attention and compute the rest with linear attention. Since natural language contains a significant amount of local information, this hybrid approach nearly recovers the performance of pretrained transformers on short-context tasks. For longer sequences, however, these methods struggle to recall important information that falls outside the window and in the hidden state. We show in Table [1](https://arxiv.org/html/2505.23666v2#S4.T1) that these models cannot perform associative recall on simple needle-in-a-haystack tasks. Other forms of sparse attention may be required alongside "low-rank" attention (Chen et al., 2021).

## 3 Mitigating Memory Collisions with Sparse Caching

#### Identifying Difficult-to-Remember KV Pairs.

As our base assumption, strong associative memory systems should allow keys to retrieve their associated values. Online functions $m_{t}$ with perfect recall interpolate the online training set, defined as

$$ $m_{t}(\bm{k}_{i})=\bm{v}_{i},\quad\forall i\leq t$ (8) $$

For perfect recall in linear attention, equation [8](https://arxiv.org/html/2505.23666v2#S3.E8) translates to

$$ $m_{t}(\bm{k}_{i})^{\top}=\frac{\phi(\bm{k}_{i})^{\top}\mathbf{H}_{t}}{\phi(\bm{k}_{i})^{\top}\bm{s}_{t}}=\frac{\phi(\bm{k}_{i})^{\top}\sum_{j=1}^{t}\phi(\bm{k}_{j})\,\bm{v}_{j}^{\top}}{\phi(\bm{k}_{i})^{\top}\sum_{j=1}^{t}\phi(\bm{k}_{j})}=\bm{v}_{i}^{\top}.$ (9) $$

In practice, however, "memory collisions" prohibit equation [9](https://arxiv.org/html/2505.23666v2#S3.E9) from holding. Non-orthogonal keys interfere with each other when forming the hidden state $\mathbf{H}_{t}$. We measure the Self-Recall Error (SRE) to measure how well a past key can retrieve its associated value with

$$ $\text{SRE}(\bm{k},\bm{v}\,|\,\mathbf{H}_{t},\bm{s}_{t})=\left\|\frac{\phi(\bm{k})^{\top}\mathbf{H}_{t}}{\phi(\bm{k})^{\top}\bm{s}_{t}}-\bm{v}\right\|_{2}=\left\|\hat{\bm{v}}-\bm{v}\right\|_{2}.$ (10) $$

This determines the error between the predicted value $\hat{\bm{v}}$ for a given key $\bm{k}$ and the ground truth value $\bm{v}$.

#### Method Overview.

We propose LoLA: a training-free inference strategy that boosts the performance of hybrid linear attention models. LoLA addresses the limitations of previous linear attention mechanisms by integrating a sparse caching strategy at inference time. This method employs three memory systems to store long term associations

- 1.
Linear Attention utilizes a finite-rank approximation to store an infinite amount of tokens.
- 2.
Sliding Window Attention provides full-rank attention scores for finite, local context.
- 3.
Sparse Caching identifies and stores key-value pairs that are challenging to remember, preventing memory collisions in linear attention.

Figure: Figure 1: LoLA stores past KV pairs in three forms memory for each attention head.
Refer to caption: x1.png

LoLA uses the self-recall error, equation [10](https://arxiv.org/html/2505.23666v2#S3.E10), to decide which KV pairs should be stored separately in full-rank. Large errors indicate the severity of the memory collision. As a result, LoLA keeps the KV pairs with the largest error in a sparse cache. This limits the corruption of past memories and improves associative recall.

Since only a finite amount of tokens can be stored at each time step to maintain efficiency, LoLA performs a greedy scoring approach. At every iteration, LoLA scores the KV pairs leaving the sliding window and re-scores all pairs currently stored in the sparse cache. At each timestep, the pairs with the lowest error are moved to the linear hidden state indefinitely. Re-scoring pairs in the sparse cache is essential since the SRE is dependent on the current hidden state. For example, a KV pair could become more aligned with the hidden state in the future after a few updates. In the generation implementation of LoLA, we define the set of pairs that are scored at time $t$ as

$$ $\mathcal{E}_{t}=G_{t-1}\,\cup\,\{(\bm{k}_{t-\eta},\bm{v}_{t-\eta})\},$ (11) $$

where $G_{t-1}$ is the set of KV pairs in the sparse cache at time $t-1$. Here, $\eta$ is the maximum number of pairs in the sliding window. We update the sparse cache by selecting the top-$\lambda$ errors in $\mathcal{E}_{t}$, i.e,

$$ $G_{t}=\operatorname*{arg\,max}_{G\subset\mathcal{E}_{t}:|G|=\lambda}\sum_{(\bm{k},\bm{v})\in G}\text{SRE}(\bm{k},\bm{v}\,|\,\mathbf{H}_{t},\bm{s}_{t}).$ (12) $$

The remaining pairs, denoted by $\mathcal{S}_{t}=\mathcal{E}_{t}\cap G_{t}^{\mathsf{c}}$ where $G_{t}^{\mathsf{c}}$ is the complement of $G_{t}$, are stored in hidden state via

$$ $\mathbf{H}_{t}=\mathbf{H}_{t-1}+\sum_{(\bm{k},\bm{v})\in\mathcal{S}_{t}}\phi(\bm{k})\bm{v}^{\top},\quad\bm{s}_{t}=\bm{s}_{t-1}+\sum_{(\bm{k},\bm{v})\in\mathcal{S}_{t}}\phi(\bm{k})$ (13) $$

Once the caches are up to date, LoLA computes the output token $\bm{y}_{t}$ as

$$ $\bm{y}_{t}=\frac{\overbrace{\phi(\bm{q}_{t})^{\top}\mathbf{H}_{t}}^{\textit{Linear Attn.}}+\overbrace{\sum_{i\in G_{t}}\exp\left(\bm{q}_{t}^{\top}\bm{k}_{i}/\sqrt{d}\right)\bm{v}_{i}}^{\textit{Sparse Cache}}+\overbrace{\sum_{j=t-\eta+1}^{t}\exp\left(\bm{q}_{t}^{\top}\bm{k}_{j}/\sqrt{d}\right)\bm{v}_{j}}^{\textit{Sliding Window}}}{\phi(\bm{q}_{t})^{\top}\bm{s}_{t}+\sum_{i\in G_{t}}\exp\left(\bm{q}_{t}^{\top}\bm{k}_{i}/\sqrt{d}\right)+\sum_{j=t-\eta+1}^{t}\exp\left(\bm{q}_{t}^{\top}\bm{k}_{j}/\sqrt{d}\right)}.$ (14) $$

Both $\eta$ and $\lambda$ are hyper-parameters for the size of the sliding window and sparse cache, respectively.

#### Chunkwise Inference.

When the input sequence is available ahead of time (e.g., prefill), LoLA is accelerated with parallelization. By partitioning the input sequence into chunks of size $C$, we can compute intra-chunk operations in parallel with dense matmuls (Yang et al., 2024b). This reduces the number of recurrent iterations by a factor of $C$ while preserving the constant-memory cost that motivates LoLA.

LoLA computes softmax attention with the current chunk of queries and previous two chunks of KV-pairs.
For small chunk sizes, softmax attention is almost equally efficient to linear attention. Artificially limiting softmax within a chunk will not improve efficiency, only hurt performance.
The past two chunks of KV-pairs are concatenated with the sparse cache in order to compute a single FlashAttention (Dao et al., 2022) pass per chunk of queries. For the linear attention portion of the forward pass, all queries within the chunk share the same hidden state.

After computing the past chunk of output tokens, we evict the oldest chunk of KV-pairs in the window, sending them to the hidden state or sparse cache. All of the evicted and sparse cache pairs are scored with the SRE, equation [10](https://arxiv.org/html/2505.23666v2#S3.E10). The $\lambda$ pairs with the largest errors in the eligible set,

$$ $\mathcal{E}_{t}=G_{t-1}\cup\{(\mathbf{k}_{i},\mathbf{v}_{i})\mid t-2C\leq i<t-C\},$ (15) $$

will remain in the sparse cache. The remainder are integrated into the hidden state through the standard outer product update as in equation [5](https://arxiv.org/html/2505.23666v2#S2.E5).

Figure: Figure 2: Illustration of where each KV pair is stored at every time step for each method.
Refer to caption: x2.png

The given hardware setup dictates the total fixed cache size of LoLA, but the ratio of chunk size to sparse cache size depends on the application.
Increasing the chunk or sliding window size reduces the number of recurrent iterations and overhead costs from sparse caching; however, this requires more VRAM. Increasing the sparse cache size $\lambda$ will better mitigate collisions in the hidden state and improve long context recall. In Appendix [B](https://arxiv.org/html/2505.23666v2#A2), we explore this trade-off for various cache hyperparameters in an efficiency analysis. Specifically we measure the total VRAM use, Time-to-First-Token, and long context performance. Furthermore, we illustrate the bounded nature of LoLA, compared to vanilla transformers.

## 4 Experiments & Results

In our experiments, we leverage the same attention distillation procedure in LoLCATs to obtain the base model, then apply our inference strategy, LoLA, at test time. To train the base model, we replace each attention module in Llama-3.1 8B (or Llama-3.2 1B) with a hybrid sliding window + linear attention module. We use a sliding window size $\eta=64$ for training and use a trainable feature map for $\phi$ as described in equation [6](https://arxiv.org/html/2505.23666v2#S2.E6). The output dimension of $\phi$ is $D=2d_{k}$. First, we freeze all non-attention layers in the linearized Transformer and only train $\phi$ with distillation for two epochs on Alpaca (Taori et al., 2023). Then, we perform LoRA (Hu et al., 2022) finetuning on the whole model for two epochs. This procedure only uses 40M training tokens with 1024-long sequences.

### 4.1 Associative Recall

To see how LoLA improves associative recall, we conduct a study on Single-Needle-in-a-Haystack (S-NIAH) tasks from RULER (Hsieh et al., 2024). In Table [1](https://arxiv.org/html/2505.23666v2#S4.T1), we compare LoLA to the base model, LoLCATs-8B, and variants with an extended sliding window for a fair comparison. We observe LoLCATs struggles to recall information outside the sliding window. Extending the sliding window size marginally improves performance. We explain the differences of each NIAH task in Appendix [C](https://arxiv.org/html/2505.23666v2#A3) and discuss these results more in depth.

Table: Table 1: Measuring long context recall with Needle-in-a-Haystack (NIAH) tasks from the RULER benchmark. We report recall accuracy for each method across different context lengths (512, 1024, etc.) for each task.

Next in Table [2](https://arxiv.org/html/2505.23666v2#S4.T2), we measure the performance of LoLA on the rest of the RULER benchmark at 4K context length. This covers much harder long context tasks, such as multi-key retrieval (MK1,MK2,MK3), multi-query (MQ), multi-value (MV), variable tracking (VT), common word extraction (CWE), frequent word extraction (FWE), Hotpot-QA (HQA), and Squad-QA (SQA). We compare LoLA against a stronger version of LoLCATs—with an equivalent, larger cache size ($\eta=896$)—and Mamba2-8B (Dao & Gu, 2024; Waleffe et al., 2024).

**Table 2: Extended RULER Benchmark on 4K context lengths. Compared to NIAH tasks, these require much stronger forms of memory and state tracking. For cache parameters, LoLA uses $\eta=128,\lambda=768$ and LoLCATs uses $\eta=896$.**
| Model | MK1 | MK2 | MK3 | MQ | MV | VT | CWE | FWE | HQA | SQA | Avg |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Mamba2-8B | 40.3 | 13.8 | 5.5 | 49.1 | 35.0 | 76.5 | 32.9 | 76.6 | 31.8 | 35.5 | 39.7 |
| LoLCATs-8B+ | 12.8 | 1.4 | 0.4 | 3.3 | 3.6 | 0.7 | 3.0 | 13.6 | 14.2 | 14.0 | 6.7 |
| LoLA-8B | 39.4 | 11.6 | 7.6 | 67.6 | 65.0 | 85.2 | 45.9 | 51.3 | 24.2 | 53.9 | 45.2 |

LoLA improves recall with minimal additional caching. Table [1](https://arxiv.org/html/2505.23666v2#S4.T1) demonstrates an improvement from the base model’s 0.6% to 97.4% accuracy on S-NIAH-1. This is achieved with a $4.6\times$ smaller cache than Llama with $\eta=256,\lambda=256$. Furthermore, we show in Table [2](https://arxiv.org/html/2505.23666v2#S4.T2) that sparse caching is essential for more difficult tasks, improving an extended form of LoLCATs from 6.7% average accuracy to 45.2%. For example, tasks such as variable tracking (VT) require understanding all of the context. Since no part of the sequence can be lost, naive metrics for sparse attention, such as (Zhang et al., 2023), cannot be used. Memory collisions must be mitigated. Sparse caching—specifically with our self-recall error—unlocks a new capability for hybrid linear attention architectures.

### 4.2 Commonsense Reasoning

We demonstrate the language modeling performance of LoLA on various zero-shot commonsense reasoning tasks using LM evaluation harness (Gao et al., 2024). To compare against previously available approaches, we use PIQA (PI) (Bisk et al., 2020), ARC-Easy (AE) & ARC-Challenge (AC) (Clark et al., 2018), HellaSwag (HS) (Zellers et al., 2019), WinoGrande (WG) (Sakaguchi et al., 2021), MMLU (MM) (Hendrycks et al., 2020), and Lambada OpenAI (Paperno et al., 2016)).

In Table [3](https://arxiv.org/html/2505.23666v2#S4.T3), we compare LoLA against other 7-9B subquadratic models (Mamba (Gu & Dao, 2024), Mamba2 (Dao & Gu, 2024), RWKV-6 (Peng et al., 2024), Hawk & Griffin (De et al., 2024), Falcon Mamba (Zuo et al., 2024), RecurrentGemma (Botev et al., 2024), Mamba-in-the-Llama (Wang et al., 2024), Llamba (Bick et al., 2025), Hedgehog (Zhang et al., 2024), and LoLCATs (Zhang et al., 2025a)). Models that use any form of unbounded global attention (e.g interleaving SSM blocks and Transformer blocks) still retain quadratic compute complexity and growing memory costs. These are outside the scope of this work. We also report the number of training tokens used to create each model in both tables. Though LoLA is a training-free inference strategy that can be used for any sliding window + linear attention model, we report the cost of distilling the base subquadratic model (Zhang et al., 2025a). In Appendix [A](https://arxiv.org/html/2505.23666v2#A1.SS0.SSS0.Px1), we show results for 1-2B subquadratic models. Additionally, we provide a direct comparison of distilled models by measuring the average accuracy relative to their teacher models.

**Table 3: Performance comparison of 7-9B parameter fixed-memory models across various common sense reasoning tasks: PIQA (PI), Arc-Easy (AE), ARC-Challenge (AC), Winogrande (WG), MMLU (MM), and Lambada-openai (LB). Bolded scores are the best and underlined scores are the second best. We report accuracy for all applicable, except AC and HS use normalized logits. MMLU (MM) uses 5-shot. Reported scores were compiled from (Bick et al., 2025; Zhang et al., 2025a; Waleffe et al., 2024; Wang et al., 2024). * indicates our reproduced score is used and is higher than reported score. LoLA uses small cache parameters, $\eta=\lambda=64$.**
| Model | Tokens (B) | PI | AE | AC | HS | WG | MM | LB |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Transformers |  |  |  |  |  |  |  |  |
| Llama-3.1-8B | 15000 | 81.1 | 81.7 | 55.1 | 79.3 | 73.9 | 68.0 | 73.0 |
| Subquadratic: Pretrained from scratch |  |  |  |  |  |  |  |  |
| Mamba-8B | 1100 | 78.9 | 75.4 | 42.2 | 75.6 | 68.3 | 28.0 | - |
| Mamba2-8B | 3500 | 79.8 | 75.9 | 48.1 | 77.7 | 71.6 | 48.7 | - |
| RWKV-6 (W2.1) 7B | 1420 | 78.7 | 76.8 | 46.3 | 75.1 | 70.0 | - | - |
| Hawk 7B | 300 | 80.0 | 74.4 | 45.9 | 77.6 | 69.9 | 35.0 | - |
| Griffin 7B | 300 | 81.0 | 75.4 | 47.9 | 78.6 | 72.6 | 39.3 | - |
| Falcon3-Mamba-7B | 7300 | 79.7 | 72.5 | 53.2 | 79.8 | 69.1 | 65.0 | 67.5 |
| RecurrentGemma-9B | 2000 | 80.6 | 78.9 | 57.1 | 80.1 | 73.7 | 55.1 | 54.1 |
| Subquadratic: Distilled from Llama-3.1-8B |  |  |  |  |  |  |  |  |
| Mamba2-Llama3-8B (L3.1-Instr.) | 20 | 76.8 | 74.1 | 48.0 | 70.8 | 58.6 | 43.2 | - |
| Hedgehog-8B (Llama-3) | 0.04 | 77.4 | 71.1 | 40.6 | 66.5 | 54.3 | 24.2 | - |
| Llamba-8B | 12 | 80.9 | 82.5 | 54.6 | 77.6 | 73.3 | 60.0 | 69.4 |
| LoLCATs-8B | 0.04 | 81.0 | 82.4 | 54.4 | 79.1 | 73.6* | 54.9 | 67.6 |
| LoLA-8B (ours) | 0.04 | 81.6 | 82.5 | 55.4 | 79.8 | 73.6 | 57.6 | 74.9 |

On short context tasks such as Winogrande, we observe additional caching is not needed as only local information is required for good performance. On the other hand, we find that gaps still exist with Lambada and MMLU. Lambada requires longer context reasoning, leading to significant improvements with sparse caching. Furthermore, Bick et al. (Bick et al., 2025) suggest that dataset selection plays a large role for MMLU performance. Though sparse caching shows significant improvement on MMLU, a more powerful distillation procedure or dataset may be needed to reach Llama’s performance (Goldstein et al., 2025).

For 1B parameter models, sparse caching provides even more utility. Since the hidden state dimensionality scales with the head dimension of the base model, 1B models can face more memory collisions. In Appendix [A](https://arxiv.org/html/2505.23666v2#A1.SS0.SSS0.Px1), LoLA demonstrates state-of-the-art performance among 1B subquadratic models, even outperforming Llama-3.2-1B on average. Overall, LoLA pushes the pareto front for training-efficient and high-performing subquadratic LLMs.

### 4.3 Understanding Memory Collisions

We visualize how memory collisions occur in practice. At each time step $t$, we measure the self-recall error from equation [10](https://arxiv.org/html/2505.23666v2#S3.E10) for every KV pair that is currently stored in the hidden state $\mathbf{H}_{t}$. We visualize the error for linear attention, sliding window + linear attention, and LoLA. We use a sliding window size of $\eta=64$ tokens and a sparse cache size of $\lambda=64$ when applicable.

Figure: Figure 3: Visualizing memory collisions by measuring SRE for stored KV pairs.
Refer to caption: x3.png

When only using linear attention, we observe large recall errors. At early time steps, the first few KV pairs receive small errors, but quickly become larger after hidden state updates. In Appendix [G](https://arxiv.org/html/2505.23666v2#A7), we additionally plot the relative error to better show how this occurs. These stored associative memories become corrupted and tough to recall in the future. Furthermore, difficult-to-memorize pairs are evident, illustrated as bright columns in Figure [3](https://arxiv.org/html/2505.23666v2#S4.F3). The additional use of sliding window attention only delays the inevitable memory collisions.

LoLA significantly reduces the errors for all KV pairs. Difficult-to-memorize pairs are appropriately stored in the sparse cache, as seen by the zero-columns in the plot. This also prevents corrupting older KV pairs that are already stored in the hidden state.

#### Scoring Method Ablation.

In our final experiment, we measure alternative scoring functions to understand which KV pairs should be sparsely cached. Traditional sparse attention metrics assume "unimportant" tokens are evicted entirely from the context (Zhang et al., 2023; Singhania et al., 2024; Zaheer et al., 2020). In our setting, these assumptions are invalid as unimportant tokens are stored in low precision through linear attention.

Table: Table 4: Ablation results for various scoring methods on S-NIAH-1 with 512 context length, $\eta=64,\lambda=64$. Extended details for the score calculation can be found in Appendix [F](https://arxiv.org/html/2505.23666v2#A6).

In Table [4](https://arxiv.org/html/2505.23666v2#S4.T4), we observe that storing keys with poor exponential dot product approximations underperforms the naive extension of sliding window attention. This hints that a better softmax approximation should not be the main objective for linear attention methods. Keys over-estimated by linear attention seem to be “more important” than local keys; however, all tested alternatives fall short of enabling associative recall.

We also compare against traditional sparse attention ideas that use softmax attention scores as a proxy for importance, such as in H2O (Zhang et al., 2023) and LESS (Dong et al., 2024). Specifically, we found that a key’s average similarity score, $\exp(\bm{q}^{\top}\bm{k})$, does not translate well in the hybrid linear attention setting. Since queries are used to influence the selection of keys in the sparse cache, we believe highly similar phrases in the haystack may interfere with finding the needle. LoLA, on the other hand, benefits from its query-agnostic metric.

## 5 Conclusion

LoLA integrates linear attention with sparse caching to effectively mitigate memory collisions. By selectively retaining KV pairs that do not align with the current hidden state, LoLA enables passkey retrieval when the base model fails. Our experimental results demonstrate that targeted sparse caching substantially improves long context performance over naively increasing the sliding window size. LoLA demonstrates strong language modeling performance over other 1B or 8B subquadratic models.

Future work. The sparse cache carries a small overhead compute cost of $\mathcal{O}(\lambda d)$ for scoring. For high-complexity, long-context tasks, we found that larger sparse cache sizes are needed to reduce interference in the hidden state. We believe that these limitations can addressed in the future with a better base subquadratic model. More advanced base architectures, such as LaCT (Zhang et al., 2025b) or Atlas (Behrouz et al., 2025), use nonlinear key-to-value maps, which may lead to smaller caches.

## Reproducibility Statement

For reproducibility of distilling base subquadratic model, see (Zhang et al., 2025a). This also includes released weights for the 8B parameter model on Huggingface. All evaluations were performed using LM Evaluation Harness (Gao et al., 2024). pseudo-code is available for reproducing LoLA’s attention operation in Appendix [H](https://arxiv.org/html/2505.23666v2#A8). This is a drop-in replacement for the hybrid linear attention in the base model. Full code will be provided in the camera-ready version. Lastly, LLMs had minor contributions to the paper writing, such as spell-check and formatting.

## References

- Arora et al. (2024)
Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, James Zou, Atri Rudra, and Christopher Re.
Simple linear attention language models balance the recall-throughput tradeoff.
In *International Conference on Machine Learning*, pp. 1763–1840. PMLR, 2024.
- Beck et al. (2024)
Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael K Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter.
xlstm: Extended long short-term memory.
In *The Thirty-eighth Annual Conference on Neural Information Processing Systems*, 2024.
- Behrouz et al. (2024)
Ali Behrouz, Peilin Zhong, and Vahab Mirrokni.
Titans: Learning to memorize at test time.
*arXiv preprint arXiv:2501.00663*, 2024.
- Behrouz et al. (2025)
Ali Behrouz, Zeman Li, Praneeth Kacham, Majid Daliri, Yuan Deng, Peilin Zhong, Meisam Razaviyayn, and Vahab Mirrokni.
Atlas: Learning to optimally memorize the context at test time.
*arXiv preprint arXiv:2505.23735*, 2025.
- Beltagy et al. (2020)
Iz Beltagy, Matthew E. Peters, and Arman Cohan.
Longformer: The long-document transformer, 2020.
URL [https://arxiv.org/abs/2004.05150](https://arxiv.org/abs/2004.05150).
- Bick et al. (2024)
Aviv Bick, Kevin Li, Eric Xing, J Zico Kolter, and Albert Gu.
Transformers to ssms: Distilling quadratic knowledge to subquadratic models.
*Advances in Neural Information Processing Systems*, 37:31788–31812, 2024.
- Bick et al. (2025)
Aviv Bick, Tobias Katsch, Nimit Sohoni, Arjun Desai, and Albert Gu.
Llamba: Scaling distilled recurrent models for efficient language processing.
*arXiv preprint arXiv:2502.14458*, 2025.
- Bisk et al. (2020)
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al.
Piqa: Reasoning about physical commonsense in natural language.
In *Proceedings of the AAAI conference on artificial intelligence*, volume 34, pp. 7432–7439, 2020.
- Botev et al. (2024)
Aleksandar Botev, Soham De, Samuel L Smith, Anushan Fernando, George-Cristian Muraru, Ruba Haroun, Leonard Berrada, Razvan Pascanu, Pier Giuseppe Sessa, Robert Dadashi, et al.
Recurrentgemma: Moving past transformers for efficient open language models.
*CoRR*, 2024.
- Chen et al. (2021)
Beidi Chen, Tri Dao, Eric Winsor, Zhao Song, Atri Rudra, and Christopher Ré.
Scatterbrain: Unifying sparse and low-rank attention.
*Advances in Neural Information Processing Systems*, 34:17413–17426, 2021.
- Choromanski et al. (2020)
Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, et al.
Rethinking attention with performers.
In *International Conference on Learning Representations*, 2020.
- Clark et al. (2018)
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord.
Think you have solved question answering? try arc, the ai2 reasoning challenge.
*arXiv preprint arXiv:1803.05457*, 2018.
- Dao & Gu (2024)
Tri Dao and Albert Gu.
Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality.
In *International Conference on Machine Learning (ICML)*, 2024.
- Dao et al. (2022)
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré.
Flashattention: Fast and memory-efficient exact attention with io-awareness.
*Advances in neural information processing systems*, 35:16344–16359, 2022.
- De et al. (2024)
Soham De, Samuel L Smith, Anushan Fernando, Aleksandar Botev, George-Cristian Muraru, Albert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, et al.
Griffin: Mixing gated linear recurrences with local attention for efficient language models.
*CoRR*, 2024.
- Dong et al. (2024)
Harry Dong, Xinyu Yang, Zhenyu Zhang, Zhangyang Wang, Yuejie Chi, and Beidi Chen.
Get more with less: Synthesizing recurrence with kv cache compression for efficient llm inference.
*arXiv preprint arXiv:2402.09398*, 2024.
- Eckart & Young (1936)
Carl Eckart and Gale Young.
The approximation of one matrix by another of lower rank.
*Psychometrika*, 1(3):211–218, 1936.
- Gao et al. (2024)
Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac’h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou.
A framework for few-shot language model evaluation, 07 2024.
URL [https://zenodo.org/records/12608602](https://zenodo.org/records/12608602).
- Glorioso et al. (2024)
Paolo Glorioso, Quentin Anthony, Yury Tokpanov, James Whittington, Jonathan Pilault, Adam Ibrahim, and Beren Millidge.
Zamba: A compact 7b ssm hybrid model.
*arXiv preprint arXiv:2405.16712*, 2024.
- Goldstein et al. (2025)
Daniel Goldstein, Eric Alcaide, Janna Lu, and Eugene Cheah.
Radlads: Rapid attention distillation to linear attention decoders at scale.
*arXiv preprint arXiv:2505.03005*, 2025.
- Grattafiori et al. (2024)
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al.
The llama 3 herd of models.
*arXiv preprint arXiv:2407.21783*, 2024.
- Gu & Dao (2024)
Albert Gu and Tri Dao.
Mamba: Linear-time sequence modeling with selective state spaces.
*arXiv preprint arXiv:2312.00752*, 2024.
URL [https://arxiv.org/abs/2312.00752](https://arxiv.org/abs/2312.00752).
- Gu et al. (2021)
Albert Gu, Karan Goel, and Christopher Re.
Efficiently modeling long sequences with structured state spaces.
In *International Conference on Learning Representations*, 2021.
- Hendrycks et al. (2020)
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt.
Measuring massive multitask language understanding.
*arXiv preprint arXiv:2009.03300*, 2020.
- Hinton et al. (2015)
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean.
Distilling the knowledge in a neural network.
*arXiv preprint arXiv:1503.02531*, 2015.
- Hsieh et al. (2024)
Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg.
Ruler: What’s the real context size of your long-context language models?
*arXiv preprint arXiv:2404.06654*, 2024.
- Hu et al. (2022)
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen, et al.
Lora: Low-rank adaptation of large language models.
*ICLR*, 1(2):3, 2022.
- Katharopoulos et al. (2020)
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret.
Transformers are rnns: Fast autoregressive transformers with linear attention.
In *International conference on machine learning*, pp. 5156–5165. PMLR, 2020.
- Lan et al. (2025)
Disen Lan, Weigao Sun, Jiaxi Hu, Jusen Du, and Yu Cheng.
Liger: Linearizing large language models to gated recurrent structures.
*arXiv preprint arXiv:2503.01496*, 2025.
- Li et al. (2023)
Yuanzhi Li, Sébastien Bubeck, Ronen Eldan, Allie Del Giorno, Suriya Gunasekar, and Yin Tat Lee.
Textbooks are all you need ii: phi-1.5 technical report.
*arXiv preprint arXiv:2309.05463*, 2023.
- Liu et al. (2024)
Aixin Liu, Bei Feng, Bin Wang, Bingxuan Wang, Bo Liu, Chenggang Zhao, Chengqi Dengr, Chong Ruan, Damai Dai, Daya Guo, et al.
Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model.
*arXiv preprint arXiv:2405.04434*, 2024.
- (32)
Jean Mercat, Igor Vasiljevic, Sedrick Scott Keh, Kushal Arora, Achal Dave, Adrien Gaidon, and Thomas Kollar.
Linearizing large language models.
In *First Conference on Language Modeling*.
- Nawrot et al. (2025)
Piotr Nawrot, Robert Li, Renjie Huang, Sebastian Ruder, Kelly Marchisio, and Edoardo M Ponti.
The sparse frontier: Sparse attention trade-offs in transformer llms.
*arXiv preprint arXiv:2504.17768*, 2025.
- Olsson et al. (2022)
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al.
In-context learning and induction heads.
*CoRR*, 2022.
- Paperno et al. (2016)
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández.
The lambada dataset: Word prediction requiring a broad discourse context.
*arXiv preprint arXiv:1606.06031*, 2016.
- Penedo et al. (2024)
Guilherme Penedo, Hynek Kydlícek, Loubna Ben Allal, and Thomas Wolf.
Fineweb: Decanting the web for the finest text data at scale.
*HuggingFace. Accessed: Jul*, 12, 2024.
- Peng et al. (2024)
Bo Peng, Daniel Goldstein, Quentin Gregory Anthony, Alon Albalak, Eric Alcaide, Stella Biderman, Eugene Cheah, Teddy Ferdinan, Kranthi Kiran GV, Haowen Hou, et al.
Eagle and finch: Rwkv with matrix-valued states and dynamic recurrence.
In *First Conference on Language Modeling*, 2024.
- Peng et al. (2020)
Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah Smith, and Lingpeng Kong.
Random feature attention.
In *International Conference on Learning Representations*, 2020.
- (39)
Zhen Qin, Weixuan Sun, Hui Deng, Dongxu Li, Yunshen Wei, Baohong Lv, Junjie Yan, Lingpeng Kong, and Yiran Zhong.
cosformer: Rethinking softmax in attention.
In *International Conference on Learning Representations*.
- Qin et al. (2024)
Zhen Qin, Songlin Yang, Weixuan Sun, Xuyang Shen, Dong Li, Weigao Sun, and Yiran Zhong.
Hgrn2: Gated linear rnns with state expansion.
In *First Conference on Language Modeling*, 2024.
- Ren et al. (2025)
Liliang Ren, Yang Liu, Yadong Lu, Chen Liang, Weizhu Chen, et al.
Samba: Simple hybrid state space models for efficient unlimited context language modeling.
In *The Thirteenth International Conference on Learning Representations*, 2025.
- Sakaguchi et al. (2021)
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi.
Winogrande: An adversarial winograd schema challenge at scale.
*Communications of the ACM*, 64(9):99–106, 2021.
- Schlag et al. (2021)
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber.
Linear transformers are secretly fast weight programmers.
In *International conference on machine learning*, pp. 9355–9366. PMLR, 2021.
- Siems et al. (2025)
Julien Siems, Timur Carstensen, Arber Zela, Frank Hutter, Massimiliano Pontil, and Riccardo Grazzi.
Deltaproduct: Increasing the expressivity of deltanet through products of householders.
In *ICLR 2025 Workshop on Foundation Models in the Wild*, 2025.
- Singhania et al. (2024)
Prajwal Singhania, Siddharth Singh, Shwai He, Soheil Feizi, and Abhinav Bhatele.
Loki: Low-rank keys for efficient sparse attention.
In *The Thirty-eighth Annual Conference on Neural Information Processing Systems*, 2024.
- Sun et al. (2024)
Yu Sun, Xinhao Li, Karan Dalal, Jiarui Xu, Arjun Vikram, Genghan Zhang, Yann Dubois, Xinlei Chen, Xiaolong Wang, Sanmi Koyejo, et al.
Learning to (learn at test time): Rnns with expressive hidden states.
*arXiv preprint arXiv:2407.04620*, 2024.
- Sun et al. (2023)
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei.
Retentive network: A successor to transformer for large language models, 2023.
URL [https://arxiv.org/abs/2307.08621](https://arxiv.org/abs/2307.08621).
- Taori et al. (2023)
Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto.
Stanford alpaca: An instruction-following llama model.
[https://github.com/tatsu-lab/stanford_alpaca](https://github.com/tatsu-lab/stanford_alpaca), 2023.
- Van Nguyen et al. (2025)
Chien Van Nguyen, Ruiyi Zhang, Hanieh Deilamsalehy, Puneet Mathur, Viet Dac Lai, Haoliang Wang, Jayakumar Subramanian, Ryan A Rossi, Trung Bui, Nikos Vlassis, et al.
Lizard: An efficient linearization framework for large language models.
*arXiv preprint arXiv:2507.09025*, 2025.
- Vaswani et al. (2017)
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.
Attention is all you need.
*Advances in neural information processing systems*, 30, 2017.
- Waleffe et al. (2024)
Roger Waleffe, Wonmin Byeon, Duncan Riach, Brandon Norick, Vijay Korthikanti, Tri Dao, Albert Gu, Ali Hatamizadeh, Sudhakar Singh, Deepak Narayanan, et al.
An empirical study of mamba-based language models.
*arXiv preprint arXiv:2406.07887*, 2024.
- Wang et al. (2024)
Junxiong Wang, Daniele Paliotta, Avner May, Alexander Rush, and Tri Dao.
The mamba in the llama: Distilling and accelerating hybrid models.
*Advances in Neural Information Processing Systems*, 37:62432–62457, 2024.
- Wang et al. (2025)
Ke Alexander Wang, Jiaxin Shi, and Emily B Fox.
Test-time regression: a unifying framework for designing sequence models with associative memory.
*arXiv preprint arXiv:2501.12352*, 2025.
- Yang et al. (2024a)
Songlin Yang, Jan Kautz, and Ali Hatamizadeh.
Gated delta networks: Improving mamba2 with delta rule.
*arXiv preprint arXiv:2412.06464*, 2024a.
- Yang et al. (2024b)
Songlin Yang, Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim.
Parallelizing linear transformers with the delta rule over sequence length.
In *The Thirty-eighth Annual Conference on Neural Information Processing Systems*, 2024b.
- Yuan et al. (2025)
Jingyang Yuan, Huazuo Gao, Damai Dai, Junyu Luo, Liang Zhao, Zhengyan Zhang, Zhenda Xie, YX Wei, Lean Wang, Zhiping Xiao, et al.
Native sparse attention: Hardware-aligned and natively trainable sparse attention.
*arXiv preprint arXiv:2502.11089*, 2025.
- Zaheer et al. (2020)
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al.
Big bird: Transformers for longer sequences.
*Advances in neural information processing systems*, 33:17283–17297, 2020.
- Zellers et al. (2019)
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi.
Hellaswag: Can a machine really finish your sentence?
In *Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics*, 2019.
- Zhang et al. (2024)
Michael Zhang, Kush Bhatia, Hermann Kumbong, and Christopher Re.
The hedgehog & the porcupine: Expressive linear attentions with softmax mimicry.
In *The Twelfth International Conference on Learning Representations*, 2024.
- Zhang et al. (2025a)
Michael Zhang, Simran Arora, Rahul Chalamala, Alan Wu, Benjamin Spector, Aaryan Singhal, Krithik Ramesh, and Christopher Ré.
Lolcats: On low-rank linearizing of large language models.
In *The Thirteenth International Conference on Learning Representations*, 2025a.
- Zhang et al. (2025b)
Tianyuan Zhang, Sai Bi, Yicong Hong, Kai Zhang, Fujun Luan, Songlin Yang, Kalyan Sunkavalli, William T Freeman, and Hao Tan.
Test-time training done right.
*arXiv preprint arXiv:2505.23884*, 2025b.
- Zhang et al. (2023)
Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, et al.
H2o: Heavy-hitter oracle for efficient generative inference of large language models.
*Advances in Neural Information Processing Systems*, 36:34661–34710, 2023.
- Zuo et al. (2024)
Jingwei Zuo, Maksim Velikanov, Dhia Eddine Rhaiem, Ilyas Chahed, Younes Belkada, Guillaume Kunsch, and Hakim Hacid.
Falcon mamba: The first competitive attention-free 7b language model.
*arXiv preprint arXiv:2410.05355*, 2024.

## Appendix A Extended Language Modeling Results

#### 1B parameter model results.

Following Section [4.2](https://arxiv.org/html/2505.23666v2#S4.SS2), we extend this comparison in Table [5](https://arxiv.org/html/2505.23666v2#A1.T5) for various 1-2B parameter subquadratic models (Grattafiori et al., 2024; Li et al., 2023; Sun et al., 2023; Qin et al., 2024; Schlag et al., 2021; Yang et al., 2024a; Gu & Dao, 2024; Dao & Gu, 2024; Beck et al., 2024; Peng et al., 2024; Bick et al., 2024; Zhang et al., 2025a; Bick et al., 2025; Ren et al., 2025). Compared to the 8B models, we observe that sparse caching is more important in the 1B parameter regime since there exist more memory collisions. This is a direct result of a smaller hidden state dimension. The size of $\mathbf{H}_{t}$ scales with the head dimension of the base model. Llama-3.2 1B’s head dimension is half that of Llama-3.1 8B.

**Table 5: Performance comparison across zero-shot commonsense reasoning tasks for various 1-2B parameter subquadratic models. * indicates normalized logits were reported instead.**
| Model | Tokens | PIQA | ARC-e | ARC-c | Hella. | Wino. | LMB. | LMB. |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | (B) | acc $\uparrow$ | acc $\uparrow$ | acc_n $\uparrow$ | acc_n $\uparrow$ | acc $\uparrow$ | acc $\uparrow$ | ppl $\downarrow$ |
| Transformers (Bick et al., 2024; Zhang et al., 2025a) |  |  |  |  |  |  |  |  |
| Llama-3.2-1B | 9000 | 74.4 | 65.5 | 35.8 | 63.7 | 60.5 | 60.1 | - |
| Phi-1.5-1.3B | 150 | 76.6 | 75.6 | 48.0 | 62.6 | 73.4 | 53.4 | - |
| Subquadratic: Pretrained from scratch on FineWeb-Edu (Yang et al., 2024a; Penedo et al., 2024) |  |  |  |  |  |  |  |  |
| RetNet-1.3B | 100 | 70.1 | 67.3 | 33.8 | 49.2 | 54.1 | 40.5 | 17.3 |
| HGRN2-1.3B | 100 | 70.5 | 69.4 | 35.3 | 49.5 | 52.8 | 39.5 | 17.7 |
| DeltaNet-1.3B | 100 | 70.7 | 68.5 | 35.7 | 50.9 | 53.4 | 42.5 | 16.9 |
| Gated-DeltaNet-1.3B | 100 | 72.3 | 71.2 | 38.4 | 55.8 | 57.5 | 46.7 | 12.2 |
| Mamba1-1.3B | 100 | 71.3 | 69.5 | 35.4 | 52.9 | 53.0 | 44.0 | 15.1 |
| Mamba2-1.3B | 100 | 71.9 | 72.5 | 37.9 | 55.7 | 55.2 | 45.7 | 12.6 |
| Subquadratic: Pretrained from scratch on various sources (Bick et al., 2024; 2025) |  |  |  |  |  |  |  |  |
| Mamba1-1.4B | 315 | 74.2 | 65.5 | 32.8 | 59.1 | 61.5 | 64.9 | - |
| Mamba2-1.3B | 315 | 73.2 | 64.3 | 33.3 | 59.9 | 60.9 | 65.7 | - |
| xLSTM-1.4B | 300 | 74.6 | 64.3 | 32.6 | 60.9 | 60.6 | 57.8 | - |
| Finch-1.6B | 1100 | 72.6 | 64.2 | 34.1 | 57.3 | 59.4 | 66.8 | - |
| RecurrentGemma-2B | 2000 | 67.2 | 35.6 | 51.2 | 60.3 | 55.7 | 52.5 | - |
| Samba-1.3B | 100 | 72.4 | 58.2 | - | 54.7 | 55.7 | 51.7 | - |
| Subquadratic: Distilled from Phi-1.5-1.3B (Bick et al., 2024; Zhang et al., 2025a) |  |  |  |  |  |  |  |  |
| Phi-Mamba-1.5B | 3 | 75.5 | 74.0 | 44.1 | 60.2 | 71.7 | 50.1 | - |
| LoLCATs-Phi-1.3B | 0.04 | 76.9 | 77.0 | 46.9 | 62.3 | 72.7 | - | - |
| Subquadratic: Distilled from Llama-3.2-1B (Bick et al., 2025; Zhang et al., 2025a) |  |  |  |  |  |  |  |  |
| Llamba-1B | 8 | 74.0* | 69.5* | 37.2 | 61.2 | 60.6 | 48.4 | - |
| LoLCATs-Llama-1B | 0.04 | 74.6 | 63.0 | 35.1 | 63.7 | 61.5 | 53.4 | 9.3 |
| LoLA-1B (ours) | 0.04 | 76.2 | 66.2 | 36.9 | 64.1 | 60.9 | 61.9 | 5.3 |

We provide additional notes for the results in Table [5](https://arxiv.org/html/2505.23666v2#A1.T5). Mamba and Mamba2 are popular architectures and have been trained many times with different datasets and hyperparameters. We report variations from two sources for robust results. In addition, LoLCATs demonstrated results on both Llama-3.2 1B and Phi-1.5. The model and code for reproducing LoLCATS-Phi-1.3B is not publicly available, so we could not produce Lambada scores. Similarly, we do not have LoLA results for this either. We were able to reproduce LoLCATs-Llama-1B, however, our achieved Winogrande accuracy was lower. We reported the score from the paper, $61.5\%$, over our reproduced $60.9\%$.

#### Cross-teacher comparison of distilled subquadratic models.

We gathered results from both Table [5](https://arxiv.org/html/2505.23666v2#A1.T5) and Table [3](https://arxiv.org/html/2505.23666v2#S4.T3) to compare language model performance relative to the teacher models. We average the performance across tasks and compute the relative average. This is calculated by dividing the model’s average by the teacher model’s average.

**Table 6: Comparison of distilled subquadratic models from different teacher models. We report the average accuracy across tasks when applicable (i.e., all scores are reported or available). We also report the relative accuracy, measured as the student average / the teacher average. Results were taken from various related works with the section header containing the sources.**
| Model | Tokens (B) | PI | AE | AC | HS | WG | LB | Avg. | Rel. Avg. |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Transformers |  |  |  |  |  |  |  |  |  |
| Phi-1.5-1.3B | 150 | 76.6 | 75.6 | 48.0 | 62.6 | 73.4 | 53.4 | 64.9 | - |
| Llama-3.2-1.3B | 9000 | 74.4 | 65.5 | 35.8 | 63.7 | 60.5 | 60.1 | 60.0 | - |
| Llama-3.1-8B | 15000 | 81.1 | 81.7 | 55.1 | 79.3 | 73.9 | 73.0 | 74.0 | - |
| Llama-3.1-8B-Instr. | 15000+ | 80.8 | 81.8 | 55.2 | 79.2 | 73.9 | - | 74.0 | - |
| Subquadratic: Distilled from Phi-1.5-1.3B |  |  |  |  |  |  |  |  |  |
| Phi-Mamba1.5B | 3 | 75.5 | 74.0 | 44.1 | 60.2 | 71.7 | 50.1 | 62.6 | $0.965\times$ |
| LoLCATs-1.3B | 0.04 | 76.9 | 77.0 | 46.9 | 62.3 | 72.7 | - | - | - |
| Subquadratic: Distilled from Llama-3.2-1.3B |  |  |  |  |  |  |  |  |  |
| Llamba-1.3B | 8 | 74.0* | 69.5* | 37.2 | 61.2 | 60.6 | 48.4 | 58.5 | $0.975\times$ |
| LoLCATs-1.3B | 0.04 | 74.6 | 63.0 | 35.1 | 63.7 | 61.5 | 53.4 | 58.6 | $0.977\times$ |
| LoLA-1.3B (ours) | 0.04 | 76.2 | 66.2 | 36.9 | 64.1 | 60.9 | 61.9 | 61.0 | $1.017\times$ |
| Subquadratic: Distilled from Llama-3.1-8B Instruct |  |  |  |  |  |  |  |  |  |
| Mamba2-Llama3-8B | 20 | 76.8 | 74.1 | 48.0 | 70.8 | 58.6 | 43.2 | 61.9 | $0.837\times$ |
| Subquadratic: Distilled from Llama-3.1-8B |  |  |  |  |  |  |  |  |  |
| Llamba-8B | 12 | 80.9 | 82.5 | 54.6 | 77.6 | 73.3 | 69.4 | 73.1 | $0.987\times$ |
| LoLCATs-8B | 0.04 | 81.0 | 82.4 | 54.4 | 79.1 | 73.6 | 67.6 | 73.0 | $0.987\times$ |
| LoLA-8B | 0.04 | 81.6 | 82.5 | 55.4 | 79.8 | 73.6 | 74.9 | 74.6 | $1.009\times$ |

In Table [6](https://arxiv.org/html/2505.23666v2#A1.T6), LoLA outperforms other distilled model approaches. We find that LoLCATs and Llamba perform similarly overall, with LoLCATs demonstrating better token efficiency. Overall, LoLA pushes the pareto front for high-performing and token-efficient models.

## Appendix B Efficiency Analysis

#### Cache Parameters.

Here, we analyze the efficiency of chunkwise LoLA. For various sliding window ($\eta$) and sparse cache ($\lambda$) sizes, we measure the peak VRAM cost and Time-to-First-Token for LoLA-8B with 4K long context on an Nvidia RTX 4090. Additionally, we show how different cache parameters lead to varying performance on RULER’s variable tracking task (Hsieh et al., 2024). There exists a trade-off between speed, memory footprint, and performance. There is no one-size-fits-all solution; we provide a short guide on navigating this tradeoff.

Figure: Figure 4: Measuring Time-to-First-Token for various sliding window and sparse cache sizes. This measurement is averaged across 100 trials and assumes data is already loaded into VRAM.
Refer to caption: x4.png

Figure [4](https://arxiv.org/html/2505.23666v2#A2.F4) illustrates that optimal throughput is achieved by maximizing the sliding window size that fits into VRAM and minimizing the sparse cache size. This reduces the number of chunks computed in sequential order, allowing for more intra-chunk parallelization.

Figure: Figure 5: Measuring Peak VRAM usage for various sliding window and sparse cache sizes. This measurement includes the base model weights, the data sequence, and online activations such as KV caches.
Refer to caption: x5.png

Figure [5](https://arxiv.org/html/2505.23666v2#A2.F5) suggests the total cache size ($\eta$ & $\lambda)$ needs to be reduced to lower VRAM cost. In our implementation, the data (4K long context sequence) already exists in VRAM, so this is included in the peak VRAM measurement. Furthermore, the TTFT does not include loading this data.

Figure: Figure 6: Cache Sizes vs. Variable Tracking performance at 4K context length from RULER (Hsieh et al., 2024)
Refer to caption: x6.png

Lastly, we show the variable tracking performance for various cache parameters in Figure [6](https://arxiv.org/html/2505.23666v2#A2.F6). Variable tracking requires understanding *all* of the context. The base subquadratic model—or even extended sliding window variants—cannot perform this task without a sparse cache. As a guideline to increase general long context performance, the sparse cache size should be maximized. This mitigates memory collisions, preserving linear attention’s hidden state on long sequences.

In summary, LoLA introduces a new trade-off for subquadratic models. Low VRAM and high performance can be achieved (maximize $\lambda$, minimize $\eta$), but the model will be slow. Low VRAM, fast models (minimize $\lambda$, moderate $\eta$) will not be able to perform well on long-context tasks (e.g. the base subquadratic model, LoLCATs). Finally, fast and high performing models will require a much larger memory footprint (maximize both $\lambda$ and $\eta$). This will extend the applicability of subquadratic models across various hardware platforms such as small inference chips or large training servers.

#### LoLA Cache vs. Transformers.

In Figure [7](https://arxiv.org/html/2505.23666v2#A2.F7), we compare LoLA’s bounded inference costs with vanilla, softmax attention. We compute the cache size as the total number of elements in all vectors and matrices stored for each attention head. For transformers, we store key and value vectors for each token in context, $t\,(d_{k}+d_{v})$. For LoLA, we add up the elements in each of the three memory systems: sliding window cache, sparse cache, and linear attention’s hidden state & normalizing state. Here, we provide both "small" and "large" cache size variants of LoLA for reference. In short context lengths, LoLA does not instantiate linear attention’s hidden state or measure the SRE. LoLA is equally efficient to softmax attention here.

Figure: Figure 7: Cache Size vs. Context Length for LoLA and vanilla transformers.
Refer to caption: x7.png

This figure illustrates why we are interested in subquadratic models to begin with. As context scales, we must linearly increase VRAM and compute per token. For example, we observed out-of-memory errors with Llama-3.1-8B at 4K context length in our experiments. LoLA, on the other hand, can be scaled with $\eta$ and $\lambda$ to maximize performance on specific hardware. The VRAM cost is agnostic to context length, meaning this model will always be able to fit.

## Appendix C Extended Long Context Tasks

To further extend our needle-in-a-haystack results from Table [1](https://arxiv.org/html/2505.23666v2#S4.T1), we provide more scores in Table [7](https://arxiv.org/html/2505.23666v2#A3.T7) across a greater variety of cache parameter combinations. For simplicity, we chose $\eta=\lambda$ and varied the total cache size, marked with "$+$". Each additional "$+$" doubles the total cache size (i.e baseline holds 64 tokens, + holds 128, ++ holds 256, etc.). Additionally, we provide longer sequences for S-NIAH-1 in Table [8](https://arxiv.org/html/2505.23666v2#A3.T8).

In these tasks, the “haystack” is synthetically constructed with various context lengths. The first task, S-NIAH-1, uses random sentences (e.g., “The grass is green.“) as the haystack, while S-NIAH-2 & 3 use essays. The needle—represented as a (word, number) pair—is placed in the haystack. At the end of the prompt, the model is tasked with returning the associated number with the special word. The first two tasks (S-NIAH-1 & 2) use a 7-digit number in the needle, and S-NIAH-3 uses a 32-digit UUID, requiring more tokens to represent the needle.

Table: Table 7: Measuring long context recall with Needle-in-a-Haystack tasks from the RULER benchmark. We report recall accuracy for each method across different context lengths (512, 1024, etc.) for each task.

**Table 8: Extended Results on S-NIAH-1. Reported as [Accuracy / Compression Rate]. We evaluated performance across 500 synthetic samples on all context lengths except 16K, which used 250.**
| Model | .5K | 1K | 2K | 4K | 8K | 16K |
| --- | --- | --- | --- | --- | --- | --- |
| LoLA-8B 4+ | 100% / 1$\times$ | 100% / 1$\times$ | 100% / 1.2$\times$ | 99.0% / 2.4$\times$ | 92.2% / 4.9$\times$ | 13.6% / 9.8$\times$ |

In general, extending the cache size can improve performance while still maintaining high compression rates over transformers. For even longer context lengths, LoLA’s cache can easily be scaled at inference to achieve the desired recall performance. This can be seen in Table [8](https://arxiv.org/html/2505.23666v2#A3.T8), where LoLA
performs well up to 8K context length, which is $8\times$ longer than the sequences seen during distillation.

For large haystacks, the accuracy of LoLCATs is roughly the proportion of context that the sliding window covers. For smaller haystacks, the performance is slightly higher than that ratio since fewer pairs are stored in the hidden state. This results in fewer collisions; though of course, this does not scale. Additionally, we observe that fewer collisions exist in essay-based haystacks (S-NIAH-2 & 3), likely as a result of being more similar to the distillation data. Lastly, we observe that harder needle-in-a-haystack tasks (e.g., S-NIAH-3 with 32-digit needles) may require more sparse caching. To further extend these results, we believe training with longer sequences should yield stronger performance.

## Appendix D Related Work

In this section, we position LoLA within the broader landscape of subquadratic models and efficient attention mechanisms.

#### Linear Attention and State Space Models (SSMs)

State Space Models (SSMs) have emerged as powerful architectures for efficient long-range sequence modeling, offering constant memory complexity irrespective of context length. Pioneered by methods like S4 (Gu et al., 2021), recent developments include various efficient architectures such as RetNet (Sun et al., 2023) and Mamba (Gu & Dao, 2024). Concurrently, original linear attention methods have explored efficient approximations of the softmax kernel (Katharopoulos et al., 2020; Choromanski et al., 2020; Qin et al.,; Peng et al., 2020; Zhang et al., 2024). Research between SSMs and linear attention has recently converged. Modern SSMs, such as DeltaNet (Schlag et al., 2021), Mamba2 (Dao & Gu, 2024) and Gated DeltaNet (Yang et al., 2024a), can be interpreted as linear attention models equipped with additional gating or delta-update mechanisms. Models like DeltaProduct (Siems et al., 2025) further generalize these linear updates through higher-rank modifications, while Titans (Behrouz et al., 2024) and TTT (Sun et al., 2024) extend the capacity for associative recall using richer hidden-state representations.

Test-time Regression (Wang et al., 2025) offers a unifying perspective for SSMs and linear attention. These sequence models perform online regression to fit hidden states to past context. Each state update can be interpreted as a gradient step in online SGD. This lens clarifies the roles of different mechanisms within these models. For example, forget gates in Mamba2 and Gated DeltaNet play an analogous role to weight decay. Similarly, momentum-based updates can be oversed in Titans (Behrouz et al., 2024).

#### Distilling transformers into subquadratic models.

The cost of pretraining LLMs is the primary obstacle in finding the successor of the transformer. To address this issue, recent approaches employ knowledge distillation, transferring the capabilities of pretrained transformers into subquadratic architectures (Bick et al., 2024; Zhang et al., 2025a; Mercat et al.,; Bick et al., 2025; Goldstein et al., 2025).
This significantly reduces training costs by recycling large pretrained models.

MOHAWK (Bick et al., 2024; 2025) demonstrated successful distillation of pretrained transformers into Mamba, maintaining competitive performance. Similarly, “Mamba in the Llama” (Wang et al., 2024) interleaves transformer and SSM blocks to retain transformer-level performance with significantly reduced inference costs. Though, this approach maintains an unbounded memory footprint due to the residual quadratic attention.

In contrast, LoLCATs use a simpler and cheaper distillation approach by using a student architecture that is more similar to the transformer. The combination of linear attention and sliding window attention significantly reduces the distillation complexity, requiring significantly fewer training tokens. LoLA directly builds on LoLCATs, leveraging its efficient distillation approach while introducing sparse caching to substantially enhance associative recall without extensive retraining.

#### Sparse attention methods.

Sparse attention methods present another orthogonal approach to reducing Transformer complexity by limiting the set of attended tokens (Nawrot et al., 2025). Methods such as Longformer (Beltagy et al., 2020) and BigBird (Zaheer et al., 2020) adopt fixed sparse patterns that incorporate sliding windows and selective global attention, efficiently capturing both local and sparse global contexts. Recent dynamic sparsification approaches, including Loki (Singhania et al., 2024) and Native Sparse Attention (NSA)(Yuan et al., 2025), employ data-dependent strategies, selectively attending to the most relevant tokens based on learned or projected keys. Native Sparse Attention, specifically, combines sparse attention with latent attention mechanisms(Liu et al., 2024), effectively approximating attention via low-rank and sparse structures.

We believe sparse attention can be complementary to linear attention. With LoLA, we encourage the use of hybrid attention techniques *within* the same attention head. This allows important tokens to leverage more computation when needed. This work contrasts the use of interleaving soft attention blocks with linear attention blocks (Ren et al., 2025; Glorioso et al., 2024) which allocates the compute costs equally between all tokens.

## Appendix E Linear Attention is a Bad Low-Rank Approximation

In this section, we analyze why linear attention struggles to closely approximate softmax attention, specifically highlighting difficulties in approximating the exponential dot product kernel. We start by defining the exponential kernel’s Gram matrix $\mathbf{G}$ as $\mathbf{G}_{i,j}=\exp(\bm{x}_{i}^{\top}\bm{x}_{j})$
for inputs $\bm{x}_{i},\bm{x}_{j}\in\mathbb{R}^{d}$. This kernel implicitly corresponds to an inner product in a potentially infinite-dimensional Hilbert space $\mathcal{H}$ via a feature map $\phi_{\text{exp}}:\mathbb{R}^{d}\rightarrow\mathcal{H}$, such that

$$ $\exp(\bm{x}_{i}^{\top}\bm{x}_{j})=\phi_{\text{exp}}(\bm{x}_{i})^{\top}\phi_{\text{exp}}(\bm{x}_{j}).$ (16) $$

Since explicitly working in an infinite-dimensional space $\mathcal{H}$ is infeasible, linear attention methods approximate this kernel using a finite-dimensional feature map $\phi:\mathbb{R}^{d}\rightarrow\mathbb{R}^{D}$. Consequently, linear attention approximates the Gram matrix as

$$ $\mathbf{G}_{i,j}\approx\widehat{\mathbf{G}}_{i,j}=\phi(\bm{x}_{i})^{\top}\phi(\bm{x}_{j}),$ (17) $$

which has a maximum rank of $D$. Ideally, $\widehat{\mathbf{G}}$ would closely approximate $\mathbf{G}$, minimizing the squared Frobenius norm error. However, linear attention’s approximation error is fundamentally lower-bounded by the truncated singular value decomposition (SVD) of $\mathbf{G}$ (Eckart & Young, 1936). Specifically, for the SVD decomposition $\mathbf{G}=\mathbf{U}\mathbf{\Sigma}\mathbf{V}^{\top}$ with singular values $\sigma_{i}$, we have:

$$ $\|\mathbf{G}-\widehat{\mathbf{G}}\|_{F}^{2}\geq\|\mathbf{G}-\mathbf{U}_{D}\mathbf{\Sigma}_{D}\mathbf{V}_{D}^{\top}\|_{F}^{2}=\sum_{i=D+1}^{\text{rank}(\mathbf{G})}\sigma_{i}^{2},$ (18) $$

where $\mathbf{U}_{D}\mathbf{\Sigma}_{D}\mathbf{V}_{D}$ is the rank $D$ truncated SVD approximation of $\mathbf{G}$.

While the truncated SVD provides the optimal low-rank approximation, it requires the entire Gram matrix to be computed and stored, making it impractical for linear attention which demands computationally efficient, online feature mappings.

To empirically demonstrate the severity of this approximation challenge, we construct the Gram matrix under different input distributions and analyze its singular values. In our simulations, we vary both $n$ (the number of independently sampled input vectors) and $d$ (input vector dimensionality) and observe how they affect singular value distributions. Specifically, we draw inputs from a scaled Gaussian distribution $\bm{x}_{i}\sim\mathcal{N}(0,d^{-1/4})$ to mimic typical transformer scaling of dot products by $\sqrt{d}$.

Figure: (a)
Refer to caption: x8.png

Figure: (a)
Refer to caption: x14.png

Figures [8](https://arxiv.org/html/2505.23666v2#A5.F8) and [9](https://arxiv.org/html/2505.23666v2#A5.F9) show that applying an exponential operation to the query-key products significantly increases the rank and complexity of the resulting Gram matrix. The singular values and approximation error increase with the number of unique input vectors $n$ (see Figure [8](https://arxiv.org/html/2505.23666v2#A5.F8)) and the input dimension $d$ (see Figure [9](https://arxiv.org/html/2505.23666v2#A5.F9)) . Practically, in transformer architectures, head dimensions are typically modest ($d=64$ for Llama-3.2 1B and $d=128$ for Llama-3.1 8B). Additionally, linear attention approaches typically select feature dimensions $D$ around $2d$ (Zhang et al., 2024; 2025a) which can be problematic without additional sparse attention or gating mechanisms.

These experiments underscore the inherent limitation of linear attention as a softmax replacement. For an arbitarily large vocabulary size, a high dimensional hidden state is needed to truly mimic softmax. We argue that future research should exploit the inherent strengths of linear attention when it makes sense to (e.g., applying linear attention on easier-to-remember tokens), rather than attempting to replicate softmax attention in all situations.

## Appendix F Scoring Ablation Extension

In section [4.3](https://arxiv.org/html/2505.23666v2#S4.SS3.SSS0.Px1), we ablated different scoring approaches for the sparse cache. Here, we describe exactly how each alternative score is computed. As a reminder, LoLA is motivated by the occurence of memory collisions in the hidden state. LoLA explicitly attempts to maintain self-recall for stored key-value pairs. An alternative perspective could be aiming for the best softmax approximation. Similar to previous kernel work (Choromanski et al., 2020), this aims to minimize the attention weight error

$$ $\sum_{i}^{N}\sum_{j}^{N}\left(\exp(\bm{q}_{i}^{\top}\bm{k}_{j})-\phi(\bm{q}_{i})^{\top}\phi(\bm{k}_{j})\right)^{2}.$ (19) $$

A natural scoring method for this objective would be to keep the keys with the highest attention weight error. As a proxy for this approach, each key’s error can be summed over all queries it sees in the sliding window with

$$ $\text{Score}(\bm{k}_{i})=\sum_{t=i}^{i+\eta-1}\left(\exp(\bm{q}_{t}^{\top}\bm{k}_{i})-\phi(\bm{q}_{t})^{\top}\phi(\bm{k}_{i})\right)^{2}.$ (20) $$

Alternatively, we can use absolute error over mean squared error instead.

From traditional sparse attention literature (Zhang et al., 2023; Dong et al., 2024), keys that are highly attended to may be important. We compute this as

$$ $\text{Score}(\bm{k}_{i})=\sum_{t=i}^{i+\eta-1}\exp(\bm{q}_{t}^{\top}\bm{k}_{i}).$ (21) $$

From a third perspective, keys that are over-represented by linear attention’s query-key interactions may seem important to cache. We compute these as

$$ $\text{Score}(\bm{k}_{i})=\sum_{t=i}^{i+\eta-1}\frac{\phi(\bm{q}_{t})^{\top}\phi(\bm{k}_{i})}{\exp(\bm{q}_{t}^{\top}\bm{k}_{i})}.$ (22) $$

Lastly, we compare these methods against a larger sliding window.

## Appendix G Extended Memory Collision Visualization

Figure: Figure 10: Visualizing the relative SRE for stored KV pairs.
Refer to caption: x20.png

In Section [4.3](https://arxiv.org/html/2505.23666v2#S4.SS3), we visualized how memory collisions can occur in practice. This was computed by measuring the SRE for all stored KV pairs. We found that the first few stored KV pairs do not achieve a large error when added to the hidden state. However, these quickly become corrupted after hidden state updates. This phenomenon was difficult to see in Figure [3](https://arxiv.org/html/2505.23666v2#S4.F3), so we provide an additional visualization with Figure [10](https://arxiv.org/html/2505.23666v2#A7.F10). Specifically, we measure the SRE of each KV pair, relative to the error of when that pair was added. For row $i$ column $j$ in the plot, the relative error is computed as

$$ $\left\|\frac{\phi(\bm{k}_{j})^{\top}\mathbf{H}_{i}}{\phi(\bm{k}_{j})^{\top}\bm{s}_{i}}-\bm{v}_{j}\right\|_{2}-\left\|\frac{\phi(\bm{k}_{j})^{\top}\mathbf{H}_{t}}{\phi(\bm{k}_{j})^{\top}\bm{s}_{t}}-\bm{v}_{j}\right\|_{2}.$ (23) $$

where $t$ is the time pair $j$ was added to the hidden state. Thus, we have $j\leq t\leq i$.

Here, we see the first few KV pairs have high relative errors for pure linear attention. These pairs observe small SREs at early time steps, but achieve much higher SREs later. On the other hand, sparse caching actively mitigates SREs, improving associative recall for stored KV pairs.

## Appendix H Algorithm Psuedo-code

We provide PyTorch-like pseudo-code for the cache during decoding or generation. In this example, we update the cache every iteration (rather than chunkwise inference) for simplicity. Lastly, this pseudo-code does not contain any optimization tricks for ease of understanding.

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: Can author provides the comparison on the decoding latency their LoLA and original LA?
soundness: 3
strengths: 1. The paper empirically demonstrates an improvement on tasks where the base model completely fails
2. The propose of a new sparse global cache is novel
summary: The paper introduces LoLA, a training-free inference strategy that enhances linear attention models by mitigating their poor memory capacity. It distributes key-value pairs into three systems: a local sliding window, the low-rank recurrent state, and a new sparse global cache that stores "difficult-to-memorize" pairs identified by a self-recall error metric, drastically improving associative recall.
weaknesses: 1. The method adds a new computational cost not present in standard linear attention, specifically for scoring and managing the sparse cache. The paper says its overhead is O(λd) which itself needs to be larger for more complex, long-context tasks.
2. LoLA moves away from the simplicity of linear attention by requiring three distinct memory systems that must be managed. 
3.  This caching strategy cannot fully compensate for the knowledge lost during the base model's efficient distillation from attention

---

confidence: 5
contribution: 2
initial_rating: 6
presentation: 3
questions: 1. The paper's goal is to mitigate memory collisions, which is the same goal as advanced recurrent models like Gated DeltaNet. How do the authors view LoLA's "sparse caching" approach in comparison to Gated DeltaNet's "gating + delta rule" approach? 

2. The TTFT analysis in Figure 4 shows a clear slowdown as $\lambda$ increases. Could the authors provide a more direct comparison of this overhead? 

3. The SRE metric relies on the $\phi$ map learned by the LoLCATs base model. Have the authors tested whether LoLA is effective when applied to other linear attention models that were not trained with LoLCATs' specific distillation process, such as a model pretrained from scratch?
soundness: 3
strengths: 1. The paper's core idea is highly original. Instead of using query similarity or softmax scores for sparse attention, it introduces the Self Recall Error (SRE) . This query agnostic metric provides a principled, data driven way to determine which tokens are "difficult to memorize" for the linear state and should be cached in full rank. This is a clever and novel approach to mitigating memory collisions.

2. The claims are supported by exceptionally strong and well targeted experiments. The method provides a near perfect fix for associative recall on Needle in a Haystack (NIAH) tasks, boosting accuracy from a failing 0.6% to 97.4%. Furthermore, the ablation study in Table 4 is critical, proving that LoLA's SRE metric (99.0% accuracy) is vastly superior to intuitive alternatives like extending the sliding window (29.4%) or using other sparse attention metrics (e.g., 11.4%, 20.0%).

3. The work is highly significant because it is a training free inference strategy. This allows it to be applied directly to existing, pretrained hybrid models like LoLCATs to immediately unlock their long context recall capabilities.
summary: This paper introduces LoLA, a training-free inference strategy designed to solve the poor memory capacity of linear attention models. Standard linear attention suffers from "memory collisions," where new information overwrites past associations in its fixed-size hidden state. LoLA addresses this by augmenting hybrid linear attention models with a novel three-part memory system: (i) a local sliding window cache for recent key-value (KV) pairs, (ii) a sparse, global cache for difficult-to-memorize pairs, and (iii) the standard recurrent hidden state for all other generic pairs. The core of LoLA is a new importance metric called Self-Recall Error (SRE), which is used to manage the sparse cache. SRE measures how well the current linear hidden state can already retrieve a KV pair's correct value. Pairs that are difficult to recall (high SRE) are kept in the full-rank sparse cache to prevent them from being corrupted, while easy-to-recall pairs (low SRE) are compressed into the hidden state. This method is highly effective, improving a base model's accuracy on pass-key retrieval tasks from 0.6% to 97.4% and achieving strong performance on general commonsense reasoning benchmarks.
weaknesses: 1. The paper states the scoring introduces a "small overhead compute cost". However, the proposed algorithm re-scores all $\lambda$ elements in the sparse cache plus the new candidate token(s) at every generation step. This overhead is non trivial, especially when $\lambda$ is large. The Time to First Token (TTFT) in Figure 4 confirms this. For a 64 token window, TTFT increases from 0.99s ($\lambda=0$) to 1.46s ($\lambda=512$). This is a roughly 47% slowdown. This trade off is not sufficiently addressed in the main paper.

2. The paper positions LoLA as a solution to memory collisions in linear attention but fails to compare its approach to other architectures designed to solve the exact same problem. Specifically, Gated DeltaNet, a highly relevant model, combines gating and a delta rule to precisely manage its recurrent state, mitigate collisions, and improve associative recall. While Gated DeltaNet is a different architecture (requiring training) and LoLA is an inference strategy, they are both direct solutions to memory collisions. The paper should have discussed, if not benchmarked against, such advanced recurrent models.

3. The method is described as "training free". While true that LoLA itself is an inference strategy, its SRE metric (Equation 10) is entirely dependent on the feature map $\phi$, which was learned during the distillation of the LoLCATs base model. It is unclear if LoLA would be as effective if applied to a linear attention model trained from scratch, or one using a different $\phi$.

---

confidence: 3
contribution: 2
initial_rating: 4
presentation: 2
questions: * **Grammar in Fig. caption.** “**LoLA stores past KV pairs in three forms memory**” → *three forms of memory*. 
* **Informal wording.** “**anyways**” → *anyway*. (Figure text) 
* **Spelling.** “**pruple**” → *purple*. (Figure text) 
* **Apostrophe.** “**Its more efficient**” → *It’s more efficient*. (Figure text) 
* **Spelling.** “**occurence**” → *occurrence*. (Appendix F text) 
* **Consistency.** Use either *commonsense* or *common sense* uniformly; currently the abstract uses “commonsense” while Table 3 caption uses “common sense.”  
* **Section label consistency.** TOC line says “**Algorithm Psuedo-code**” (again typo); ensure capitalization and spelling consistent with the section header. 
* Reframe “training-free” to “inference-only augmentation **for subquadratic models**” and clearly delineate the *pre-training/distillation* of the base from the *LoLA* mechanism. Provide a small *plug-compatibility* study across multiple linear-attention bases to evidence generality. 
* Make the cache-size factor claims auditable: define the exact KV token counts behind “× smaller” in each table/setting, and add a per-token throughput plot (not only TTFT) with and without sparse rescoring, end-to-end (I/O included).
soundness: 3
strengths: The SRE criterion is intuitive and easy to compute given $\Phi(k)$, H, s; the paper supplies pseudo-code and a useful efficiency study (TTFT and VRAM) versus sliding-window size η and sparse-cache size $\lambda$, which practitioners can adopt to tune deployments.
summary: The paper proposes **LoLA**, a *training-free* inference-time augmentation for hybrid sliding-window + linear attention models. LoLA keeps three memory systems: (i) a local sliding-window cache, (ii) a sparse global cache, and (iii) a recurrent linear-attention hidden state. It scores KV pairs using a **self-recall error (SRE)** how poorly linear attention alone can reconstruct a value from its key and promotes “hard” pairs into the sparse cache while easier ones flow into the recurrent state. On synthetic long-context recall (RULER/S-NIAH) and some zero-shot commonsense tasks, LoLA improves over the authors’ distilled LoLCATs baselines under fixed cache budgets.
weaknesses: 1. **Positioning / novelty is narrow and tied to a special base model.** Although billed as “training-free,” LoLA **assumes** a specific *subquadratic* base (sliding-window + linear attention) obtained via distillation/LoRA (40M tokens) before LoLA can be used. It is therefore not a drop-in for standard Transformers and the headline “training-free” risks misinterpretation. Please clarify scope and re-title accordingly; also separate the cost/benefit of distillation from LoLA’s cache policy.  

2. **Claims vs. baselines and calculation clarity.** The abstract’s “**4.6$\times$ smaller cache than Llama-3.1-8B on 4K**” is not explained precisely; Table 1 lists the Transformer as full-attention (effectively a 4K-token KV cache at 4K), while LoLA uses η+λ (e.g., 256+256). Please specify exactly which Llama cache size is used for that factor and how it’s computed; as written, it’s hard to reconcile. More broadly, comparisons mix distilled subquadratic models and different training budgets; ensure apples-to-apples controls (same teacher, tokens, and evaluation protocol).   

3. **Evidence focuses on synthetic recall; limited breadth on real tasks and ablations.** The large wins are concentrated on RULER/S-NIAH; while there are commonsense results, the absolute gains vs. LoLCATs are modest in places and still trail Llama on some tasks. Scoring ablations for SRE are shown but only on S-NIAH @ 0.5K; please broaden ablations (e.g., alternative importance metrics, cache-update cadence, chunk size sensitivity) on harder and real-world long-context tasks (multi-needle, multi-hop QA with citations).

---

confidence: 4
contribution: 2
initial_rating: 4
presentation: 2
questions: N/A
soundness: 2
strengths: The proposed method is effective at NIAH recall.
summary: The authors proposes LoLA, a training-free enhancement to linear attention that improves long-term associative memory. It introduces a three-part memory system: (i) a local sliding window for recent key-value pairs, (ii) a sparse global cache for hard-to-recall items, and (iii) a low-rank recurrent hidden state for generic context, which is guided by a self-recall error metric to detect and mitigate memory collisions.
weaknesses: 1. Limited evaluation: The current evaluation focuses on synthetic task and short-context NLP benchmarks. Could authors also evaluate their method on at least one long-context NLP benchmark, e.g. LongBench [1]?
2. Work related to this paper to be discussed:
   - The authors claim that current hybrid approaches only focus on local information with sliding window attention, which is not true. Could authors discuss the recent hybrid architectures with sparse attention, e.g. B'MOJO [2], Se-Attn [3], HAX [4]?
   - For the minimization of Self-Recall Error (Eq. 10) in the paper, could authors discuss the correlation of their method with recent test-time training method, e.g. GatedDeltaNet [5], TTT [6], Titans [7], Lattice [8]?

[1]. LongBench: A Bilingual, Multitask Benchmark for Long Context Understanding. ACL2024.

[2]. B'MOJO: Hybrid State Space Realizations of Foundation Models with Eidetic and Fading Memory. NeurIPS 2024.

[3]. Expansion Span: Combining Fading Memory and Retrieval in Hybrid State Space Models. NeuS 2025.

[4]. Overcoming Long-Context Limitations of State-Space Models via Context-Dependent Sparse Attention. NeurIPS 2025.

[5]. Gated Delta Networks: Improving Mamba2 with Delta Rule. NeurIPS 2024.

[6]. Learning to (Learn at Test Time): RNNs with Expressive Hidden States. ICML 2025.

[7]. Titans: Learning to Memorize at Test Time. https://arxiv.org/abs/2501.00663

[8]. Lattice: Learning to Efficiently Compress the Memory. CoLM 2025.


Decision: Reject