Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Linear Attention › The Linear Attention Formula

Michael BrenndoerferUpdated June 28, 20256 min read

Language AI HandbookAll chapters

Part 17 of 67Efficient Attention

Chapter 123Linear Attention › The Linear Attention Formula

Reading section 3 of 15 · The Linear Attention Formula · 6 minSection 3 of 15 · 6 min

Section 3 of 15 · 6 min read

Linear Attention › The Linear Attention FormulaOne section · Sections in this chapter3 / 15 ↓

Sections in this chapter15

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Having established the mathematical foundation, we can now synthesize everything into the complete linear attention formula. The journey from standard attention to linear attention follows a clear logical path: we identified the softmax as the source of the quadratic bottleneck, replaced it with a decomposable kernel function, and discovered that the resulting expression can be rearranged to avoid computing the full attention matrix. Now let's see how these insights combine into a practical algorithm.

From Element-wise to Matrix FormLink Copied

We derived earlier that each position's output can be written as:

\text{output}_i = \frac{\phi(q_i)^T S}{\phi(q_i)^T Z}

To express this for all positions simultaneously, we stack the feature-mapped queries into a matrix and apply the same operations in parallel. For a sequence of $n$ tokens with queries $Q$ , keys $K$ , and values $V$ , the complete formula is:

\text{LinearAttention}(Q, K, V) = \frac{\phi(Q)(\phi(K)^T V)}{\phi(Q) \phi(K)^T \mathbf{1}}

Let's unpack each component to understand how the pieces fit together:

$\phi(Q) \in \mathbb{R}^{n \times D}$ : The feature map $\phi$ applied row-wise to the query matrix. Each of the $n$ query vectors is transformed from $d_k$ dimensions to $D$ dimensions.
$\phi(K) \in \mathbb{R}^{n \times D}$ : Similarly, the feature map applied to each key vector. After transformation, we transpose this to get $\phi(K)^T \in \mathbb{R}^{D \times n}$ .
$\phi(K)^T V$ : This is the key-value aggregate $S$ from our derivation. The matrix product has shape $(D \times n) \cdot (n \times d_v) = (D \times d_v)$ . Notice that $n$ disappears from the result shape, and this is precisely why linear attention avoids the quadratic bottleneck.
$\phi(Q)(\phi(K)^T V)$ : Multiplying the queries by the aggregate produces the numerator. Shape: $(n \times D) \cdot (D \times d_v) = (n \times d_v)$ .
$\phi(K)^T \mathbf{1}$ : Summing the transformed keys (equivalent to $Z$ in our derivation). The vector of ones $\mathbf{1} \in \mathbb{R}^n$ sums across all $n$ key positions, producing a $D$ -dimensional vector.
The division is performed element-wise per row, normalizing each output position by its corresponding scalar denominator.

The Order of Operations MattersLink Copied

The formula highlights the critical insight: the order in which we perform matrix multiplications determines the complexity. Consider the numerator $\phi(Q)(\phi(K)^T V)$ :

If we compute $\phi(Q)\phi(K)^T$ first (as in standard attention), we get an $n \times n$ matrix. Multiplying by $V$ then costs $O(n^2 d_v)$ .
If we compute $\phi(K)^T V$ first (the linear attention approach), we get a $D \times d_v$ matrix. Multiplying $\phi(Q)$ by this costs only $O(n D d_v)$ .

The second approach avoids ever creating the $n \times n$ matrix. This is matrix associativity in action: $(AB)C = A(BC)$ , but the intermediate sizes differ dramatically.

Computational comparison between standard and linear attention showing how reordering matrix operations eliminates the quadratic intermediate matrix.
Operation	Standard Attention	Linear Attention
Core computation	$(QK^T)V$	$Q(K^TV)$
Intermediate matrix	$n \times n$	$D \times d_v$
Complexity	$O(n^2 d)$	$O(nd^2)$ or $O(ndD)$

When the feature dimension $D$ is proportional to the model dimension $d$ (and $d \ll n$ for long sequences), linear attention provides substantial savings. For a 10,000-token sequence with $d = 64$ , standard attention requires a 100-million-element intermediate matrix, while linear attention needs only 4,096 elements.

Implementation: Translating Math to CodeLink Copied

Let's implement linear attention step by step, mapping each line of code to the mathematical operations we've discussed.

In[3]:

Code

import numpy as np


def linear_attention(Q, K, V, feature_map=None, eps=1e-6):
    """
    Compute linear attention using the associative property.

    Args:
        Q: Query matrix of shape (n, d_k)
        K: Key matrix of shape (n, d_k)
        V: Value matrix of shape (n, d_v)
        feature_map: Function to apply to Q and K. If None, uses ELU + 1.
        eps: Small constant for numerical stability

    Returns:
        Output of shape (n, d_v)
    """
    if feature_map is None:
        # ELU + 1 ensures non-negative features (commonly used)
        feature_map = lambda x: np.maximum(x, 0) + np.exp(np.minimum(x, 0))

    # Apply feature map to queries and keys
    Q_prime = feature_map(Q)  # (n, d_k)
    K_prime = feature_map(K)  # (n, d_k)

    # Compute the key-value aggregate: K'^T @ V
    # Shape: (d_k, d_v)
    KV = K_prime.T @ V

    # Compute the key aggregate for normalization: sum of K' along sequence
    # Shape: (d_k,)
    K_sum = K_prime.sum(axis=0)

    # Compute outputs using the efficient formulation
    # Numerator: Q' @ KV, shape (n, d_v)
    numerator = Q_prime @ KV

    # Denominator: Q' @ K_sum, shape (n,)
    denominator = Q_prime @ K_sum + eps

    # Normalize each position
    output = numerator / denominator[:, np.newaxis]

    return output

The implementation follows our derivation exactly:

Apply the feature map: Transform $Q$ and $K$ to get $\phi(Q)$ and $\phi(K)$ . We use ELU+1 as a default because it ensures all features are positive, which is necessary for the normalization to produce valid weights.
Compute the key-value aggregate: K_prime.T @ V computes $\phi(K)^T V$ , the $D \times d_v$ matrix that summarizes how each feature dimension relates to the values. This is computed once and reused for all queries.
Compute the normalization aggregate: K_prime.sum(axis=0) computes $\sum_j \phi(k_j)$ , equivalent to $\phi(K)^T \mathbf{1}$ . This is the denominator's aggregate.
Query the aggregates: Each query retrieves its output by multiplying against the precomputed aggregates. The numerator Q_prime @ KV and denominator Q_prime @ K_sum are both $O(nDd)$ operations.
Normalize: Element-wise division produces the final attention output.

Seeing the Efficiency GainLink Copied

Let's run this implementation and verify that it produces the expected shapes and efficiency characteristics.

In[4]:

Code

# Demonstrate the efficiency difference
n, d = 128, 64

Q = np.random.randn(n, d).astype(np.float32)
K = np.random.randn(n, d).astype(np.float32)
V = np.random.randn(n, d).astype(np.float32)

output = linear_attention(Q, K, V)

Out[5]:

Console

Input shapes: Q=(128, 64), K=(128, 64), V=(128, 64)
Output shape: (128, 64)

Key computation: K'^T @ V creates a 64x64 matrix
Standard attention would create a 128x128 matrix
Ratio: 4.0x smaller intermediate representation

The output confirms the efficiency gain: instead of a 128×128 attention matrix (16,384 elements), linear attention uses a 64×64 key-value aggregate (4,096 elements). This 4x reduction might seem modest, but remember that the ratio scales as $(n/d)^2$ . At 1,024 tokens with $d = 64$ , the reduction is 256x. At 8,192 tokens, it's over 16,000x. The quadratic-to-linear improvement becomes increasingly dramatic as sequences grow longer.