Module 05: Attention

Introduction

Attention made transformers revolutionary. Each token examines every other token and gathers relevant information.

Attention enables each token to ask: “Which tokens in this sequence matter to me?”

Why attention matters for LLMs:

  • Long-range dependencies: Token 100 can attend to token 1—solving the vanishing gradient problem that cripples RNNs
  • Parallelization: All positions compute simultaneously during training, unlike in RNNs
  • Interpretability: Attention weights reveal what the model examines
  • Dynamic context: Each token’s representation is context-dependent, not fixed

Self-attention is the key innovation: tokens attend to other tokens within the same sequence - queries, keys, and values all come from the same input. Cross-attention (used in encoder-decoder models) draws queries from one sequence and keys/values from another.

What You’ll Learn

After this module, you can:

  • Understand Query, Key, Value projections and their roles
  • Implement scaled dot-product attention from scratch
  • Apply causal masking for autoregressive models
  • Build multi-head attention and understand why it’s beneficial
  • Recognize attention patterns and what they reveal
  • Shrink the KV cache with grouped-query and multi-query attention (GQA/MQA)
  • Cancel attention noise with differential attention (a difference of two softmaxes)

Prerequisites

This module requires familiarity with:

Note: Attention treats tokens as an unordered set. Positional embeddings (Module 04) supply the sense of order.

Attention as Three Questions

Every token in a sequence asks three questions. These questions unlock attention.

import numpy as np

# A simple sentence
tokens = ["The", "cat", "sat"]

# Each token has an embedding (we'll use random ones for illustration)
np.random.seed(42)
embed_dim = 4
embeddings = {tok: np.random.randn(embed_dim).round(2) for tok in tokens}

print("Each token has an embedding vector:")
for tok, emb in embeddings.items():
    print(f"  '{tok}': {emb}")
Each token has an embedding vector:
  'The': [ 0.5  -0.14  0.65  1.52]
  'cat': [-0.23 -0.23  1.58  0.77]
  'sat': [-0.47  0.54 -0.46 -0.47]

The Three Questions:

Question Vector What it asks
Query (Q) “What am I looking for?” Token seeks relevant context
Key (K) “What do I contain?” Token advertises its content
Value (V) “What do I return if matched?” Token’s actual information 
# Each token projects its embedding into Q, K, V
# These are learned linear transformations

# For "sat", the query might encode: "I need a subject (who sat?)"
# For "cat", the key might encode: "I'm a noun, a subject candidate"
# For "cat", the value carries: the actual semantic content of "cat"

print("When 'sat' attends to 'cat':")
print("  Q_sat . K_cat = high score (sat is looking for a subject, cat is one)")
print("  The output for 'sat' includes V_cat weighted by this score")
When 'sat' attends to 'cat':
  Q_sat . K_cat = high score (sat is looking for a subject, cat is one)
  The output for 'sat' includes V_cat weighted by this score

Q, K, and V are learned projections: the model learns what to seek (Q), how to advertise content (K), and what information to transmit (V).

Intuition: Query, Key, Value

Think of attention as a “soft lookup” - like a database query, but differentiable:

Query:   "What information do I need?"     (the question)
Keys:    "What information do I have?"     (index/labels for content)
Values:  "Here's my actual information"    (the content itself)

Attention = softmax(Query . Keys) x Values

Analogy: Imagine a library where:

  • Your query is “books about cats”
  • Each book has a key (its topic/keywords)
  • Each book has a value (its actual content)
  • You get a weighted average of book contents based on how well they match your query

For the sentence “The cat sat on the mat”:

  • “sat” might attend strongly to “cat” (who sat?) and weakly to “mat” (where?)
  • “mat” might attend strongly to “the” and “on” (which mat? on what?)

Softmax normalizes each row of attention weights to sum to 1.

The Math: Scaled Dot-Product Attention

The attention formula:

Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) x V

Where:

  • Q (Query): What am I looking for? Shape: (seq, d_k)
  • K (Key): What do I have to offer? Shape: (seq, d_k)
  • V (Value): What information do I carry? Shape: (seq, d_v)
  • d_k: Dimension of keys (for scaling)

Step by Step

Building Attention by Hand

Before using PyTorch, let’s build attention with NumPy to see what happens at each step.

import numpy as np

def attention_from_scratch(x, W_q, W_k, W_v):
    """
    Single-head attention implemented with pure NumPy.

    Args:
        x: Input embeddings (seq_len, embed_dim)
        W_q, W_k, W_v: Projection matrices (embed_dim, head_dim)

    Returns:
        output: Attended values (seq_len, head_dim)
        weights: Attention weights (seq_len, seq_len)
    """
    # Step 1: Project input into Q, K, V
    Q = x @ W_q  # (seq, head_dim) - "What am I looking for?"
    K = x @ W_k  # (seq, head_dim) - "What do I contain?"
    V = x @ W_v  # (seq, head_dim) - "What do I return?"

    print(f"Input x shape: {x.shape}")
    print(f"Q = x @ W_q: {Q.shape}")
    print(f"K = x @ W_k: {K.shape}")
    print(f"V = x @ W_v: {V.shape}")

    # Step 2: Compute attention scores
    # Each query attends to all keys: Q @ K.T
    d_k = K.shape[-1]
    scores = Q @ K.T  # (seq, seq) - similarity between every pair
    scores = scores / np.sqrt(d_k)  # Scale to prevent softmax saturation

    print(f"\nScores = Q @ K.T / sqrt({d_k}): {scores.shape}")
    print(f"Score matrix (who attends to whom):")
    print(scores.round(2))

    # Step 3: Softmax to get attention weights
    # Each row sums to 1: how much each position attends to others
    def softmax(x):
        exp_x = np.exp(x - x.max(axis=-1, keepdims=True))  # Numerical stability
        return exp_x / exp_x.sum(axis=-1, keepdims=True)

    weights = softmax(scores)

    print(f"\nAttention weights (each row sums to 1):")
    print(weights.round(3))
    print(f"Row sums: {weights.sum(axis=-1).round(3)}")

    # Step 4: Weighted sum of values
    output = weights @ V  # (seq, head_dim)

    print(f"\nOutput = weights @ V: {output.shape}")

    return output, weights

# Demo with a tiny example
np.random.seed(42)
seq_len, embed_dim, head_dim = 3, 4, 2

x = np.random.randn(seq_len, embed_dim)
W_q = np.random.randn(embed_dim, head_dim) * 0.5
W_k = np.random.randn(embed_dim, head_dim) * 0.5
W_v = np.random.randn(embed_dim, head_dim) * 0.5

print("=" * 50)
print("ATTENTION FROM SCRATCH")
print("=" * 50)
output, weights = attention_from_scratch(x, W_q, W_k, W_v)
==================================================
ATTENTION FROM SCRATCH
==================================================
Input x shape: (3, 4)
Q = x @ W_q: (3, 2)
K = x @ W_k: (3, 2)
V = x @ W_v: (3, 2)

Scores = Q @ K.T / sqrt(2): (3, 3)
Score matrix (who attends to whom):
[[ 0.05  0.2   0.4 ]
 [ 0.48  0.72 -0.05]
 [ 0.18  0.22 -0.18]]

Attention weights (each row sums to 1):
[[0.278 0.324 0.397]
 [0.348 0.445 0.206]
 [0.365 0.381 0.255]]
Row sums: [1. 1. 1.]

Output = weights @ V: (3, 2)

Key Insight: Attention reduces to four matrix multiplications: 1. Q = x @ W_q - Project to queries 2. K = x @ W_k - Project to keys 3. scores = Q @ K.T / sqrt(d_k) - Compute similarities 4. output = softmax(scores) @ V - Weighted sum of values

Why Scale by sqrt(d_k)?

Scaling prevents large d_k from producing dot products that saturate softmax and shrink gradients to near-zero. Here’s the intuition:

The Problem: When Q and K have elements drawn from a distribution with mean 0 and variance 1, their dot product has variance proportional to d_k. For d_k = 64, dot products can easily reach values like 8 or -10.

Why This Matters: Softmax of large values produces near-one-hot distributions:

  • softmax([10, 0, 0]) = [0.9999, 0.00005, 0.00005]

This causes:

  1. Vanishing gradients: The gradient of softmax approaches 0 at extremes
  2. Loss of information: We want soft attention, not hard selection

The Solution: Dividing by sqrt(d_k) normalizes variance back to ~1.

import torch
import torch.nn.functional as F
import math

# Show the effect of scaling
d_k = 64
q = torch.randn(1, d_k)
k = torch.randn(1, d_k)

dot_product = (q @ k.T).item()
scaled = dot_product / math.sqrt(d_k)

print(f"d_k = {d_k}")
print(f"Raw dot product: {dot_product:.2f}")
print(f"Scaled by sqrt({d_k}) = {math.sqrt(d_k):.1f}: {scaled:.2f}")
print(f"\nScaling keeps values in a reasonable range for softmax")

# Demonstrate the gradient problem
scores_large = torch.tensor([[10.0, 1.0, 1.0]], requires_grad=True)
scores_normal = torch.tensor([[1.0, 0.5, 0.5]], requires_grad=True)

weights_large = F.softmax(scores_large, dim=-1)
weights_normal = F.softmax(scores_normal, dim=-1)

print(f"\nLarge scores [10, 1, 1] -> softmax: {weights_large.detach().numpy().round(4)}")
print(f"Normal scores [1, 0.5, 0.5] -> softmax: {weights_normal.detach().numpy().round(3)}")
d_k = 64
Raw dot product: -2.46
Scaled by sqrt(64) = 8.0: -0.31

Scaling keeps values in a reasonable range for softmax

Large scores [10, 1, 1] -> softmax: [[9.998e-01 1.000e-04 1.000e-04]]
Normal scores [1, 0.5, 0.5] -> softmax: [[0.452 0.274 0.274]]

Numerical Stability in Softmax

The naive softmax implementation hides a danger: overflow.

import numpy as np

# The naive softmax
def naive_softmax(x):
    """This will overflow for large values!"""
    exp_x = np.exp(x)
    return exp_x / exp_x.sum()

# Try it with large values
large_scores = np.array([1000.0, 1001.0, 1002.0])

print("Large scores:", large_scores)
print("exp(1000) =", np.exp(1000))  # This is inf!

try:
    result = naive_softmax(large_scores)
    print("Naive softmax:", result)  # Will be [nan, nan, nan]
except:
    print("Overflow error!")
Large scores: [1000. 1001. 1002.]
exp(1000) = inf
Naive softmax: [nan nan nan]
/var/folders/hl/bw75m5hd5xvfyx8j9qd71vjw0000gn/T/ipykernel_65533/757484777.py:13: RuntimeWarning: overflow encountered in exp
  print("exp(1000) =", np.exp(1000))  # This is inf!
/var/folders/hl/bw75m5hd5xvfyx8j9qd71vjw0000gn/T/ipykernel_65533/757484777.py:6: RuntimeWarning: overflow encountered in exp
  exp_x = np.exp(x)
/var/folders/hl/bw75m5hd5xvfyx8j9qd71vjw0000gn/T/ipykernel_65533/757484777.py:7: RuntimeWarning: invalid value encountered in divide
  return exp_x / exp_x.sum()

The Problem: exp(1000) is astronomically large - it overflows to infinity. Even exp(100) is about 2.7 x 10^43.

The Solution: The max-subtraction trick. Subtract the maximum value before exponentiating.

def stable_softmax(x):
    """
    Numerically stable softmax using the max-subtraction trick.

    Key insight: softmax(x) = softmax(x - c) for any constant c
    We choose c = max(x) to keep values small.
    """
    # Subtract max for numerical stability
    x_shifted = x - x.max()

    print(f"Original: {x}")
    print(f"After subtracting max ({x.max()}): {x_shifted}")
    print(f"Now exp() won't overflow: exp({x_shifted}) = {np.exp(x_shifted)}")

    exp_x = np.exp(x_shifted)
    return exp_x / exp_x.sum()

print("Stable softmax with max-subtraction trick:")
print("=" * 50)
large_scores = np.array([1000.0, 1001.0, 1002.0])
result = stable_softmax(large_scores)
print(f"\nResult: {result}")
print(f"Sum: {result.sum()}")  # Should be 1.0
Stable softmax with max-subtraction trick:
==================================================
Original: [1000. 1001. 1002.]
After subtracting max (1002.0): [-2. -1.  0.]
Now exp() won't overflow: exp([-2. -1.  0.]) = [0.13533528 0.36787944 1.        ]

Result: [0.09003057 0.24472847 0.66524096]
Sum: 0.9999999999999999

Why does this work mathematically?

softmax(x)_i = exp(x_i) / sum(exp(x_j))
            = exp(x_i - c) / sum(exp(x_j - c))    [multiply by exp(-c)/exp(-c)]

For c = max(x), all exponents are <= 0, so exp() stays bounded.
# Verify the math: both give same result
normal_scores = np.array([2.0, 1.0, 0.1])

naive_result = naive_softmax(normal_scores)
stable_result = stable_softmax(normal_scores)

print(f"\nNaive:  {naive_result}")
print(f"Stable: {stable_result}")
print(f"Same? {np.allclose(naive_result, stable_result)}")
Original: [2.  1.  0.1]
After subtracting max (2.0): [ 0.  -1.  -1.9]
Now exp() won't overflow: exp([ 0.  -1.  -1.9]) = [1.         0.36787944 0.14956862]

Naive:  [0.65900114 0.24243297 0.09856589]
Stable: [0.65900114 0.24243297 0.09856589]
Same? True
WarningAlways Use Stable Softmax

PyTorch’s F.softmax applies the max-subtraction trick automatically. Naive softmax fails silently, returning NaN when scores grow large.

Code: Scaled Dot-Product Attention

Let’s implement attention step by step. This follows the exact algorithm from the attention.py module:

def scaled_dot_product_attention(query, key, value, mask=None):
    """
    Compute scaled dot-product attention.

    Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) x V

    Args:
        query: (batch, seq, d_k) or (..., seq, d_k)
        key: (batch, seq, d_k)
        value: (batch, seq, d_v)  # d_v can differ from d_k
        mask: Optional mask where 0 = masked, 1 = attend

    Returns:
        output: (batch, seq, d_v)
        attention_weights: (batch, seq, seq)

    Note: The mask convention (0 = masked) matches the module implementation.
    Masked positions get -inf before softmax, becoming 0 after.
    """
    d_k = query.size(-1)

    # Step 1: Compute similarity scores
    # QK^T: (..., seq, d_k) @ (..., d_k, seq) -> (..., seq, seq)
    scores = torch.matmul(query, key.transpose(-2, -1))

    # Step 2: Scale by sqrt(d_k)
    scores = scores / math.sqrt(d_k)

    # Step 3: Apply mask (if provided)
    # Masked positions get -inf, which becomes 0 after softmax
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))

    # Step 4: Softmax (each row sums to 1)
    attention_weights = F.softmax(scores, dim=-1)

    # Step 5: Weighted sum of values
    output = torch.matmul(attention_weights, value)

    return output, attention_weights

# Test it
batch, seq, d_k = 1, 4, 8
Q = torch.randn(batch, seq, d_k)
K = torch.randn(batch, seq, d_k)
V = torch.randn(batch, seq, d_k)

output, weights = scaled_dot_product_attention(Q, K, V)

print(f"Query shape: {Q.shape}")
print(f"Key shape: {K.shape}")
print(f"Value shape: {V.shape}")
print(f"Output shape: {output.shape}")
print(f"Attention weights shape: {weights.shape}")
print(f"\nAttention weights (each row sums to 1):")
print(weights[0].round(decimals=2).numpy())
print(f"\nRow sums: {weights[0].sum(dim=-1).numpy()}")
Query shape: torch.Size([1, 4, 8])
Key shape: torch.Size([1, 4, 8])
Value shape: torch.Size([1, 4, 8])
Output shape: torch.Size([1, 4, 8])
Attention weights shape: torch.Size([1, 4, 4])

Attention weights (each row sums to 1):
[[0.2  0.16 0.27 0.36]
 [0.69 0.08 0.03 0.2 ]
 [0.08 0.29 0.52 0.12]
 [0.24 0.37 0.21 0.18]]

Row sums: [1. 1. 1. 1.]

Visualizing Attention Patterns

# Extract attention weights for OJS visualization
attention_weights_for_viz = weights[0].detach().numpy().tolist()
ojs_define(attention_weights_viz = attention_weights_for_viz)

Causal Masking for Language Models

In autoregressive models (like GPT), each token attends only to previous tokens, never future ones. A causal mask enforces this constraint:

NoteMask Convention

In this lesson, we use an additive mask where:

  • 0 = attend (score unchanged)
  • -inf = masked (softmax converts to 0)

Some libraries use a boolean mask (True = attend, False = masked) which is converted internally. The key insight: positions with -inf before softmax become 0 attention weight.

The Causal Mask from Scratch

The causal mask is elegantly simple: we add -inf to positions we want to mask, and softmax turns those into zeros.

import numpy as np

def causal_mask_from_scratch(seq_len):
    """
    Create a causal mask using np.triu (upper triangular).

    The mask has -inf above the diagonal (future positions)
    and 0 on and below the diagonal (past/current positions).
    """
    # np.triu with k=1 gives us the strictly upper triangular part
    # (everything above the main diagonal)
    mask = np.triu(np.ones((seq_len, seq_len)), k=1)

    # Convert 1s to -inf (positions to mask)
    mask = mask * (-1e9)  # Use large negative instead of -inf for visualization

    return mask

# Visualize the mask
seq_len = 5
mask = causal_mask_from_scratch(seq_len)

print("Causal Mask (0 = attend, -inf = masked):")
print(mask.round(0))

print("\nHow it works:")
print("  Position 0: sees only position 0")
print("  Position 1: sees positions 0, 1")
print("  Position 4: sees positions 0, 1, 2, 3, 4")
Causal Mask (0 = attend, -inf = masked):
[[-0.e+00 -1.e+09 -1.e+09 -1.e+09 -1.e+09]
 [-0.e+00 -0.e+00 -1.e+09 -1.e+09 -1.e+09]
 [-0.e+00 -0.e+00 -0.e+00 -1.e+09 -1.e+09]
 [-0.e+00 -0.e+00 -0.e+00 -0.e+00 -1.e+09]
 [-0.e+00 -0.e+00 -0.e+00 -0.e+00 -0.e+00]]

How it works:
  Position 0: sees only position 0
  Position 1: sees positions 0, 1
  Position 4: sees positions 0, 1, 2, 3, 4
def attention_with_causal_mask(x, W_q, W_k, W_v):
    """
    Causal attention from scratch - each position only attends to past.
    """
    Q = x @ W_q
    K = x @ W_k
    V = x @ W_v

    seq_len = x.shape[0]
    d_k = K.shape[-1]

    # Compute scores
    scores = Q @ K.T / np.sqrt(d_k)

    print("Scores before masking:")
    print(scores.round(2))

    # Add causal mask: -inf for future positions
    mask = np.triu(np.ones((seq_len, seq_len)), k=1) * (-1e9)
    scores = scores + mask

    print("\nScores after adding causal mask:")
    print(scores.round(2))

    # Softmax: -inf becomes 0
    def softmax(x):
        exp_x = np.exp(x - x.max(axis=-1, keepdims=True))
        return exp_x / exp_x.sum(axis=-1, keepdims=True)

    weights = softmax(scores)

    print("\nAttention weights (upper triangle is 0!):")
    print(weights.round(3))

    output = weights @ V
    return output, weights

# Demo
np.random.seed(42)
seq_len, embed_dim, head_dim = 4, 4, 2
x = np.random.randn(seq_len, embed_dim)
W_q = np.random.randn(embed_dim, head_dim) * 0.5
W_k = np.random.randn(embed_dim, head_dim) * 0.5
W_v = np.random.randn(embed_dim, head_dim) * 0.5

print("=" * 50)
print("CAUSAL ATTENTION FROM SCRATCH")
print("=" * 50)
output, weights = attention_with_causal_mask(x, W_q, W_k, W_v)
==================================================
CAUSAL ATTENTION FROM SCRATCH
==================================================
Scores before masking:
[[-1.08 -0.42  0.22  0.84]
 [-1.26 -0.68  0.22  1.97]
 [ 0.11  0.11 -0.01 -0.41]
 [ 2.12  0.79 -0.44 -1.52]]

Scores after adding causal mask:
[[-1.08000000e+00 -1.00000000e+09 -1.00000000e+09 -9.99999999e+08]
 [-1.26000000e+00 -6.80000000e-01 -1.00000000e+09 -9.99999998e+08]
 [ 1.10000000e-01  1.10000000e-01 -1.00000000e-02 -1.00000000e+09]
 [ 2.12000000e+00  7.90000000e-01 -4.40000000e-01 -1.52000000e+00]]

Attention weights (upper triangle is 0!):
[[1.    0.    0.    0.   ]
 [0.359 0.641 0.    0.   ]
 [0.348 0.345 0.307 0.   ]
 [0.731 0.193 0.057 0.019]]
NoteKey Insight

Causal masking is just adding -inf before softmax. That is the entire trick.

  • softmax([2.0, 1.0, -inf]) = [0.73, 0.27, 0.00]
  • The -inf position gets exactly 0 weight
  • No information flows from future to past
def create_causal_mask(seq_len):
    """Create a lower triangular causal mask."""
    return torch.tril(torch.ones(seq_len, seq_len))

# Show the mask
seq_len = 6
mask = create_causal_mask(seq_len)

print("Causal Mask (1 = can attend, 0 = masked):")
print()
tokens = ["The", "cat", "sat", "on", "the", "mat"]
for i in range(seq_len):
    row = ['#' if mask[i, j] == 1 else '.' for j in range(seq_len)]
    print(f"  {tokens[i]:4s}: {''.join(row)}")

print(f"\nPosition 0 can only see position 0")
print(f"Position 5 can see all previous positions")
Causal Mask (1 = can attend, 0 = masked):

  The : #.....
  cat : ##....
  sat : ###...
  on  : ####..
  the : #####.
  mat : ######

Position 0 can only see position 0
Position 5 can see all previous positions
# Apply causal mask
Q = torch.randn(1, 6, 8)
K = torch.randn(1, 6, 8)
V = torch.randn(1, 6, 8)

# Without mask (bidirectional)
output_bi, weights_bi = scaled_dot_product_attention(Q, K, V)

# With causal mask
output_causal, weights_causal = scaled_dot_product_attention(Q, K, V, mask=mask)

# Pass to OJS for visualization
ojs_define(
    weights_bi_viz = weights_bi[0].detach().numpy().tolist(),
    weights_causal_viz = weights_causal[0].detach().numpy().tolist(),
    mask_tokens = tokens
)

print("Notice: In causal attention, the upper triangle is 0 (can't attend to future)")

Multi-Head Attention

Multiple attention heads, instead of one, learn different patterns:

MultiHead(Q, K, V) = Concat(head_1, ..., head_h) x W_O

where head_i = Attention(Q x W_Q_i, K x W_K_i, V x W_V_i)

Why Multiple Heads?

A single attention head computes one weighted average, which limits what relationships it can capture. Multiple heads provide:

  • Diverse patterns: Different heads focus on different relationships: syntax, semantics, position, coreference
  • Subspace attention: Each head operates in a lower-dimensional subspace (head_dim = embed_dim / num_heads), allowing specialized representations
  • Computational efficiency: Despite having multiple heads, the total computation is similar to single-head attention with full dimensionality (same number of parameters)

Typical configurations:

  • GPT-2: 12 heads, 768 embed_dim, 64 head_dim
  • GPT-3: 96 heads, 12288 embed_dim, 128 head_dim
  • Llama 2 (7B): 32 heads, 4096 embed_dim, 128 head_dim

What Different Heads Learn

Trained models show head specialization:

  • Head 0: “Who did what?” - attends to subject-verb pairs
  • Head 1: “What comes before?” - attends to previous token
  • Head 2: “What’s similar?” - attends to semantically similar words
  • Head 3: “Syntax patterns” - attends to grammatical structure
import torch.nn as nn

# Simplified implementation for learning - see attention.py for production version
class MultiHeadAttention(nn.Module):
    """Multi-head attention with separate Q, K, V projections (simplified for illustration)."""

    def __init__(self, embed_dim, num_heads, dropout=0.0):
        super().__init__()
        assert embed_dim % num_heads == 0

        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads

        # Projections for Q, K, V
        self.q_proj = nn.Linear(embed_dim, embed_dim)
        self.k_proj = nn.Linear(embed_dim, embed_dim)
        self.v_proj = nn.Linear(embed_dim, embed_dim)

        # Output projection
        self.out_proj = nn.Linear(embed_dim, embed_dim)

        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None, return_attention=False):
        batch_size, seq_len, _ = x.shape

        # Project Q, K, V
        q = self.q_proj(x)  # (batch, seq, embed)
        k = self.k_proj(x)
        v = self.v_proj(x)

        # Reshape for multi-head: (batch, seq, embed) -> (batch, heads, seq, head_dim)
        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)

        # Compute attention
        d_k = q.size(-1)
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)

        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))

        attention_weights = F.softmax(scores, dim=-1)
        attention_weights = self.dropout(attention_weights)

        # Apply attention to values
        attn_output = torch.matmul(attention_weights, v)

        # Reshape back: (batch, heads, seq, head_dim) -> (batch, seq, embed)
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.view(batch_size, seq_len, self.embed_dim)

        # Final projection
        output = self.out_proj(attn_output)

        if return_attention:
            return output, attention_weights
        return output

# Test multi-head attention
embed_dim = 64
num_heads = 8
mha = MultiHeadAttention(embed_dim=embed_dim, num_heads=num_heads)

x = torch.randn(2, 10, embed_dim)
output, weights = mha(x, return_attention=True)

print(f"Multi-Head Attention Configuration:")
print(f"  Embedding dimension: {embed_dim}")
print(f"  Number of heads: {num_heads}")
print(f"  Dimension per head: {embed_dim // num_heads}")
print(f"\nInput shape: {x.shape}")
print(f"Output shape: {output.shape}")
print(f"Attention weights shape: {weights.shape}")
print(f"  (batch, heads, query_pos, key_pos)")
print(f"\nTotal parameters: {sum(p.numel() for p in mha.parameters()):,}")
Multi-Head Attention Configuration:
  Embedding dimension: 64
  Number of heads: 8
  Dimension per head: 8

Input shape: torch.Size([2, 10, 64])
Output shape: torch.Size([2, 10, 64])
Attention weights shape: torch.Size([2, 8, 10, 10])
  (batch, heads, query_pos, key_pos)

Total parameters: 16,640

Visualizing Multi-Head Attention

# Pass multi-head attention weights to OJS
multi_head_weights = [weights[0, h].detach().numpy().tolist() for h in range(num_heads)]
ojs_define(mha_weights = multi_head_weights, mha_num_heads = num_heads)

Using Our Attention Module

The attention.py module provides production-ready implementations:

from attention import (
    CausalMultiHeadAttention,
    demonstrate_attention,
    demonstrate_causal_attention
)

# Run built-in demonstrations
print("=" * 60)
print("MULTI-HEAD ATTENTION DEMONSTRATION")
print("=" * 60)
demonstrate_attention(seq_len=6, embed_dim=32, num_heads=4)
============================================================
MULTI-HEAD ATTENTION DEMONSTRATION
============================================================
============================================================
ATTENTION DEMONSTRATION
============================================================

Multi-Head Attention:
  Embed dim: 32
  Num heads: 4
  Head dim: 8

Input shape: (1, 6, 32)
  (batch=1, seq_len=6, embed_dim=32)

Output shape: (1, 6, 32)
Attention weights shape: (1, 4, 6, 6)
  (batch, heads, seq, seq)

Attention weights for head 0, position 0:
  [0.16073556244373322, 0.23629829287528992, 0.31348538398742676, 0.11531048268079758, 0.07284009456634521, 0.1013302132487297]
  Sum: 1.0000 (should be 1.0)
MultiHeadAttention(
  (q_proj): Linear(in_features=32, out_features=32, bias=True)
  (k_proj): Linear(in_features=32, out_features=32, bias=True)
  (v_proj): Linear(in_features=32, out_features=32, bias=True)
  (out_proj): Linear(in_features=32, out_features=32, bias=True)
  (attention): ScaledDotProductAttention()
  (dropout): Dropout(p=0.0, inplace=False)
)
print("\n" + "=" * 60)
print("CAUSAL ATTENTION DEMONSTRATION")
print("=" * 60)
demonstrate_causal_attention(seq_len=6)

============================================================
CAUSAL ATTENTION DEMONSTRATION
============================================================
============================================================
CAUSAL ATTENTION DEMONSTRATION
============================================================

Causal mask for seq_len=6:
(1 = can attend, 0 = masked)
  Position 0: █·····
  Position 1: ██····
  Position 2: ███···
  Position 3: ████··
  Position 4: █████·
  Position 5: ██████

Interpretation:
  Position 0: can only see position 0
  Position 1: can see positions 0, 1
  Position 5: can see all positions

Without mask (position 0 attends to all):
  [0.1269863396883011, 0.30605772137641907, 0.04593325033783913, 0.3218763768672943, 0.039763979613780975, 0.15938231348991394]

With causal mask (position 0 only attends to itself):
  [1.0, 0.0, 0.0, 0.0, 0.0, 0.0]
# Causal multi-head attention (what GPT uses)
causal_mha = CausalMultiHeadAttention(
    embed_dim=64,
    num_heads=8,
    max_seq_len=512,
    dropout=0.0
)

x = torch.randn(2, 10, 64)
output, weights = causal_mha(x, return_attention=True)

print(f"\nCausal Multi-Head Attention:")
print(f"  Input shape: {x.shape}")
print(f"  Output shape: {output.shape}")
print(f"  Attention weights shape: {weights.shape}")

Causal Multi-Head Attention:
  Input shape: torch.Size([2, 10, 64])
  Output shape: torch.Size([2, 10, 64])
  Attention weights shape: torch.Size([2, 8, 10, 10])

PyTorch’s Optimized Attention

Now that we understand attention from scratch, PyTorch’s production-optimized implementations offer a faster path.

F.scaled_dot_product_attention

PyTorch 2.0+ provides F.scaled_dot_product_attention - a single function that replaces our manual implementation and automatically uses the best available backend.

import torch
import torch.nn.functional as F

# Our inputs
batch, num_heads, seq_len, head_dim = 2, 8, 64, 32
query = torch.randn(batch, num_heads, seq_len, head_dim)
key = torch.randn(batch, num_heads, seq_len, head_dim)
value = torch.randn(batch, num_heads, seq_len, head_dim)

# The manual way (what we implemented)
def manual_attention(q, k, v, is_causal=False):
    d_k = q.size(-1)
    scores = q @ k.transpose(-2, -1) / (d_k ** 0.5)
    if is_causal:
        mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool()
        scores.masked_fill_(mask, float('-inf'))
    weights = F.softmax(scores, dim=-1)
    return weights @ v

# PyTorch's optimized version
output_manual = manual_attention(query, key, value, is_causal=True)
output_pytorch = F.scaled_dot_product_attention(query, key, value, is_causal=True)

print(f"Manual output shape: {output_manual.shape}")
print(f"PyTorch SDPA shape: {output_pytorch.shape}")
print(f"Results match: {torch.allclose(output_manual, output_pytorch, atol=1e-5)}")
Manual output shape: torch.Size([2, 8, 64, 32])
PyTorch SDPA shape: torch.Size([2, 8, 64, 32])
Results match: True

Flash Attention: The Speed Revolution

F.scaled_dot_product_attention invokes Flash Attention when the backend supports it - a breakthrough algorithm that:

  1. Avoids materializing the full attention matrix (O(n^2) memory -> O(n) memory)
  2. Uses tiling to keep computation in fast GPU SRAM
  3. Fuses operations to minimize memory bandwidth bottleneck
# Check which backends are available
print("PyTorch Attention Backends:")
print(f"  Flash Attention: {torch.backends.cuda.flash_sdp_enabled() if torch.cuda.is_available() else 'N/A (no CUDA)'}")
print(f"  Memory-efficient: {torch.backends.cuda.mem_efficient_sdp_enabled() if torch.cuda.is_available() else 'N/A (no CUDA)'}")
print(f"  Math (fallback): Always available")

# The beautiful thing: same API, automatic optimization
# PyTorch picks the fastest available backend
PyTorch Attention Backends:
  Flash Attention: N/A (no CUDA)
  Memory-efficient: N/A (no CUDA)
  Math (fallback): Always available
# Benchmark: manual vs PyTorch SDPA
import time

def benchmark(fn, name, warmup=5, runs=20):
    # Warmup
    for _ in range(warmup):
        _ = fn()

    # Timed runs
    start = time.perf_counter()
    for _ in range(runs):
        _ = fn()
    elapsed = (time.perf_counter() - start) / runs * 1000

    print(f"{name}: {elapsed:.2f} ms per call")
    return elapsed

# Benchmark on CPU (GPU would show more dramatic difference)
batch, num_heads, seq_len, head_dim = 1, 8, 256, 64
q = torch.randn(batch, num_heads, seq_len, head_dim)
k = torch.randn(batch, num_heads, seq_len, head_dim)
v = torch.randn(batch, num_heads, seq_len, head_dim)

print(f"\nBenchmark (seq_len={seq_len}, {num_heads} heads, head_dim={head_dim}):")
t_manual = benchmark(lambda: manual_attention(q, k, v, is_causal=True), "Manual attention")
t_sdpa = benchmark(lambda: F.scaled_dot_product_attention(q, k, v, is_causal=True), "PyTorch SDPA")
print(f"\nSpeedup: {t_manual/t_sdpa:.1f}x")

Benchmark (seq_len=256, 8 heads, head_dim=64):
Manual attention: 0.73 ms per call
PyTorch SDPA: 0.27 ms per call

Speedup: 2.7x
TipFrom Scratch to Production
What we learned What PyTorch provides
Q @ K.T / sqrt(d_k) Fused kernel, no intermediate storage
Causal mask with -inf Built-in is_causal=True flag
Stable softmax Numerically stable implementation
Manual loops Flash Attention tiling

Use F.scaled_dot_product_attention in production. The scratch implementation aids debugging, but the optimized version runs 2-10x faster on GPU.

Exercises

Exercise 1: Verify Attention Row Sums

# Verify that each row of attention weights sums to 1
Q = torch.randn(1, 5, 16)
K = torch.randn(1, 5, 16)
V = torch.randn(1, 5, 16)

output, weights = scaled_dot_product_attention(Q, K, V)

print("Attention weights row sums (should all be 1.0):")
print(weights[0].sum(dim=-1))
Attention weights row sums (should all be 1.0):
tensor([1.0000, 1.0000, 1.0000, 1.0000, 1.0000])

Exercise 2: Effect of Temperature

# Temperature scaling affects attention sharpness
# Higher temperature = more uniform, Lower = more peaked

def attention_with_temperature(Q, K, V, temperature=1.0):
    d_k = Q.size(-1)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / (math.sqrt(d_k) * temperature)
    weights = F.softmax(scores, dim=-1)
    output = torch.matmul(weights, V)
    return output, weights

# Generate fixed Q, K, V for temperature comparison
torch.manual_seed(42)
Q_temp = torch.randn(1, 4, 8)
K_temp = torch.randn(1, 4, 8)
V_temp = torch.randn(1, 4, 8)

# Compute attention at different temperatures
temp_weights = {}
for temp in [0.5, 1.0, 2.0]:
    _, w = attention_with_temperature(Q_temp, K_temp, V_temp, temperature=temp)
    temp_weights[f"temp_{str(temp).replace('.', '_')}"] = w[0].detach().numpy().tolist()

ojs_define(
    temp_weights_05 = temp_weights["temp_0_5"],
    temp_weights_10 = temp_weights["temp_1_0"],
    temp_weights_20 = temp_weights["temp_2_0"]
)

print("Lower temperature = sharper attention (more peaked)")
print("Higher temperature = softer attention (more uniform)")

Exercise 3: Compare Single-Head vs Multi-Head

# Single head with full dimension vs multiple heads with smaller dimensions
embed_dim = 64
seq_len = 8

# Single head: one attention over 64 dimensions
single_head = MultiHeadAttention(embed_dim=embed_dim, num_heads=1)

# Multi head: 8 attention heads over 8 dimensions each
multi_head = MultiHeadAttention(embed_dim=embed_dim, num_heads=8)

x = torch.randn(1, seq_len, embed_dim)

out_single, w_single = single_head(x, return_attention=True)
out_multi, w_multi = multi_head(x, return_attention=True)

print(f"Single-head attention:")
print(f"  Attention weights shape: {w_single.shape}")
print(f"  One pattern to rule them all")

print(f"\nMulti-head attention:")
print(f"  Attention weights shape: {w_multi.shape}")
print(f"  8 different patterns, each can specialize")

# Compute average of first 4 heads for visualization
multi_combined = torch.zeros(seq_len, seq_len)
for h in range(4):
    multi_combined += w_multi[0, h].detach()
multi_combined /= 4

# Pass to OJS for visualization
ojs_define(
    single_head_weights = w_single[0, 0].detach().numpy().tolist(),
    multi_head_avg_weights = multi_combined.numpy().tolist(),
    single_multi_seq_len = seq_len
)

Complexity and Optimizations

Attention has:

  • Time complexity: O(n^2 * d) where n = sequence length, d = embedding dimension
  • Memory complexity: O(n^2) for storing the attention matrix

For long sequences (n = 10000), O(n²) complexity demands 100 million attention entries!

KV Cache for Efficient Inference

During autoregressive generation, we compute attention one token at a time. Without caching, we’d recompute K and V for all previous tokens at each step.

KV Cache: Store computed K and V values for previous tokens:

  • At step t, only compute K_t and V_t for the new token
  • Concatenate with cached K_{1:t-1} and V_{1:t-1}
  • Query only needs the new token’s Q_t

This reduces per-token generation from O(n²) to O(n).

Modern Optimizations

Flash Attention (Dao et al., 2022):

  • Avoids materializing the full n x n attention matrix
  • Uses tiling and recomputation to be memory-efficient
  • 2-4x faster than standard attention on modern GPUs
  • PyTorch 2.0+ and most frameworks now default to Flash Attention

Sparse Attention patterns:

  • Local attention: Each token only attends to nearby tokens
  • Strided attention: Attend to every k-th token
  • Block-sparse: Combine local and strided patterns

Linear Attention approximations:

  • Replace softmax(QK^T)V with kernel feature maps
  • Achieves O(n) complexity but may sacrifice quality

Grouped-Query Attention (GQA & MQA)

The KV cache above is what makes generation fast — but it is also what makes generation expensive to remember. At each step we store a key and a value vector for every past token, in every layer, for every head. For a 70B-class model at a few thousand tokens of context, that cache runs to tens of gigabytes and quickly dominates the memory budget, capping how many requests a GPU can serve at once. The bottleneck is not compute — it is the size of the KV cache.

Grouped-Query Attention (GQA) attacks it with one idea: let several query heads share a single key/value head. The queries stay diverse — each head still asks its own question — but they consult a smaller set of keys and values, so there is less to cache.

  • Multi-Head Attention (MHA) — every query head has its own K/V head. Maximum quality, maximum cache.
  • Multi-Query Attention (MQA) — all query heads share one K/V head. The cache shrinks by a factor of num_heads, but quality can dip.
  • Grouped-Query Attention (GQA) — the middle ground: split the query heads into num_kv_heads groups, one shared K/V head per group. Near-MHA quality at a fraction of the cache. This is the setting LLaMA 2/3, Mistral, Qwen, and Gemma actually ship.

Intuition: Share the Keys and Values

Picture the query heads on the left and the key/value heads on the right. In MHA the two columns are the same size and wired one-to-one. GQA keeps every query head but collapses the right column, wiring each group of queries to one shared K/V head. MQA collapses it all the way to a single K/V head. Step through the three variants:

Notice the query column never shrinks — the model keeps all 8 distinct queries. Only the key/value column collapses. That is the whole trick: diversity where it is cheap (queries are recomputed each step and never cached), sharing where it is expensive (keys and values are cached for every past token).

The Math: Counting the KV Cache

The cache stores one key and one value vector per token, per layer, per key/value head. Its size is

\text{KV bytes} = 2 \cdot L \cdot n \cdot n_{kv} \cdot d_\text{head} \cdot b

where L is the number of layers, n the context length, n_{kv} the number of key/value heads, d_\text{head} the head dimension, b the bytes per element (2 for fp16/bf16), and the leading 2 counts K and V.

The quantity we control is n_{kv}. Full MHA sets n_{kv} = n_\text{heads}; MQA sets n_{kv} = 1; GQA picks something in between. So GQA shrinks the cache by exactly

\frac{n_\text{heads}}{n_{kv}}

and it shrinks the K and V projection parameters by the same factor (the query projection and the number of query heads are untouched). Attention still runs over all n_\text{heads} query heads — GQA changes only how many distinct keys and values those heads look at.

NoteKey Insight

MHA gives each of n_\text{heads} query heads its own K/V; GQA gives each group of query heads one shared K/V; MQA gives them all one. Because the KV cache — not the matrix multiply — is the real inference bottleneck, dividing n_{kv} by 8 divides the cache by 8 at almost no quality cost. Llama-2-70B uses 64 query heads and 8 K/V heads: an 8× smaller cache than the equivalent MHA model.

The GroupedQueryAttention you build here is a general attention module. In Module 09: Efficient Attention you’ll build its cache-aware, causal counterpart for decode-time generation — and see FlashAttention attack the compute side of the same problem.

Code: Grouped-Query Attention From Scratch

The implementation is MHA with two changes: K and V project to the smaller num_kv_heads, and we repeat_interleave them back up to num_heads so the same scaled_dot_product_attention runs unchanged. This is GroupedQueryAttention in attention.py.

import torch
from attention import GroupedQueryAttention, kv_cache_bytes, demonstrate_gqa

# 8 query heads sharing 2 key/value heads (a group size of 4).
gqa = GroupedQueryAttention(embed_dim=256, num_heads=8, num_kv_heads=2)
x = torch.randn(4, 32, 256)                     # (batch, seq, embed)
out, attn = gqa(x, return_attention=True)

print(f"Output shape:        {tuple(out.shape)}")
print(f"Attn weights shape:  {tuple(attn.shape)}   (all 8 query heads)")
print(f"Q projection:        {gqa.q_proj.weight.shape[0]} rows (8 heads × 32)")
print(f"K projection:        {gqa.k_proj.weight.shape[0]} rows (2 heads × 32)")
print(f"n_rep (group size):  {gqa.n_rep} query heads per K/V head")
Output shape:        (4, 32, 256)
Attn weights shape:  (4, 8, 32, 32)   (all 8 query heads)
Q projection:        256 rows (8 heads × 32)
K projection:        64 rows (2 heads × 32)
n_rep (group size):  4 query heads per K/V head

The two endpoints fall straight out of num_kv_heads:

mha = GroupedQueryAttention(embed_dim=256, num_heads=8, num_kv_heads=8)  # MHA
mqa = GroupedQueryAttention(embed_dim=256, num_heads=8, num_kv_heads=1)  # MQA

for name, layer in [("MHA", mha), ("GQA", gqa), ("MQA", mqa)]:
    kv_params = layer.k_proj.weight.numel() + layer.v_proj.weight.numel()
    print(f"{name}: {layer.num_kv_heads} K/V head(s), "
          f"{kv_params:>6,} K+V params, n_rep {layer.n_rep}")
MHA: 8 K/V head(s), 131,072 K+V params, n_rep 1
GQA: 2 K/V head(s), 32,768 K+V params, n_rep 4
MQA: 1 K/V head(s), 16,384 K+V params, n_rep 8

The K/V heads are expanded with repeat_interleave, so query head i reads shared head i // n_rep, where n_rep = num_heads // num_kv_heads is the number of query heads sharing each K/V head — the group size. This is the attribute name LLaMA’s repeat_kv uses, and the one Module 09 reuses on its cache-aware build, so the two modules read the same. Because the expansion is exact, MHA (num_kv_heads == num_heads) is a genuine special case — no approximation.

# The payoff, on a realistic decoder (32 layers, 32 heads, d_head=128, 8K ctx):
demonstrate_gqa(embed_dim=4096, num_heads=32, seq_len=8192, num_layers=32)
============================================================
GROUPED-QUERY ATTENTION DEMONSTRATION
============================================================
embed_dim=4096, num_heads=32, head_dim=128
KV cache at seq_len=8192, num_layers=32 (fp16):

   MHA: 32 KV heads →   4.00 GB  (1x smaller than MHA)
   GQA:  8 KV heads →   1.00 GB  (4x smaller than MHA)
   MQA:  1 KV heads →   0.12 GB  (32x smaller than MHA)
{'MHA': {'num_kv_heads': 32, 'kv_cache_bytes': 4294967296},
 'GQA': {'num_kv_heads': 8, 'kv_cache_bytes': 1073741824},
 'MQA': {'num_kv_heads': 1, 'kv_cache_bytes': 134217728}}
Warningnum_kv_heads must divide num_heads

The query heads split into equal groups, so num_heads % num_kv_heads == 0 must hold — GroupedQueryAttention asserts it. 8 query heads can share 1, 2, 4, or 8 K/V heads, but not 3. Production kernels (F.scaled_dot_product_attention with enable_gqa=True, FlashAttention) skip the physical repeat_interleave for speed, but compute the identical result.

How Much Memory Does It Save?

Slide the number of key/value heads and watch the KV cache grow with context length. The dashed line is full MHA (32 K/V heads); the solid line is your GQA choice. Fewer K/V heads → a flatter, cheaper curve — the same context for a fraction of the memory.

TipTry This
  1. Drop to 1 K/V head (MQA) — the curve flattens to 1/32 of MHA. That is the whole cache budget of a long-context request, reclaimed.
  2. Set 8 K/V heads (Llama-2-70B’s choice) — a 4× cut versus MHA while keeping 8 distinct K/V subspaces. Note how it tracks well below the dashed MHA line.
  3. Set 32 K/V heads — the solid line lands exactly on the dashed MHA line: GQA with num_kv_heads == num_heads is MHA.

QK-Norm: Bounding the Attention Logits

Back when we introduced the scaling factor, we divided the scores by \sqrt{d_k} so the dot products would not blow up and saturate the softmax. But look closely at what that argument assumed: it assumed the entries of Q and K have unit variance — true at initialization, when the weights are freshly sampled. Nothing keeps it true during training.

As a model trains, the query and key activations are free to grow. Their norms \lVert q \rVert and \lVert k \rVert drift upward, the logits q \cdot k grow with them, and the softmax sharpens toward a one-hot spike. Attention rows lose almost all their entropy — every query attends to a single key — and the gradient through the softmax goes to zero. Training destabilizes and the loss diverges. This failure mode has a name: attention entropy collapse.

It is not hypothetical. When Google trained a 22-billion-parameter Vision Transformer, the run diverged after a few thousand steps once the model passed ~8B parameters — the attention logits had grown into the range of 10^4, producing near-one-hot attention with near-zero entropy.

Query-Key Normalization (QK-Norm) removes the assumption instead of hoping it holds. Before taking the dot product, it L2-normalizes each per-head query and key vector so they lie on the unit sphere. Now the score is a cosine similarity — always in [-1, 1] — no matter how large the raw activations became. A single learned scale g (a “temperature”) then decides how sharp attention is allowed to get:

\sqrt{d_k} assumes the scores are well-behaved. QK-Norm enforces it.

The Math: normalize, then temper

Standard scaled dot-product attention is \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V. QK-Norm replaces it with

\text{QKNorm}(Q, K, V) = \text{softmax}\!\left(g \cdot \hat{Q}\hat{K}^\top\right)V, \qquad \hat{q} = \frac{q}{\lVert q \rVert_2}.

Because each \hat{q} and \hat{k} is a unit vector, \hat{q}\cdot\hat{k} = \cos\theta \in [-1, 1], so every pre-softmax logit is bounded in [-g,\, g] — for the entire lifetime of training, regardless of the activation norms. The learned scale g is initialized following Henry et al. to g_0 = \log_2(L^2 - L) for a reference sequence length L, large enough that a uniform row can still be sharpened to near one-hot when the model wants it.

Walk the pipeline one step at a time — watch the logit range collapse from “unbounded” to a hard interval:

NoteKey Insight

The \sqrt{d_k} factor is a fixed constant chosen for the init-time statistics; QK-Norm’s g is a learned constant that stays valid because the normalization guarantees the scores it multiplies never leave [-1, 1]. One is a hope, the other a bound.

Code: QK-Norm from scratch

The whole mechanism is l2_normalize on Q and K, then a learned scale — a few lines on top of the attention you already built. This follows the implementation in attention.py (qk_norm_attention and the QKNormAttention layer).

import torch
import torch.nn.functional as F
from attention import l2_normalize, qk_norm_attention, scaled_dot_product_attention

torch.manual_seed(0)
# One head's worth of query/key/value: (batch, heads, seq, d_k)
q = torch.randn(1, 1, 6, 32)
k = torch.randn(1, 1, 6, 32)
v = torch.randn(1, 1, 6, 32)

# Unit-length queries and keys → the score is a cosine similarity in [-1, 1]
cosine = l2_normalize(q) @ l2_normalize(k).transpose(-2, -1)
print(f"Cosine score range: [{cosine.min():.3f}, {cosine.max():.3f}]  (always ⊆ [-1, 1])")

g = 10.0  # the learned scale (fixed here for illustration)
out, weights = qk_norm_attention(q, k, v, scale=g)
print(f"Pre-softmax logits live in [-{g}, {g}] regardless of activation size")
print(f"Output shape: {tuple(out.shape)}, weights sum: {weights.sum(-1)[0,0].tolist()[:3]}")
Cosine score range: [-0.344, 0.328]  (always ⊆ [-1, 1])
Pre-softmax logits live in [-10.0, 10.0] regardless of activation size
Output shape: (1, 1, 6, 32), weights sum: [1.0, 1.0000001192092896, 0.9999999403953552]

The headline property is scale invariance. Blow up the query activations by 100× — as unconstrained training eventually would — and QK-Norm’s attention weights do not move at all, because only the direction of q survives normalization. Plain attention, by contrast, swings wildly:

# Same q, k, v — but scale the queries up 100x (simulating activation drift)
_, w_qk_small = qk_norm_attention(q, k, v, scale=g)
_, w_qk_big   = qk_norm_attention(q * 100.0, k, v, scale=g)
print(f"QK-Norm   max |Δweight| after 100x rescale: {(w_qk_small - w_qk_big).abs().max():.2e}")

_, w_plain_small = scaled_dot_product_attention(q, k, v)
_, w_plain_big   = scaled_dot_product_attention(q * 100.0, k, v)
print(f"Plain     max |Δweight| after 100x rescale: {(w_plain_small - w_plain_big).abs().max():.2e}")
QK-Norm   max |Δweight| after 100x rescale: 1.19e-07
Plain     max |Δweight| after 100x rescale: 7.21e-01

Now let the activation scale sweep across a whole range and record what happens to the largest logit and to the attention entropy, for both schemes (qk_norm_logit_growth in attention.py computes this):

from attention import qk_norm_logit_growth

growth = qk_norm_logit_growth(
    scale_factors=(1.0, 2.0, 4.0, 8.0, 16.0, 32.0, 64.0),
    seq_len=32, d_k=64, g=12.0,
)
ojs_define(qkn_growth = growth)

Watch it collapse — or not

Drive the two knobs yourself. Activation growth simulates the query/key norms drifting up during training; learned scale g is QK-Norm’s temperature. Watch the plain-attention logits run away and its entropy collapse toward zero, while QK-Norm holds its logits at g and keeps a healthy, non-degenerate attention distribution.

TipTry This
  1. Crank activation growth to 64×. Plain attention’s max logit rockets into the hundreds and its entropy bar drains toward empty — that is entropy collapse, the thing that diverged ViT-22B. QK-Norm’s readout does not budge.
  2. Now shrink g toward 1. QK-Norm’s entropy bar fills up: with a tiny scale, cosine scores in [-1,1] can’t sharpen, so attention goes nearly uniform. This is exactly why the learned scale exists — and why initializing it to \log_2(L^2-L) matters.
  3. Set g around 12–15. You get a healthy middle: bounded logits and an attention distribution sharp enough to be useful.

Two things called “QK-norm”

There are two normalizations that both go by this name, and conflating them is the most common mistake:

Variant What it does Bound Used by
ℓ2 / cosine (built here) Normalize q,k to unit length; score = g\cos\theta Hard [-g, g] Henry et al. 2020
QK-LayerNorm / RMSNorm Apply LayerNorm (or RMSNorm) to q,k before the dot product Controls scale, not a hard interval ViT-22B, OLMo 2, Gemma 3

The cosine version gives the clean, provable bound we just built, so it’s the best one to learn from. The variant that ships in most production LLMs is the LayerNorm/RMSNorm one — it reuses the normalization layer you’ll build in Module 06 and controls the logit scale without forcing the exact [-1,1] cosine geometry. Both cure the same disease: attention-logit growth.

WarningPitfalls
  • Dropping the learned scale. Cosine scores in [-1,1] make a very flat softmax. Without g, attention can never sharpen — QK-Norm needs its temperature.
  • Normalizing the wrong axis. Normalize each vector over the head dimension (d_k), per query and per key — not across the sequence.
  • Normalizing V too. Only Q and K are normalized; the values carry the content and are left untouched.
  • Assuming it’s free. It adds two normalization ops in the attention hot path. Cheap, but not zero — worth it precisely when stability is at risk (very large or long-training models).

Exercise: measure the collapse

Fill in the two lines that compute the mean attention entropy for plain vs QK-normalized attention as the activations grow. You should find plain entropy falling toward 0 while QK-Norm’s stays flat.

import torch
import torch.nn.functional as F
import math
from attention import l2_normalize

torch.manual_seed(0)
q = torch.randn(16, 64)
k = torch.randn(16, 64)

def mean_entropy(weights):
    p = weights.clamp_min(1e-12)
    return float((-(p * p.log2()).sum(-1)).mean())

for f in (1.0, 8.0, 64.0):
    qf, kf = q * f, k * f
    # Your implementation here:
    # plain_w = F.softmax( ... / math.sqrt(64), dim=-1)
    # qk_w    = F.softmax( l2_normalize(qf) @ ... * 12.0, dim=-1)
    plain_w = F.softmax(qf @ kf.T / math.sqrt(64), dim=-1)
    qk_w = F.softmax(l2_normalize(qf) @ l2_normalize(kf).T * 12.0, dim=-1)
    print(f"scale {f:5.0f}x | plain entropy {mean_entropy(plain_w):.3f} | "
          f"QK-Norm entropy {mean_entropy(qk_w):.3f} bits")
scale     1x | plain entropy 3.388 | QK-Norm entropy 2.961 bits
scale     8x | plain entropy 0.062 | QK-Norm entropy 2.961 bits
scale    64x | plain entropy 0.000 | QK-Norm entropy 2.961 bits
NoteSee also: the weight-space twin

QK-Norm bounds the logits by changing the forward pass — always on. There is a second cure for the same explosion that leaves the forward pass untouched and instead clips the weights after each optimizer step: QK-Clip, the stability half of the MuonClip optimizer behind Kimi K2. When the logit explosion is driven by the optimizer (Muon’s larger updates), that is often the better fit — see Module 7: Muon at Scale.

Differential Attention: Cancelling the Noise

QK-Norm fixed how big the logits get. There is a second softmax pathology it leaves untouched: attention noise. A softmax row is forced to sum to one, so every query must spend some of its attention budget on tokens it does not actually care about — a diffuse, content-independent floor smeared across the whole sequence. In the extreme, that mass piles onto the first token or two — the attention sink you’ll meet again in Module 10. Either way it is noise: budget the query cannot help spending, drowning out the signal.

You cannot remove it by making one softmax sharper — that just trades a broad floor for a spiky one. But you can cancel it. Think of noise-cancelling headphones: they record the ambient hum with a second microphone and subtract it. Differential attention does the same to attention. It computes two softmax maps and outputs their difference:

\text{DiffAttn}(X) = \Big(\underbrace{\text{softmax}\!\big(\tfrac{Q_1 K_1^\top}{\sqrt{d_k}}\big)}_{\text{signal + noise}} - \lambda\,\underbrace{\text{softmax}\!\big(\tfrac{Q_2 K_2^\top}{\sqrt{d_k}}\big)}_{\text{noise}}\Big)\,V.

Both maps carry the same common-mode floor. Subtract a \lambda-scaled copy of the second from the first and that shared floor cancels, leaving a sparser, higher signal-to-noise map focused on the tokens that matter. The values V are untouched — only the weighting changes.

QK-Norm bounds the softmax. Differential attention denoises it.

The Math: two maps, one difference

Each head projects the input to twice the usual width and splits it in half. With X \in \mathbb{R}^{N \times d_{\text{model}}} and head dimension d_k:

[Q_1; Q_2] = X W^Q, \qquad [K_1; K_2] = X W^K, \qquad V = X W^V,

so Q_1, Q_2, K_1, K_2 \in \mathbb{R}^{N \times d_k} are the two halves of the query and key. The two attention maps A_1 = \text{softmax}(Q_1 K_1^\top/\sqrt{d_k}) and A_2 = \text{softmax}(Q_2 K_2^\top/\sqrt{d_k}) are ordinary distributions — each row sums to one. Their difference A_1 - \lambda A_2 is not a distribution (its rows sum to 1-\lambda), and that is the point: a component of attention present in both maps is suppressed in proportion to \lambda, while a target only the first map sees survives.

Two limits make the mechanism concrete, and both are exact:

  • \lambda = 0 — the subtraction vanishes and you are back to ordinary attention on (Q_1, K_1, V).
  • identical maps (Q_1{=}Q_2,\ K_1{=}K_2) — the two maps cancel perfectly: the output is exactly (1-\lambda) times ordinary attention. Pure common mode, nothing left but a scaled echo.

Walk the subtraction one step at a time:

NoteKey Insight

Sharpening a single softmax cannot remove the floor — a peaked distribution still sums to one and still leaks mass onto irrelevant tokens. Only a difference of two softmaxes can push a weight below zero and truly cancel the common noise. That is why one map is not enough.

The learnable λ, and why it grows with depth

\lambda is a scalar, but the paper does not learn it directly. It learns four vectors \lambda_{q_1}, \lambda_{k_1}, \lambda_{q_2}, \lambda_{k_2} \in \mathbb{R}^{d} and reparameterizes

\lambda = \exp(\lambda_{q_1}\!\cdot\lambda_{k_1}) - \exp(\lambda_{q_2}\!\cdot\lambda_{k_2}) + \lambda_{\text{init}}.

The two \exp terms keep the gradient of \lambda smooth, and at initialization (small vectors ⇒ dot products \approx 0 ⇒ both \exp \approx 1) they cancel, so \lambda \approx \lambda_{\text{init}}. The base value follows a depth schedule:

\lambda_{\text{init}} = 0.8 - 0.6\,\exp\!\big(-0.3\,(l-1)\big), \qquad l \in [1, L].

Shallow layers (l=1) start at \lambda_{\text{init}} = 0.2 and barely subtract; deep layers climb toward 0.8 and cancel aggressively, as their representations sharpen and the noise floor becomes the thing worth removing.

One more piece keeps the output well-scaled: each head is normalized with a headwise GroupNorm (an RMSNorm applied per head) and then multiplied by the fixed (1-\lambda_{\text{init}}), restoring the variance the subtraction removed. And the whole layer is free: by halving the head count and doubling each head’s width (h = d_{\text{model}}/2d), the Q/K/V/O projections stay d_{\text{model}} \times d_{\text{model}}, so DIFF attention has the same parameters and FLOPs as the multi-head attention you already built — plus four tiny \lambda vectors.

Code: differential attention from scratch

The core is three lines on top of the attention you already have — two softmaxes and a subtraction. This follows diff_attention.py (differential_attention, compute_lambda, lambda_init_schedule, and the DifferentialAttention layer).

import torch
import torch.nn.functional as F
from diff_attention import differential_attention, lambda_init_schedule

torch.manual_seed(0)
# Two independent (query, key) pairs per head, plus values: (batch, seq, d)
q1, k1 = torch.randn(1, 6, 16), torch.randn(1, 6, 16)
q2, k2 = torch.randn(1, 6, 16), torch.randn(1, 6, 16)
v = torch.randn(1, 6, 16)

out, diff_w = differential_attention(q1, k1, q2, k2, v, lam=0.8)
print(f"Output shape: {tuple(out.shape)}")
print(f"Differential-weight row sums (= 1 − λ = 0.2): {diff_w.sum(-1)[0].tolist()[:3]}")
Output shape: (1, 6, 16)
Differential-weight row sums (= 1 − λ = 0.2): [0.19999998807907104, 0.19999995827674866, 0.19999997317790985]

The two structural identities are worth seeing directly — they are exact:

import math

def ordinary(q, k, val):
    scores = (q @ k.transpose(-2, -1)) / math.sqrt(q.shape[-1])
    return F.softmax(scores, dim=-1) @ val

# λ = 0  →  ordinary attention on (q1, k1, v)
out0, _ = differential_attention(q1, k1, q2, k2, v, lam=0.0)
print(f"λ=0        max|DiffAttn − ordinary|       = {(out0 - ordinary(q1, k1, v)).abs().max():.2e}")

# identical maps  →  exactly (1 − λ) × ordinary attention
lam = 0.7
out_same, _ = differential_attention(q1, k1, q1, k1, v, lam=lam)
print(f"same maps  max|DiffAttn − (1−λ)·ordinary| = {(out_same - (1 - lam) * ordinary(q1, k1, v)).abs().max():.2e}")
λ=0        max|DiffAttn − ordinary|       = 0.00e+00
same maps  max|DiffAttn − (1−λ)·ordinary| = 5.96e-08

Now the payoff. differential_noise_demo plants a signal on one key over a shared sink floor and measures the signal-to-noise ratio of the ordinary map (map A alone) versus the denoised differential map:

from diff_attention import differential_noise_demo

demo = differential_noise_demo()
o, dd = demo["ordinary"], demo["differential"]
print(f"Ordinary map      signal={o['signal']:.3f}  noise={o['noise']:.3f}  SNR={o['snr']:.2f}")
print(f"Differential map  signal={dd['signal']:.3f}  noise={dd['noise']:.3f}  SNR={dd['snr']:.2f}")
print(f"→ {dd['snr'] / o['snr']:.0f}x better SNR — the sink floor is gone")
Ordinary map      signal=0.199  noise=0.801  SNR=0.25
Differential map  signal=0.173  noise=0.027  SNR=6.35
→ 26x better SNR — the sink floor is gone

The full multi-head layer, parameter-matched to CausalMultiHeadAttention, lives in DifferentialAttention. Note it takes a layer_idx — the depth that sets \lambda_{\text{init}}:

from diff_attention import DifferentialAttention

attn = DifferentialAttention(embed_dim=256, num_heads=4, layer_idx=6)
x = torch.randn(2, 32, 256)                      # (batch, seq, embed)
y = attn(x)
print(f"Head dim d = embed/(2·heads) = {attn.head_dim}   (heads are half as many, twice as wide)")
print(f"λ_init at layer 6: {attn.lambda_init:.3f}   current λ ≈ {float(attn.current_lambda().detach()):.3f}")
print(f"Output shape: {tuple(y.shape)}")
Head dim d = embed/(2·heads) = 32   (heads are half as many, twice as wide)
λ_init at layer 6: 0.666   current λ ≈ 0.745
Output shape: (2, 32, 256)

Drive the cancellation

The bridge below hands the two real softmax maps from differential_noise_demo to the widget. Map A carries the target on key 7 plus the sink floor on keys 0–1; Map B is the floor alone. Slide \lambda from 0 (ordinary attention, noisy) upward and watch the sink columns drain out of the difference while the target holds — the signal-to-noise readout climbing as they go.

from diff_attention import differential_noise_demo

viz = differential_noise_demo(seq_len=12, target=7)
ojs_define(
    diff_mapA = viz["ordinary"]["weights"],       # softmax(Q1 K1ᵀ): signal + noise
    diff_mapB = viz["differential"]["map2"],       # softmax(Q2 K2ᵀ): the noise floor
    diff_target = viz["target"],
)
TipTry This
  1. Start at λ = 0. Map A and the difference are identical — plain, noisy attention. Keys 0–1 hold a big chunk of the weight the query never wanted.
  2. Raise λ toward 0.8. The sink columns fade out of the difference and the differential SNR readout rockets past the ordinary one. That is the noise cancelling.
  3. Push λ past ~1. Over-subtract and the target itself starts to go negative — the difference stops being a sensible weighting. This is why \lambda is learned and bounded by the schedule, not cranked to the maximum.

Deeper layers subtract more

The \lambda_{\text{init}} schedule is a gentle curve from 0.2 to 0.8. Early layers keep most of both maps (little cancellation); later layers, whose features are sharp enough that the residual is mostly noise, subtract hard.

from diff_attention import lambda_init_schedule

sched = [{"layer": l, "lambda_init": lambda_init_schedule(l)} for l in range(1, 25)]
ojs_define(diff_schedule = sched)
WarningPitfalls
  • Using one λ for the whole model. \lambda_{\text{init}} is per layer and follows the depth schedule; a single global value throws away the shallow-keeps, deep-cancels structure.
  • Forgetting the (1-\lambda_{\text{init}}) scale. The subtraction removes variance; without the headwise GroupNorm and the fixed rescale the residual stream drifts and training destabilizes.
  • Comparing head counts, not parameters. DIFF attention halves the heads and doubles their width. Match it to a standard block by parameter count (d_{\text{model}} \times d_{\text{model}} projections), not by head count.
  • Reading the difference as a distribution. Its rows sum to 1-\lambda and can go negative — it is a denoised weighting, not a probability.

Exercise: measure the denoising

Fill in the differential weighting and confirm the noise mass on the sink keys collapses as \lambda rises, while the target weight barely moves.

import torch
import torch.nn.functional as F

# Map A: sink floor on keys 0-1, signal on key 7.  Map B: floor only.
floor = torch.tensor([2.5, 2.5, 0, 0, 0, 0, 0, 0.0])
A = F.softmax(floor.clone().index_add(0, torch.tensor([7]), torch.tensor([2.0])), dim=-1)
B = F.softmax(floor, dim=-1)

for lam in (0.0, 0.4, 0.8):
    # Your implementation here:
    # diff = A - lam * B
    diff = A - lam * B
    sink_mass = diff[:2].abs().sum().item()
    print(f"λ={lam:.1f} | sink noise mass {sink_mass:.3f} | target weight {diff[7]:.3f}")
λ=0.0 | sink noise mass 0.663 | target weight 0.201
λ=0.4 | sink noise mass 0.342 | target weight 0.188
λ=0.8 | sink noise mass 0.021 | target weight 0.175

Interactive Exploration

Experiment with attention in real-time. Adjust the temperature to see how it affects the attention distribution:

  • Low temperature → Sharp, focused attention (nearly one-hot)
  • High temperature → Soft, diffuse attention (more uniform)
TipTry This
  1. Set temperature to 0.1 — notice how attention becomes nearly one-hot (picks one token)
  2. Set temperature to 3.0 — notice how attention becomes almost uniform
  3. Compare how “sat” attends (looks at “cat”) vs how “the” attends (looks at other “the”)

Common Pitfalls

When implementing attention, watch out for these issues:

  1. Forgetting to scale: Without /sqrt(d_k), training becomes unstable with large head dimensions
  2. Wrong mask dimensions: Mask should broadcast correctly over batch and head dimensions
  3. NaN from all-masked rows: If an entire row is masked, softmax produces NaN (log(0)). Handle with nan_to_num or ensure at least one position is unmasked
  4. Memory leaks with attention weights: Storing attention weights for visualization can exhaust memory. Only compute when needed

Summary

Key takeaways:

  1. Attention computes weighted sums: Each position’s output is a weighted combination of all (allowed) positions’ values
  2. Q, K, V: Query asks “what do I need?”, Key says “what do I have?”, Value carries the information
  3. Scaling prevents gradient issues: Dividing by sqrt(d_k) keeps softmax from saturating
  4. Causal masking enables generation: In LLMs, we mask future tokens so the model learns to predict the next token
  5. Multiple heads learn different patterns: Each head can specialize in different linguistic relationships
  6. Complexity is O(n^2): Attention’s quadratic cost limits sequence length, motivating optimizations like Flash Attention and KV caching
  7. Grouped-Query Attention shrinks the KV cache: Sharing key/value heads across query heads (GQA, or MQA at the extreme) cuts the cache by num_heads / num_kv_heads at near-MHA quality — the setting every modern open model ships
  8. QK-Norm bounds the logits: 1/sqrt(d_k) only assumes well-behaved scores; L2-normalizing Q and K makes each logit a cosine similarity times a learned scale g, so it stays in [-g, g] throughout training and prevents attention entropy collapse
  9. Differential attention cancels the noise: a softmax must sum to one, so it always leaks mass onto irrelevant tokens; subtracting a λ-scaled second softmax map (softmax(Q1K1ᵀ/√d_k) − λ·softmax(Q2K2ᵀ/√d_k)) cancels that common-mode floor like noise-cancelling headphones, at the same parameter and FLOP cost as ordinary MHA

Going Deeper

Core Papers:

Practical Resources:

What’s Next

Module 06: Transformer combines attention with feed-forward networks, layer normalization, and residual connections to build a complete transformer decoder block.