Linear Attention › Comparing Standard and Linear Attention

Michael BrenndoerferUpdated June 28, 20253 min read
Language AI HandbookAll chapters
Part 17 of 67Efficient Attention
Chapter 123Linear Attention › Comparing Standard and Linear Attention
Reading section 5 of 15 · Comparing Standard and Linear Attention · 3 minSection 5 of 15 · 3 min

Book contents

Language AI Handbook67 parts · 513 chaptersView the whole book →
Active chapter 123 of 513

Section 5 of 15 · 3 min read

Linear Attention › Comparing Standard and Linear AttentionOne section · Sections in this chapter5 / 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

Let's implement both attention mechanisms and compare their outputs to understand what linear attention preserves and what it changes.

In[10]:
Code
def softmax_attention(Q, K, V):
    """Standard scaled dot-product attention."""
    d_k = Q.shape[-1]
    scores = Q @ K.T / np.sqrt(d_k)

    # Softmax normalization
    exp_scores = np.exp(scores - np.max(scores, axis=-1, keepdims=True))
    attention_weights = exp_scores / np.sum(exp_scores, axis=-1, keepdims=True)

    return attention_weights @ V, attention_weights


def compare_attention_mechanisms(n, d, seed=42):
    """Compare outputs of standard and linear attention."""
    np.random.seed(seed)

    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)

    # Standard attention
    output_standard, weights_standard = softmax_attention(Q, K, V)

    # Linear attention with ELU
    output_linear = linear_attention(Q, K, V, feature_map=elu_feature_map)

    # Compute similarity metrics
    cosine_sim = np.mean(
        [
            np.dot(output_standard[i], output_linear[i])
            / (
                np.linalg.norm(output_standard[i])
                * np.linalg.norm(output_linear[i])
            )
            for i in range(n)
        ]
    )

    mse = np.mean((output_standard - output_linear) ** 2)

    return {
        "output_standard": output_standard,
        "output_linear": output_linear,
        "weights_standard": weights_standard,
        "cosine_similarity": cosine_sim,
        "mse": mse,
    }
In[11]:
Code
results = compare_attention_mechanisms(n=64, d=32)
Out[12]:
Console
Comparison of Standard vs Linear Attention (ELU feature map):

Sequence length: 64, Dimension: 32
Cosine similarity between outputs: 0.9846
Mean squared error: 0.000780

Output magnitude ratio: 0.975

The cosine similarity shows that linear attention with a simple ELU feature map produces outputs that point in a similar direction to standard attention but aren't identical. The MSE quantifies the actual difference in output values.

Out[13]:
Visualization
Heatmap of softmax attention weights showing concentrated dark spots indicating strong attention to specific positions.
Standard softmax attention produces sharp, focused patterns. High-weight positions (dark blue) capture specific dependencies, while most positions receive near-zero attention.
Heatmap of linear attention showing more uniform, spread-out attention pattern with less contrast between positions.
Linear attention with ELU features produces smoother, more diffuse patterns. The lack of softmax normalization means attention weights don't sum to 1 per row, changing the weighting dynamics.

The visualizations reveal a fundamental difference in attention behavior. Softmax attention concentrates weight on a few key positions, creating the "sharp" attention patterns that transformers are known for. Linear attention with simple feature maps spreads weight more uniformly, producing softer patterns that may miss fine-grained dependencies.