Attention, Explained from Scratch

Attention is the engine inside nearly every large language model we talk to today, and it is what makes the transformer a transformer. On first contact its mathematics can look forbidding, yet the idea underneath rests on a single, intuitive question: when processing one word in a sentence, which of the other words do I need to pay attention to in order to get its meaning right? In this post we'll formalize that question step by step, build the query–key–value triple, and see why attention set off an architectural revolution.

Framing the problem

The word "bank" in "she sat by the bank" means something entirely different from "bank" in "she withdrew money from the bank," even though it's the same word. For a model to capture that difference, it must represent each word not in isolation but in relation to the others around it. Attention establishes exactly that relation: it adds to a word's representation the information "gathered" from the other relevant words in the sentence.

Query, key, and value

To keep this gathering orderly, the model produces three separate vectors for each token (word piece). An analogy with searching a library sharpens the intuition. The query ($Q$) is the token's question: "what am I looking for?" The key ($K$) is each token's label: "here is what I offer." The value ($V$) is the actual content a token carries. You compare one token's query against every token's key to measure how relevant each is, turn those relevance scores into weights, and take a weighted average of the values. The result is that token's new, context-enriched representation.

The whole process fits into a single equation:

Let's take it apart. The product $QK^\top$ computes the similarity between every query and every key — through a dot product, which is once again linear algebra — yielding a score for each pair of tokens. The $\sqrt{d_k}$ in the denominator rescales these scores, which grow with the vector dimension, to keep them numerically stable. Softmax then converts the raw scores into weights that sum to one, giving us a distribution that can say "attend sixty percent to this word, forty percent to that one." Finally we multiply the values by these weights and sum.

If we squeeze this whole story into a single line, the heart of attention is this formula:

Read it left to right and you'll recognize the very intuition we just built: $Q K^\top$ multiplies every query against every key to produce similarity scores, $\sqrt{d_k}$ calms those scores down so they don't blow up as the dimension grows, softmax turns them into weights that sum to one, and the final multiplication by $V$ applies those weights to the actual content — yielding a fresh, context-enriched representation for each token. This is also exactly why the mechanism is so beloved: since it's nothing but a handful of matrix multiplications, it runs in parallel on a GPU, blazingly fast.

A few lines of code

Writing the mechanism in its plainest form dissolves the abstraction of the equation:

import torch
import torch.nn.functional as F
 
def attention(Q, K, V):
    d_k = Q.size(-1)
    scores = Q @ K.transpose(-2, -1) / d_k**0.5  # similarity for every token pair
    weights = F.softmax(scores, dim=-1)          # who attends to whom, how much
    return weights @ V                           # weighted sum of the values
 
# toy example: 5 tokens, each 16-dimensional
Q = K = V = torch.randn(5, 16)
print(attention(Q, K, V).shape)  # torch.Size([5, 16])

Each row of scores holds the attention one token pays to all the others, and after softmax weights is its probability-distribution form. The product weights @ V lets every token inherit, in weighted fashion, the content of the tokens most relevant to it.

Why it was a revolution

The previous generation of models processed text strictly left to right, word by word. In that setup, relating a subject at the start of a sentence to a pronoun at its end requires information to be carried along a long chain, and the longer the chain, the more that information fades. Attention removes this constraint entirely: every token can look directly at every other token in a single step. The pronoun "it" can connect instantly to the noun ten words earlier, independent of any chain. Better still, because these operations are independent of one another, they parallelize on modern hardware — which is exactly why transformers are both more accurate and trainable at enormous scale.

Where to go from here

We've covered attention in its single-head, plainest form. Real transformers run the mechanism as several parallel "heads" (multi-head attention) so the model can capture different kinds of relationships at once, and they add positional encoding to tell the model about word order. Upcoming posts will open up both pieces with the same care.