Attention, from scratch

A LITTLE HISTORY

From translation to transformers

The attention mechanism is the technical breakthrough that led to transformers, and transformers are what made modern large language models possible. It started with machine translation. A basic approach from converting English to French is to take each English word and look up the French word it translates into. But that doesn't work well:

• Word order may not be the same in English and French.
• Sentences can be different lengths. For example, "they arrived late" is five words in French.

Around 2014, Yoshua Bengio's group at the University of Montreal and Chris Manning's group at Stanford independently invented an early form of attention mechanism inside an encoder-decoder model:

• The encoder read one word at a time and produced an output vector per word. Earlier approaches compressed the whole sentence into a single dense vector; these papers preserved every word's vector and handed all of them to the decoder.
• The per-word vectors captured the meaning of each word in the context of the sentence now called contextual embeddings.
• The decoder generated the output one word at a time, weighting — attending to — each input word's embedding independently, based on where that word sits in the input and where the decoder is in producing the output. The model attends to the words most relevant for that step of the translation.

In 2017, "Attention Is All You Need" (Vaswani et al., Google Brain) introduced the transformer and a more general form of attention. The number one criterion for its design choices: will this scale on a GPU.

• Its encoder creates contextual embeddings for the whole input in a single pass; its decoder produces the output one word at a time, with each output fed back as context for the next step.
• The encoder went on to become BERT ("Bidirectional Encoder Representations from Transformers"), the basis of nearly all embedding models used for RAG and recommenders today.
• The decoder became the GPT ("Generative Pre-trained Transformer") family and most other popular models like Claude, ChatGPT, LLaMa.
• The original paper used six layers of attention; Llama 3.2-405B uses 126. The basic architecture is the same.

THE SETUP

Three parts of a transformer

Transformers can look complicated, but fundamentally they require three main parts:

• Word embedding converts words, bits of words, and symbols — collectively called tokens — into numbers.
• Positional encoding keeps track of word order. "Dogs chase cats" and "cats chase dogs" use the exact same words with very different meanings.
• Attention establishes relationships among words.

In this post, we will build the three types of attention a transformer uses — self-attention, masked self-attention, and encoder-decoder attention also called cross-attention. For each type, we will code it from scratch and compare the matrix computations by hand to get a deeper understanding of how it works.

In the scheme of a transformer architecture, this is how they all sit together.

The Transformer Architecture: From the original paper by Vaswani et al. (2017), with the three attention blocks highlighted: blue self-attention in the encoder, orange masked self-attention in the decoder, and the gradient block bridging them with encoder-decoder attention.

TERMINOLOGY

Queries, keys, and values

In "the coffee spilled on the laptop and it stopped working", "it" could refer to the laptop or to the coffee — attention is the mechanism that correctly associates "it" with the laptop.

The most basic type is self-attention: it computes how similar each word is to every other word in the sentence, including itself, and it does this for every word. The similarities then determine how the transformer encodes each word: if "it" was more commonly associated with laptops than with coffee in the training sentences, the similarity score for the laptop gives it a larger impact on how "it" is encoded.

The attention equation uses three matrices: Q for query, K for key, and V for value. The terms come from database terminology.

A hotel database pairs each guest's last name with a room number. A guest checks in and the clerk types a misspelled last name, Merser instead of Mercer:

• The query — the misspelling the computer uses to search the database.
• The keys — the actual names stored in the database. The computer compares the query to all of the keys and ranks each one.
• The value — what the database returns as the result of the search: room 537.

In a transformer, the same roles appear. Given the prompt "draw a map", the model converts each word into a word embedding, adds positional encoding, and gets the encodings representing each word — two numbers per word in this example, though 512 or more is far more common.

To create Q, K, and V:

• Stack the encodings in a matrix and multiply it by a 2×2 matrix of query weights — two query numbers per word. Likewise multiply by the key weights for K and the value weights for V.
• The weight matrices are 2×2 because we started with two encoded values per word; with 512, a common choice is 512×512. The only rule that matters: the matrix math has to be possible.
• The weight matrices carry a transpose symbol because PyTorch prints its weights in a form that has to be transposed before the math works out.

THE EQUATION

Self-attention

Self-attention learns the relationships between each word and every other word in the sentence, including itself, by considering the entire context

So here is the equation (looks indimidating). Let's break it step by step:

Step 1 — QK^T: similarity between every pair of words

Multiply the query matrix Q by the transpose of the key matrix K. Transposing K is needed for the multiplication to work — but more importantly, it yields the dot product of each query with each key: multiply pairs of corresponding numbers, add them up. A dot product is an unscaled measure of similarity between two things, closely related to cosine similarity (cosine similarity scales the result into [−1, 1]; the dot product is not scaled). QK^T gives the unscaled dot-product similarities between all possible combinations of queries and keys.

Step 2 — scale by √d_k

Scale the similarities by the square root of d_k — the dimension of the key matrix, i.e. the number of values per token. Dividing by the square root does not scale the similarities in any systematic way, but even with this limited scaling, the original authors reported it improved performance.

Step 3 — softmax each row

Take the softmax of each row. Softmax makes every row sum to 1, so the values read as a summary of relationships among the tokens: for "draw a map", the row for "draw" might be 36% similar to itself, 40% to "a", and 24% to "map". Section 6 opens up how softmax pulls this off.

The softmax row is a recipe. The percentages fix how much influence each word's value has on the final encoding of "draw".

Step 4 — multiply by V

Multiply the percentages by the values in V. The percentages fix how much influence each word has on the final encoding of any given word: 36% of draw's value + 40% of a's value + 24% of map's value gives the first self-attention score for "draw". Repeat for every row and every column of V.

The four steps as the original paper draws them. Redrawn from Figure 2 of Vaswani et al. (2017). V skips the similarity math and joins only at the last MatMul. The dashed Mask box stays idle until §10.

A CLOSER LOOK

Softmax, demystified

Softmax turns raw scores into shares of a whole — positive percentages that sum to 100% — and it has to survive negative inputs.

The obvious trick fails. Say three judges score three desserts 2.0, 1.0, and −1.0, and we want each dessert's share of the praise. Dividing by the sum gives 100%, 50%, and −50% — a negative share of praise, which means nothing. And scores that sum to zero cannot be divided at all.

Softmax adds one move before dividing: push every score through the exponential, e^x. That buys exactly the two properties needed:

• e^x is always positive — e⁻¹ is 0.368, not a negative number — so negative scores become small shares instead of nonsense.
• e^x amplifies gaps — a score 1 higher gets an e ≈ 2.72× bigger share, not a slightly bigger one.

For the three desserts:

"Soft" max. A hard max hands the winner 100% and everyone else 0%. Softmax gives the winner the biggest share but keeps every option alive — which is what lets gradients flow during training.

Exponentiate, then divide by the total. That is the whole equation:

Mapped back to the desserts: x_i is one score (2.0), e^x_i is its positivized, gap-amplified version (7.389), the denominator is the total (10.475), and the ratio is the share (70.5%). Every term is positive and the denominator is their sum, so the outputs are guaranteed positive and guaranteed to sum to 1. No special cases.

In attention, the items are the keys competing for one query's attention: each row of the scaled similarity matrix goes through this recipe independently, which is why every row of the softmax matrix sums to 1. The exponential is also what makes masking exact — e^−∞ = 0, so a masked position does not get a small share. It gets none.

IMPLEMENTATION

Coding self-attention in PyTorch

Three imports cover everything: torch for tensors — multi-dimensional lists optimized for neural networks — plus its helper functions; torch.nn for the Module and Linear classes; and torch.nn.functional for softmax.

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

The class:

• Inherits from nn.Module, the base class for every neural network module you make with PyTorch and calls the parent's __init__ which initializes a host of features including gradient computation
• n_embed is the number of word embedding values per token; it sets the size of the weight matrices that create the queries, keys, and values.
• Each weight matrix is an nn.Linear, where in_features sets the rows and out_features sets the columns (both n_embed here). The Linear object does not just store the weights — it also does the math when the time comes.
• The original transformer manuscript does not add bias terms when calculating attention, so bias=False.
• row and col are convenience parameters: usually the first dimension is the batch size, this example does not use batches, and these let us adjust later.

class SelfAttention(nn.Module):

    def __init__(self, n_embed=2, row=0, col=1):
        super().__init__()
        self.n_embed = n_embed
        self.WQ = nn.Linear(n_embed, n_embed, bias=False)
        self.WK = nn.Linear(n_embed, n_embed, bias=False)
        self.WV = nn.Linear(n_embed, n_embed, bias=False)
        self.row = row
        self.col = col

    def forward(self, x):
        Q = self.WQ(x)
        K = self.WK(x)
        V = self.WV(x)

        q_dot_k = Q @ K.transpose(dim0=self.row, dim1=self.col)
        scaled_q_dot_k = q_dot_k / torch.tensor(K.size(self.col))**0.5
        softmax = F.softmax(scaled_q_dot_k, dim=self.col)
        attention_scores = softmax @ V
        return attention_scores

forward is the four steps in order: pass the token encodings (word embeddings plus positional encoding) through W_Q, W_K, W_V to get Q, K, V; multiply Q by K transposed; scale by the square root of the number of values per key; softmax to get the attention percents; multiply by V for the attention scores. There are no rules about the weight matrix shapes beyond the multiplication working out.

To test: a matrix of encodings for three tokens, two numbers each. torch.manual_seed seeds the random number generator so everyone gets the same results. Calling the object passes the matrix to forward — that is what inheriting from nn.Module buys.

c = torch.tensor([[1.16, 0.23],
                  [0.57, 1.36],
                  [4.41, -2.16]])

torch.manual_seed(42)
sa = SelfAttention(n_embed=2, row=0, col=1)
sa(c)

tensor([[1.0100, 1.0641],
        [0.2040, 0.7057],
        [3.4989, 2.2427]], grad_fn=<MmBackward0>)

The grad_fn bit of the tensor is used for training the weights with backpropagation. We are coding a self-attention class, not a full transformer, so no training happens here.

VERIFICATION

Checking the math by hand

A tensor of numbers proves nothing by itself. The spreadsheet redoes the whole calculation and must land on the same answer.

First, extract the weights the object is actually using — print them and pass the encodings directly to each Linear to validate the math the class performs. PyTorch prints the weights in a form that has to be transposed before the math works on paper; the matrices below carry the transpose. Every step from here is a plain row × column sum of products.

sa.WQ.weight, sa.WK.weight, sa.WV.weight

(tensor([[ 0.5406,  0.5869], [-0.1657,  0.6496]]),   # WQ
 tensor([[-0.1549,  0.1427], [-0.3443,  0.4153]]),   # WK
 tensor([[ 0.6233, -0.5188], [ 0.6146,  0.1323]]))   # WV

Encodings × weights → Q, K, V

Every entry is a sum of products. The first query number: 1.16 × 0.5406 + 0.23 × 0.5869 = 0.7621:

QK^T, then scale by √2

Dot product of each query row with each key row, then divide by √2 ≈ 1.4142:

Softmax, in two moves

Softmax as defined: exponentiate every entry, then divide it by its row sum. Every row of the result sums to 1:

The third row already tells a story: the last token puts 89.6% of its attention on itself. Its encoding (4.41, −2.16) sits far from the other two, so it mostly keeps its own value.

Softmax × V → attention scores

The first score: 0.357 × 0.6037 + 0.401 × (−0.3503) + 0.242 × 3.8694 = 1.0100:

ENCODERS VS DECODERS

Why mask anything?

Self-attention looks both ways. Masked self-attention only looks at the words that came before it.

The difference is easiest to see by asking why embeddings exist at all:

• Assigning each word one random number works in theory, but similar words like "great" and "awesome" end up with very different numbers — the network needs far more complexity and training, because learning to process "great" will not help it use "awesome".
• Similar words used in similar ways should get similar numbers, so learning one helps learn the other.
• The same word appears in different contexts — a sincere "this coffee is great" versus a sarcastic "my laptop just died, great" — so each word should get more than one number.

Years before transformers, standalone neural networks created word embeddings:

• One input per unique word, one output per word, inputs connected to activation functions. The number of activation functions decides how many numbers represent each word; the weights on the connections from inputs to activation functions are the embedding values, initialized randomly and trained.
• The training objective: each word in the training data should predict the next word. After training, "great" and "awesome" cluster together on a graph of the embeddings — similar words used in similar contexts.
• Feeding in more context (say, the preceding four words predicting the next) improves the embeddings but ignores word order: "the coffee spilled on" produces the same inputs as the jumbled "coffee on spilled the".
• A transformer's positional encoding layer fixes that, and the attention layer then establishes relationships among the words.

Encoder-only transformers

When self-attention factors in all of the words — including those that came after the word of interest — the output is a new kind of embedding called a context aware embedding, or contextualized embedding. Where word embeddings only cluster individual words, context aware embeddings can cluster similar sentences and even similar documents. They can also feed a normal neural network that classifies sentiment, or serve as variables in a logistic regression model that does classification. Transformers that only use self-attention are called encoder-only transformers.

Decoder-only transformers

A decoder-only transformer starts with the same word embedding and positional encoding, but uses masked self-attention instead of self-attention. The big difference: self-attention can look at words before and after the word of interest; masked self-attention ignores every word that comes after it. For the first word of "the coffee spilled on the laptop and it stopped working", it may use only "the". At "it", it sees everything up to "it" and nothing after.

Because a decoder-only transformer can never look ahead at what comes next, it can be trained to generate responses to prompts: give it the first part of a sentence and modify the weights until it generates the rest. This is why ChatGPT — a decoder-only transformer — is called a generative model: it was specifically trained to generate the text that comes after a prompt. Where an encoder-only transformer creates context aware embeddings, a decoder-only transformer creates generative inputs that plug into a simple neural network that generates new tokens.

THE EQUATION, MASKED

The mask, step by step

The only difference between the self-attention equation and the masked version: add a new matrix M — M for mask — to the scaled similarities. Q, K, V, the similarities, and the √d_k scaling are calculated just like before, so the hand calculation resumes from the scaled-similarity matrix above.

The purpose of the mask is to prevent any token from including anything that comes after it. For "draw a map":

• "draw" only includes itself.
• "a" includes itself and "draw".
• "map" includes everything.

M adds zeros to the values we want to include and negative infinity to any value we need to mask out. Adding zero changes nothing. Adding −∞ turns the scaled similarity into −∞, and after softmax those positions get exactly 0% — since e^−∞ = 0.

Lower triangle open, upper triangle blocked. Each row may attend to itself and anything to its left — never to anything after it.

The masked hand calculation

Add M to the scaled similarities (the spreadsheet stands in −1×10²³ for −∞), then exponentiate — the blocked entries collapse to zero:

Row-normalize and multiply by the same V as before:

Three things to notice. The first token attends 100% to itself, so its score is exactly its own row of V. The second row mixes only the first two tokens. The third row is identical to unmasked self-attention — the last token was always allowed to see everything.

IMPLEMENTATION

Coding masked self-attention

The class is the previous one plus a mask switch — off, and it computes the original self-attention; one class, both behaviors. Building the mask takes three moves:

1. torch.ones creates a 3×3 matrix of ones, because there are three tokens.
2. torch.tril turns the ones in the upper triangle into zeros and leaves the lower triangle alone.
3. Comparing with == 0 converts the zeros into Trues — the positions to mask out — and the ones into Falses.

masked_fill then replaces every True position in the scaled similarities with a very large negative number standing in for −∞, and leaves every False position alone. Print the mask before trusting it.

class MaskedSelfAttention(nn.Module):

    def __init__(self, n_embed=2, row=0, col=1, mask=False):
        super().__init__()
        self.n_embed = n_embed
        self.WQ = nn.Linear(n_embed, n_embed, bias=False)
        self.WK = nn.Linear(n_embed, n_embed, bias=False)
        self.WV = nn.Linear(n_embed, n_embed, bias=False)
        self.row = row
        self.col = col
        self.mask = mask

    def forward(self, x):
        Q = self.WQ(x)
        K = self.WK(x)
        V = self.WV(x)

        q_dot_k = Q @ K.transpose(dim0=self.row, dim1=self.col)
        scaled_q_dot_k = q_dot_k / torch.tensor(K.size(self.col))**0.5

        # masking
        if self.mask:
            tril = torch.tril(torch.ones(scaled_q_dot_k.size()))
            scaled_q_dot_k = scaled_q_dot_k.masked_fill(tril == 0, float('-inf'))
        softmax = F.softmax(scaled_q_dot_k, dim=self.col)
        attention_scores = softmax @ V
        return attention_scores

Same encodings, same seed, same parameters, mask on:

torch.manual_seed(42)
msa = MaskedSelfAttention(n_embed=2, row=0, col=1, mask=True)
msa(c)

tensor([[ 0.6038,  0.7434],
        [-0.0062,  0.6072],
        [ 3.4989,  2.2427]], grad_fn=<MmBackward0>)

Exactly the hand-calculated matrix: first row equal to V's first row, last row unchanged. With mask=False the same object reproduces the original self-attention output — a free check that masking is the only thing that changed.

CROSS-ATTENTION

Encoder-decoder attention

Before encoder-only and decoder-only transformers existed, the first transformer ever made had both parts: an encoder that used self-attention and a decoder that used masked self-attention, connected so they could calculate a third thing — encoder-decoder attention.

The recipe changes in one place: the encoder's output calculates the keys and values, and the queries are calculated from the decoder's masked self-attention output. Once Q, K, and V exist, encoder-decoder attention is calculated just like self-attention, using every similarity — no mask. This first transformer was a seq-to-seq (encoder-decoder) model built to translate one language into another: the encoder calculates self-attention over "the sky is blue", and the decoder uses the encoder's output to calculate encoder-decoder attention, which generates "el cielo es azul".

Trace it on the architecture map in the overview: K and V climb out of the encoder tower into the gradient block; Q rises from below.

The names fall out of this history. The encoder alone proved useful — encoder-only transformers. The decoder alone could generate text, including translations — decoder-only transformers.

Seq-to-seq models have faded for pure language modeling, but the pattern survives in multimodal models: an encoder trained on images or sound produces context aware embeddings, and a text decoder consumes them through encoder-decoder attention to caption images or answer audio prompts.

SCALING UP

Multi-head attention

For a simple example, one attention unit works fine. To correctly establish how words relate in longer, more complicated sentences and paragraphs, transformers apply attention to the encoded values multiple times simultaneously. Each attention unit is a head, and each head has its own sets of weights for calculating the queries, keys, and values. Multiple heads calculating attention is multi-head attention. The manuscript that first described transformers used eight heads.

More heads mean more outputs: three heads producing two attention values each give six values where the encoding had two. Three ways back down to the original number of encoded values:

1. Connect all the attention values to a fully connected layer with that many outputs.
2. Shrink the value weight matrix — with one column of weights, each head outputs one value, so two heads land back on two.
3. Code the transformer to be flexible about these sizes.

Multi-head attention as the original paper draws it. Redrawn from Figure 2 of Vaswani et al. (2017). The stacked outlines are the h heads. Each head's Linear projections are its own W_Q, W_K, W_V; Concat is our torch.cat; the final Linear brings the concatenated values back to the original size.

IMPLEMENTATION

One class for all three

One Attention class implements all three attention types. The __init__ is unchanged. The forward changes twice:

• It accepts different encodings for the queries, keys, and values.
• It passes those potentially different encodings to the matrices that create Q, K, and V. Everything else is the same.

The same tensor three times is self-attention. Encoder output for keys and values, decoder output for queries is cross-attention. mask=True is masked.

class Attention(nn.Module):

    def __init__(self, n_embed=2, row=0, col=1):
        super().__init__()
        self.n_embed = n_embed
        self.WQ = nn.Linear(n_embed, n_embed, bias=False)
        self.WK = nn.Linear(n_embed, n_embed, bias=False)
        self.WV = nn.Linear(n_embed, n_embed, bias=False)
        self.row = row
        self.col = col

    def forward(self, enc_q, enc_k, enc_v, mask=False):
        Q = self.WQ(enc_q)
        K = self.WK(enc_k)
        V = self.WV(enc_v)

        q_dot_k = Q @ K.transpose(dim0=self.row, dim1=self.col)
        scaled_q_dot_k = q_dot_k / torch.tensor(K.size(self.col))**0.5

        # masking
        if mask:
            tril = torch.tril(torch.ones(scaled_q_dot_k.size()))
            scaled_q_dot_k = scaled_q_dot_k.masked_fill(tril == 0, float('-inf'))
        softmax = F.softmax(scaled_q_dot_k, dim=self.col)
        attention_scores = softmax @ V
        return attention_scores

MultiHeadAttention adds one new parameter, num_heads. A for loop creates that many Attention objects — each initialized with the same arguments, each carrying its own weights — stored in an nn.ModuleList, a list of modules we can index. The forward pass loops the encoding matrices through every head and concatenates what comes back.

class MultiHeadAttention(nn.Module):

    def __init__(self, n_embed=2, row=0, col=1, num_heads=1):
        super().__init__()
        self.heads = nn.ModuleList([Attention(n_embed, row, col) for _ in range(num_heads)])
        self.col = col

    def forward(self, enc_q, enc_k, enc_v, mask=False):
        attention_scores_concatenated = torch.cat([
            head(enc_q, enc_k, enc_v) for head in self.heads
        ], dim=self.col)
        return attention_scores_concatenated

Using the same encodings for queries, keys, and values keeps the results comparable to everything above. One head with the same seed must reproduce the single-head answer. Two heads must produce twice as many values:

q, k, v = c, c, c

torch.manual_seed(42)
ca = Attention(n_embed=2, row=0, col=1)
ca(q, k, v)          # same output as SelfAttention — 1.0100, 1.0641, ...

torch.manual_seed(42)
mha = MultiHeadAttention(n_embed=2, row=0, col=1, num_heads=2)
mha(q, k, v)

tensor([[ 1.0100,  1.0641, -0.7081, -0.8268],
        [ 0.2040,  0.7057, -0.7417, -0.9193],
        [ 3.4989,  2.2427, -0.7190, -0.8447]], grad_fn=<CatBackward0>)

The first two columns are the single-head answer — the first head drew the same seed-42 weights. The last two are the second head's view of the same tokens with its own freshly initialized weights. Six values per token from two in: exactly the growth the fully connected layer (or a slimmer value matrix) brings back down inside a real transformer.