Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Linear Attention › Causal Linear Attention

Michael BrenndoerferUpdated June 28, 20254 min read

Language AI HandbookAll chapters

Part 17 of 67Efficient Attention

Chapter 123Linear Attention › Causal Linear Attention

Reading section 8 of 15 · Causal Linear Attention · 4 minSection 8 of 15 · 4 min

Section 8 of 15 · 4 min read

Linear Attention › Causal Linear AttentionOne section · Sections in this chapter8 / 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

For autoregressive language models, we need causal attention: each position can only attend to previous positions. Standard causal attention applies a triangular mask before softmax. Linear attention requires a different approach.

The key insight is that the key-value aggregate can be computed incrementally. Instead of summing over all positions, we maintain a running sum up to the current position:

S_i = \sum_{j=1}^{i} \phi(k_j) v_j^T

Z_i = \sum_{j=1}^{i} \phi(k_j)

\text{output}_i = \frac{\phi(q_i)^T S_i}{\phi(q_i)^T Z_i}

where:

$S_i \in \mathbb{R}^{D \times d_v}$ : the cumulative key-value aggregate up to position $i$ , computed by summing outer products $\phi(k_j) v_j^T$ for all positions $j \leq i$
$Z_i \in \mathbb{R}^{D}$ : the cumulative key sum up to position $i$ , used for normalization
$\phi(q_i)^T S_i$ : retrieves a weighted combination of values based on the query at position $i$
$\phi(q_i)^T Z_i$ : computes the normalization factor

This formulation naturally respects causality: position $i$ only accesses information from positions $1$ through $i$ . The running sum structure enables efficient incremental computation during inference.

In[21]:

Code

def causal_linear_attention(Q, K, V, feature_map=None, eps=1e-6):
    """
    Compute causal linear attention where each position attends only to previous positions.

    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
        eps: Numerical stability constant

    Returns:
        Output of shape (n, d_v)
    """
    if feature_map is None:
        feature_map = elu_feature_map

    n, d_k = Q.shape
    d_v = V.shape[1]

    Q_prime = feature_map(Q)  # (n, d_k)
    K_prime = feature_map(K)  # (n, d_k)

    # Initialize running sums
    S = np.zeros((d_k, d_v), dtype=Q.dtype)  # Key-value aggregate
    Z = np.zeros(d_k, dtype=Q.dtype)  # Key sum for normalization

    outputs = np.zeros((n, d_v), dtype=Q.dtype)

    for i in range(n):
        # Update running sums with current position
        S += np.outer(K_prime[i], V[i])
        Z += K_prime[i]

        # Compute output for position i
        numerator = Q_prime[i] @ S
        denominator = Q_prime[i] @ Z + eps
        outputs[i] = numerator / denominator

    return outputs

In[22]:

Code

np.random.seed(42)
n, d = 32, 16
Q = np.random.randn(n, d).astype(np.float32) * 0.5
K = np.random.randn(n, d).astype(np.float32) * 0.5
V = np.random.randn(n, d).astype(np.float32)

output_causal = causal_linear_attention(Q, K, V)

# Compare with non-causal
output_full = linear_attention(Q, K, V, feature_map=elu_feature_map)

# The outputs should differ because causal restricts attention
diff = np.mean(np.abs(output_causal - output_full))

Out[23]:

Console

Causal vs Non-causal Linear Attention:
Mean absolute difference: 0.2188

This difference is expected: causal attention sees less context per position.

The causal formulation is particularly elegant because the running sum structure means that during inference, we can process one token at a time with $O(1)$ memory overhead per step (just updating the $d \times d$ aggregate). This makes linear attention ideal for autoregressive generation.

Out[24]:

Visualization

Lower-triangular heatmap of causal attention weights where positions only attend to earlier positions in the sequence. — Causal linear attention pattern showing the lower-triangular structure where each position attends only to itself and previous positions. The running sum formulation naturally enforces this constraint without explicit masking.