Linear Attention & Kernel Methods

Learning Objectives
Notation
Core Intuition
From Softmax to Kernels
The Linear Attention Trick
Feature Map Choices
Causal Linear Attention as RNN
RetNet: Retention Mechanism
RWKV: Linear Attention Language Model
Expressiveness Analysis
Common Pitfalls
Summary
Exercises

Learning Objectives

Derive linear attention from the kernel view of softmax.
Prove the $O(Td^2)$ complexity of linear attention.
Show the equivalence between causal linear attention and an RNN.
Explain why linear attention struggles with precise retrieval.
Compare RetNet and RWKV architectures.

Notation

$\phi: \mathbb{R}^d \to \mathbb{R}^D$ — feature map
$\mathbf{S} \in \mathbb{R}^{D \times d_v}$ — accumulated state matrix
$\mathbf{z} \in \mathbb{R}^D$ — normalizer state

Core Intuition

Standard attention computes pairwise similarities between all $T^2$ pairs, then normalizes. But what if the similarity function $\kappa(\mathbf{q}, \mathbf{k}) = \phi(\mathbf{q})^T\phi(\mathbf{k})$ can be decomposed into a product of features? Then we can change the order of computation: instead of "for each query, compare to all keys," we compute "aggregate all key-value pairs into a summary, then query the summary." This changes $O(T^2)$ to $O(T)$ .

Linear Attention

Feat d

Exact softmaxLinear approx

Explore: Linear attention replaces softmax with feature maps φ, enabling O(N) complexity. Higher feature dim improves approximation quality.

From Softmax to Kernels

Standard softmax attention for a single query:

\text{Attn}(\mathbf{q}) = \frac{\sum_{j=1}^T \exp(\mathbf{q}^T\mathbf{k}_j/\sqrt{d_k}) \, \mathbf{v}_j}{\sum_{j=1}^T \exp(\mathbf{q}^T\mathbf{k}_j/\sqrt{d_k})}. \tag{1}

View the exponential as a kernel: $\kappa(\mathbf{q}, \mathbf{k}) = \exp(\mathbf{q}^T\mathbf{k}/\sqrt{d_k})$ .

Kernel trick: If $\kappa(\mathbf{q}, \mathbf{k}) = \phi(\mathbf{q})^T\phi(\mathbf{k})$ for some feature map $\phi$ , then:

\text{Attn}(\mathbf{q}) = \frac{\phi(\mathbf{q})^T \sum_j \phi(\mathbf{k}_j)\mathbf{v}_j^T}{\phi(\mathbf{q})^T \sum_j \phi(\mathbf{k}_j)}. \tag{2}

The Linear Attention Trick

Key insight: The sums $\sum_j \phi(\mathbf{k}_j)\mathbf{v}_j^T$ and $\sum_j \phi(\mathbf{k}_j)$ don't depend on the query — compute them once.

Define:

\mathbf{S} = \sum_{j=1}^T \phi(\mathbf{k}_j)\mathbf{v}_j^T \in \mathbb{R}^{D \times d_v}, \quad \mathbf{z} = \sum_{j=1}^T \phi(\mathbf{k}_j) \in \mathbb{R}^D. \tag{3}

Then for any query:

\text{Attn}(\mathbf{q}_i) = \frac{\phi(\mathbf{q}_i)^T\mathbf{S}}{\phi(\mathbf{q}_i)^T\mathbf{z}}. \tag{4}

Complexity:

Build $\mathbf{S}$ : $O(T \cdot D \cdot d_v)$ .
Query each of $T$ positions: $O(T \cdot D \cdot d_v)$ .
Total: $O(T D d_v)$ . If $D = O(d)$ : $O(Td^2)$ — linear in $T$ .

Feature Map Choices

The softmax kernel $\exp(\mathbf{q}^T\mathbf{k})$ has an infinite-dimensional exact feature map. Practical choices approximate it:

1. ELU + 1 (Katharopoulos et al., 2020):

\phi(\mathbf{x}) = \text{elu}(\mathbf{x}) + 1 = \begin{cases}x + 1 & x \geq 0 \\ e^x & x < 0\end{cases}. \tag{5}

Simple, non-negative, but doesn't approximate softmax well.

2. Random Fourier Features (Performers):

\phi(\mathbf{x}) = \frac{1}{\sqrt{D}}[\cos(\boldsymbol{\omega}_1^T\mathbf{x}), \sin(\boldsymbol{\omega}_1^T\mathbf{x}), \ldots], \quad \boldsymbol{\omega}_i \sim \mathcal{N}(0, \mathbf{I}). \tag{6}

Unbiased estimator of the Gaussian kernel: $\mathbb{E}[\phi(\mathbf{q})^T\phi(\mathbf{k})] = \exp(-\|\mathbf{q}-\mathbf{k}\|^2/2)$ .

3. Positive random features (FAVOR+): Ensure non-negativity:

\phi(\mathbf{x}) = \frac{\exp(-\|\mathbf{x}\|^2/2)}{\sqrt{D}}[\exp(\boldsymbol{\omega}_1^T\mathbf{x}), \ldots, \exp(\boldsymbol{\omega}_D^T\mathbf{x})]. \tag{7}

Unbiased estimator of softmax kernel.

Causal Linear Attention as RNN

For autoregressive (causal) models, the state updates incrementally:

\mathbf{S}_t = \mathbf{S}_{t-1} + \phi(\mathbf{k}_t)\mathbf{v}_t^T, \quad \mathbf{z}_t = \mathbf{z}_{t-1} + \phi(\mathbf{k}_t), \tag{8}

\mathbf{o}_t = \frac{\phi(\mathbf{q}_t)^T\mathbf{S}_t}{\phi(\mathbf{q}_t)^T\mathbf{z}_t}. \tag{9}

This is an RNN with state $(\mathbf{S}_t, \mathbf{z}_t)$ :

Fixed-size state: $D \times d_v + D$ regardless of sequence length.
Constant-time per-step update: $O(Dd_v)$ .
No KV-cache needed.

Implication: Linear attention converts a transformer into an RNN at inference time (linear generation), while allowing parallel training (using the cumulative sum formulation).

RetNet: Retention Mechanism

RetNet adds exponential decay to linear attention:

\mathbf{S}_t = \gamma \mathbf{S}_{t-1} + \phi(\mathbf{k}_t)\mathbf{v}_t^T, \tag{10}

where $\gamma \in (0, 1)$ is a decay factor.

Effect: Old information is gradually forgotten (exponential decay). This acts as a learned forgetting mechanism and helps with:

Reducing the influence of distant irrelevant tokens.
Providing implicit positional encoding (closer tokens have higher weight).

Three computation modes:

Parallel (training): $O(Td^2)$ using matrix form.
Recurrent (inference): $O(d^2)$ per step using state update.
Chunk-wise (hybrid): Process chunks in parallel, connect chunks recurrently.

RWKV: Linear Attention Language Model

RWKV combines ideas from RNNs and transformers:

Time-mixing (attention analog):

\text{wkv}_t = \frac{\sum_{j=1}^{t-1}e^{-(t-1-j)w + k_j}v_j + e^{u+k_t}v_t}{\sum_{j=1}^{t-1}e^{-(t-1-j)w + k_j} + e^{u+k_t}}, \tag{11}

where $w$ is a learned decay and $u$ is a bonus for the current token.

Properties:

Linear complexity: $O(Td)$ per layer.
Can be computed as an RNN (constant memory per step).
Competitive with transformers up to ~14B parameters.

Expressiveness Analysis

Theorem (Limitation). Linear attention cannot implement "sharp" attention (near-one-hot) patterns.

Argument: In standard attention, softmax can produce distributions arbitrarily close to one-hot (by scaling up logits). In linear attention, the output is:

\mathbf{o} = \frac{\phi(\mathbf{q})^T\mathbf{S}}{\phi(\mathbf{q})^T\mathbf{z}}. \tag{12}

This is always a smooth average of values weighted by $\phi(\mathbf{q})^T\phi(\mathbf{k}_j)$ — cannot perfectly select a single key.

Consequence: Tasks requiring precise retrieval from memory (e.g., copying a specific token, looking up a fact at a specific position) are harder for linear attention.

Common Pitfalls

Pitfall 1. Assuming linear attention always beats quadratic for long sequences. At small $T$ , the overhead of feature computation and the larger constant in $O(Td^2)$ (vs $O(T^2d)$ with $d < T$ ) can make linear attention slower.

Pitfall 2. Using a feature map that produces negative values. Negative attention weights are semantically meaningless and cause instability in the normalization.

Pitfall 3. Expecting linear attention to handle in-context learning well. Tasks requiring precise token recall (common in ICL) are fundamentally limited by the fixed-size state.

Summary

Linear attention replaces $O(T^2)$ with $O(Td^2)$ using kernel decomposition.
Causal linear attention = RNN with state $\mathbf{S} \in \mathbb{R}^{D \times d_v}$ .
No KV-cache needed; constant memory for generation.
Feature maps: ELU+1, random Fourier features, FAVOR+.
RetNet adds exponential decay; RWKV adds position-dependent decay.
Tradeoff: linear complexity at the cost of reduced expressiveness (no sharp retrieval).

Exercises

Exercise 1. Derive the parallel form of causal linear attention using cumulative sums.

Exercise 2. Prove that with $\phi(\mathbf{x}) = \mathbf{x}$ (identity), linear attention reduces to a linear layer.

Exercise 3. Compute the state size $|\mathbf{S}|$ for a model with $d_k = d_v = 128$ and compare to the KV-cache size at $T = 4096$ .

Exercise 4. Show that RetNet with $\gamma = 1$ reduces to standard linear attention.

Exercise 5. Design an experiment to demonstrate the retrieval limitation of linear attention vs softmax attention.