Self-Attention

The Problem: Words Are Islands

After the embedding layer, each word in a sentence becomes a vector. But that vector is the same regardless of context. “Bank” gets the same vector whether the sentence is about a river or about money. “It” gets the same vector whether it refers to a cat or a mat.

This is a real problem. Suppose the model is trying to predict the next word after “The cat sat on the mat because it was ___.” To predict well, it needs to know what “it” refers to. But “it” is just a fixed vector that means “generic pronoun.” The model has no way to figure out which earlier words matter.

Seeing the Problem

Play 3 quick rounds below: watch how changing a single word at the end of the sentence completely changes which noun "it" must attend to:

Same pronoun, same position, completely different referent. The model needs a mechanism that lets each word look at the other words and update its representation based on what it finds relevant. That mechanism is called self-attention.

Why “self”? Because the words are attending to other words in the same sentence. Later we’ll see cross-attention, where tokens from one sequence (like a French translation) attend to a different sequence (the English original). Self-attention is the simpler, more fundamental case.

What Is “Attention”?

Step through the animation below. It shows the three things attention does for a single word:

That’s all attention is: look at every word, score relevance, blend the winners. Each word’s representation stops being about that word in isolation and starts carrying context from the words it paid attention to.

The Math, in Three Small Steps

Take the sentence “I love cats.” Each word starts as a vector (from the embedding layer). Right now each vector only knows about its own word. “cats” has no idea that someone loves it. We want to produce a new vector for “cats” that blends in information from the relevant words. Three operations get us there:

That’s the entire recipe: dot product to score, softmax to normalize, weighted sum to blend. But we used the raw embedding vectors directly. The model needs a way to learn which aspects of each vector to compare. That’s what the Q, K, and V projections do.

The Q, K, V Framework

Each token gets three different learned roles: a question it asks, a description it advertises, and information it carries. The analogy that makes this click: a library search.

Why Three Separate Projections?

Why not just use the raw embeddings for everything? Toggle between the two modes to see what goes wrong:

But how do the matrices “know” they’re Q, K, V? They don’t. $W^Q$, $W^K$, and $W^V$ all start as random numbers. What makes them different is where they sit in the formula: $W^Q$‘s output is the left side of the dot product (the one asking), $W^K$‘s output is the right side (the one being compared against), and $W^V$‘s output feeds into the weighted sum (the one providing information). Because each matrix occupies a different position, backpropagation sends each one different gradients, and they specialize over training. The “query, key, value” labels describe the architectural role, not something the matrices inherently possess.

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:

How It Works, Step by Step

Here’s what self-attention does in plain language (matching the steps above):

Why does comparing queries and keys measure relevance? Each token’s query encodes what it’s looking for (“something alive that can feel tired”), and each token’s key encodes what it offers (“I’m alive, I can feel things”). When a query and key point in the same direction, the match score is high. When they’re unrelated, the score is low or negative.

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

What Did We Just Build?

Step through the animation below to see how self-attention fits into the bigger picture: from static embeddings to context-aware representations, stacked layers, and finally next-word prediction.

Self-attention is one building block. Multi-Head Attention runs multiple self-attention operations in parallel for different perspectives. The full transformer architecture shows how all the pieces fit together.

Why Scale the Scores?

You may have noticed the ”÷ √d_k” in step 4 of the visualization above. Each attention score is a dot product, a sum of many small numbers (one per dimension). The more dimensions, the bigger the total. With 64 dimensions, scores can easily reach values like 20 or 30.

Why is that a problem?

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

When softmax saturates, the model puts nearly all its attention on a single token and ignores everything else. Worse, the gradients become near-zero, so the network can’t learn to fix this. Dividing by the square root of the dimension keeps scores in a reasonable range where softmax can spread attention across multiple tokens.

Interactive: Explore Attention Patterns

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

Things to try:

• 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: each token can only attend to itself and the tokens before it. It can’t see the future.

Think of it like writing an exam. You can only use information you’ve already seen, not answers that haven’t been revealed yet.

The Mask

We block future positions by setting their scores to an extremely negative number before softmax, so they get zero attention weight:

Each row shows what a token can attend to. The lower-left triangle is visible; the upper-right is blocked.

Notice how the upper-right triangle of the attention matrix is zeroed out:

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

In self-attention, Q, K, and V all come from the same sequence. But there’s a variant called cross-attention where Q comes from one sequence and K, V come from another. This is how translation models let the decoder “read” the encoder’s output.

Self-attention is the fundamental building block. We’ll explore cross-attention in detail later in Seq2Seq Attention →, where we trace the historical origin of the attention idea.

The Cost of Full Context

Self-attention’s power comes from comparing every token to every other token. But that power has a cost. Drag the slider and watch the attention matrix grow:

What to look for: Start at 4 tokens and slowly drag to 64. The grid area grows much faster than the slider moves. At 4 tokens, there are 16 comparisons. At 8 tokens, there are 64 (not 32). At 64 tokens, there are 4,096. Every time you double the sequence length, the number of comparisons quadruples.

Think of it like seating at a round table. If 4 people need to introduce themselves to everyone else, that’s manageable. But double the guests to 8, and the introductions don’t just double, they quadruple. With 64 guests, you’d need 4,096 introductions. That’s the $n^2$ pattern: the work grows as the square of the input size.

For real models, this gets serious. A 128,000-token context window means over 16 billion entries in the attention matrix, per head, per layer. This is why researchers developed techniques like Flash Attention (exact attention with better memory access patterns), sliding window attention (attend only to nearby tokens), and sparse attention (skip most of the matrix). The quadratic nature is fundamental to the mechanism, but clever engineering works around the worst of the cost.

Quick Summary

Next: Multi-Head Attention

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