Attention Is All You Need - But What Is Attention?

Why Attention Was Invented

Before transformers, NLP relied on RNNs (recurrent neural networks) to process sequences. RNNs read text one token at a time, maintaining a hidden state that carries information forward. The problem: by the time the model reaches token 500, information from token 1 has been compressed through 499 state transitions. Long-range dependencies - "the animal in the field that the farmer mentioned earlier" - were hard to capture.

Attention solves this by allowing every token to directly attend to every other token, regardless of distance. No information bottleneck. No sequential dependency.

The Query-Key-Value Framework

Attention is computed using three matrices: Query (Q), Key (K), and Value (V). These are learned linear projections of the input.

The intuition:

• Query: "What am I looking for?" - what this position needs from others
• Key: "What do I contain?" - what information this position is advertising
• Value: "What will I contribute?" - the actual content, if selected

For each position, attention computes how well its query matches every key (a dot product), normalizes the scores into a probability distribution (softmax), and returns the weighted sum of all values.

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

def attention(Q, K, V, mask=None):
    d_k = Q.shape[-1]  # Dimension of keys/queries

    # Compute attention scores: how well each query matches each key
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)

    # Apply causal mask for decoder (prevents attending to future positions)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))

    # Softmax to normalize scores into a probability distribution
    attention_weights = F.softmax(scores, dim=-1)

    # Weighted sum of values
    output = torch.matmul(attention_weights, V)
    return output, attention_weights

The Scaling Factor: Why Divide by √d_k?

Without scaling, when d_k (the key dimension) is large, the dot products can become very large, pushing the softmax into regions with very small gradients (the distribution becomes nearly one-hot). Dividing by √d_k keeps the dot products at a reasonable scale regardless of the dimension.

Multi-Head Attention: Learning Multiple Relationship Types

A single attention head learns one type of relationship. Multi-head attention runs several attention computations in parallel, each with its own Q, K, V projections.

pythonCopy
class MultiHeadAttention(torch.nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.n_heads = n_heads
        self.d_k = d_model // n_heads

        self.W_q = torch.nn.Linear(d_model, d_model)
        self.W_k = torch.nn.Linear(d_model, d_model)
        self.W_v = torch.nn.Linear(d_model, d_model)
        self.W_o = torch.nn.Linear(d_model, d_model)

    def forward(self, x, mask=None):
        batch, seq_len, d_model = x.shape

        Q = self.W_q(x).view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        K = self.W_k(x).view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        V = self.W_v(x).view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2)

        out, weights = attention(Q, K, V, mask)

        # Concatenate heads and project
        out = out.transpose(1, 2).contiguous().view(batch, seq_len, d_model)
        return self.W_o(out)

Different heads learn to attend to different things: one head might capture syntactic relationships (subject-verb agreement), another might capture coreference (pronoun to noun), another might capture semantic similarity.

Causal (Masked) Attention vs. Bidirectional Attention

Causal attention (GPT, Claude, Llama): Each token can only attend to previous tokens. Required for autoregressive text generation - the model cannot "see" future tokens when generating.

Bidirectional attention (BERT): Each token attends to all other tokens in both directions. Better for understanding tasks (classification, named entity recognition) but cannot generate text autoregressively.

The choice determines the architecture's appropriate use cases.

What Attention Weights Tell You

The attention weights for each head show which tokens each position attends to most strongly. These are sometimes visualized to interpret model behavior - though interpretability from attention weights is limited and often misleading.

More useful: understanding that attention enables:

• Coreference resolution: "it" attending to the noun it refers to
• Syntactic agreement: verb attending to its subject
• Long-range dependencies: the last word of a sentence attending to the first word that constrains it

This mechanism - direct token-to-token connections, unbounded by distance - is the structural reason transformers outperformed RNNs on virtually every language task once scaled up.

Where to Go Next

The attention mechanism is the core of Module 4 in the non-programmer path (LLMs Under the Hood) at a conceptual level, and Modules 19 and 20 of the College Student path at full technical depth. Once attention makes sense, context windows, prompt design, and fine-tuning choices become much easier to reason about.

Common Mistakes

Forgetting the sqrt(d_k) scaling. The dot product between queries and keys grows in magnitude as the key dimension d_k increases. Without dividing by sqrt(d_k), the softmax receives very large inputs, saturates toward one-hot distributions, and gradients vanish during training. Always include the scaling factor and understand that its sole job is keeping the softmax in a numerically stable operating range.

Confusing encoder self-attention with decoder masked self-attention. In encoder self-attention every token attends to every other token in both directions. In decoder self-attention a causal mask zeros out all future positions so the model cannot "see ahead" during autoregressive generation. Mixing up the two leads to data-leakage bugs in custom transformer implementations that are invisible until you inspect attention weight matrices.

Thinking attention "selects" a single token. Attention computes a weighted sum over all value vectors; it blends information from every position in proportion to query-key similarity. Describing it as selection encourages wrong mental models when you need to reason about information flow or debug attention patterns. Think of it as a soft, differentiable lookup table, not a pointer.

What to Practice Next

• Implement scaled dot-product attention in NumPy from scratch - compute Q, K, V projections manually, apply the scaling and softmax, and verify output shapes match expected dimensions.
• Load a pretrained transformer (e.g., DistilBERT via Hugging Face) and visualize the attention weight matrices for a real sentence using BertViz; identify which heads appear to track syntactic versus semantic relationships.
• Write a short explanation of why multi-head attention uses H independent projection matrices instead of one large one, focusing on the representational diversity argument.