Self-Attention

The Problem: What Does “It” Mean?

Read this sentence:

“The cat sat on the mat because it was tired.”

What does “it” refer to? Obviously the cat — not the mat. You know this because you understand the meaning of “tired” and can connect it back to a living thing.

Now consider:

“The cat sat on the mat because it was dirty.”

Now “it” refers to the mat! Same sentence structure, different meaning based on context.

For a neural network to understand language, it needs a mechanism to figure out which other words are relevant to each word. That mechanism is self-attention.

The Key Insight

Self-attention lets every token in a sequence “look at” every other token and decide:

• “Who should I pay attention to?”
• “How much should I pay attention to them?”
• “What information should I take from them?”

The Q, K, V Framework

Self-attention uses three concepts borrowed from information retrieval:

The Library Analogy

Imagine you’re in a library:

1. Query (Q): Your question — “I need information about sleeping animals”
2. Keys (K): Each book’s title/tags — “Cat Behavior”, “Mat Materials”, “Animal Sleep”
3. Values (V): Each book’s actual content — the detailed information inside

You match your query against all keys to find the most relevant books, then read the content (values) of the matching ones.

Self-attention works the same way:

• Each token generates a Query: “What am I looking for?”
• Each token generates a Key: “What information do I have?”
• Each token generates a Value: “What will I contribute if attended to?”

Why Three Separate Projections?

Why not just use the raw embeddings?

Without projections, the dot product of a vector with itself is always the largest — every token would attend mostly to itself. Separate Q, K, V projections let the model learn:

• What aspects to query for (Q)
• What aspects to advertise (K)
• What information to provide (V)

These can be completely different! A word might query for “what noun am I modifying?” (Q) while advertising “I’m an adjective” (K) and providing “here’s my semantic meaning” (V).

Step-by-Step: The Math with Real Numbers

Let’s trace through self-attention with actual numbers. Click “Next” to advance through each computation step:

The Formula

The complete self-attention operation in one line:

$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$

Let’s break this down:

def scaled_dot_product_attention(Q, K, V):
    """
    Args:
        Q: Queries (seq_len, d_k)
        K: Keys    (seq_len, d_k)
        V: Values  (seq_len, d_v)
    Returns:
        output: (seq_len, d_v)
        weights: (seq_len, seq_len) attention matrix
    """
    d_k = Q.shape[-1]
    
    # Step 1: Compute dot product scores
    scores = Q @ K.T                    # (seq_len, seq_len)
    
    # Step 2: Scale
    scores = scores / np.sqrt(d_k)
    
    # Step 3: Softmax each row
    weights = softmax(scores, axis=-1)  # Each row sums to 1
    
    # Step 4: Weighted sum of values
    output = weights @ V                # (seq_len, d_v)
    
    return output, weights

Why √d_k Scaling?

This is one of the most common “why?” questions about attention. Here’s the intuition:

The Problem

When $d_k$ is large, dot products have large variance:

$Q \cdot K = \sum_{i=1}^{d_k} q_i k_i$

If each $q_i, k_i \sim N(0, 1)$, then $Q \cdot K$ has variance $d_k$.

For $d_k = 64$: standard deviation $= \sqrt{64} = 8$. Scores can easily be ±20 or more.

The Consequence

Large scores push softmax into saturation — outputs become nearly one-hot:

Saturated softmax → vanishing gradients → network can’t learn. Scaling by $\sqrt{d_k}$ keeps the variance at 1, and softmax stays in its healthy range.

Interactive: Explore Attention Patterns

Play with the attention visualizer below. Try different sentences and see how tokens attend to each other:

Things to try:

• Toggle Causal Mask to see how decoder-style attention works (each token can only see previous tokens)
• Increase temperature to see how it softens the attention distribution
• Click a token to see its complete attention pattern
• Try sentences with ambiguous words like “bank” to see context effects

Causal (Masked) Self-Attention

For language generation (GPT, LLaMA), we need a critical constraint: a token at position $i$ should only attend to positions $0, 1, \ldots, i$. It can’t see the future!

The Mask

We add $-\infty$ to blocked positions before softmax, so they get zero attention weight:

Each row shows what a token can see: ✅ = visible, ❌ = blocked (future).

When to Mask

def causal_mask(seq_len):
    # Upper triangle = blocked (-infinity so softmax gives 0)
    mask = np.triu(np.ones((seq_len, seq_len)), k=1)
    return np.where(mask == 1, -np.inf, 0)

Self-Attention vs. Cross-Attention

Self-attention is the fundamental building block. Cross-attention is an extension used only in encoder-decoder architectures.

Key Takeaways

1. Self-attention lets every token gather information from every other token
2. Q, K, V projections let the model learn different roles for finding and providing information
3. √d_k scaling prevents softmax saturation and gradient vanishing
4. Causal masking prevents tokens from seeing the future (essential for generation)
5. Self-attention is O(n²) in sequence length — this is both its power and its limitation

Next: Multi-Head Attention

A single attention head gives one perspective. But language is multi-faceted — we need multiple perspectives simultaneously. That’s Multi-Head Attention →