Every large language model you have used is, at its core, dot products, a softmax, and a weighted sum. That is the attention mechanism, and in this workshop you build it from scratch in PyTorch, tensor by tensor. You will compute Query, Key, and Value by hand on a four-word sentence, read a real heatmap, assemble multi-head attention, add causal masking, and train a tiny model that predicts text.

“If you can compute it by hand once, you can debug it forever.”

Last Updated: July 2026 | Tested on: Python 3.14.6, PyTorch 2.12.1, NumPy 2.4.6 | Difficulty: Expert | Reading Time: 17 minutes

This is the finale of the PyTorch track, and it is the bridge you cross before the full Transformer. Interviewers love to ask you to reproduce this exact calculation on a whiteboard, so by the end you will be able to write it out from memory. The math here is evergreen: dot products, a softmax, and a weighted sum have not changed since the 2017 paper and will not change. Only the library function names drift over time, so those are pinned to versions you can check.

Prerequisites

Why Attention: The RNN Bottleneck

Picture a friend named Anvay reading a long recipe out loud to you over the phone, one word at a time, and you are only allowed to remember a single short note that you keep rewriting after each word. By the time he reaches the last step you have almost certainly forgotten the quantity of rice from the first line. That is exactly how a Recurrent Neural Network (RNN) works. It walks through a sentence left to right, squeezing everything it has seen so far into one fixed-size hidden state. The further back a detail sits, the more likely it has been overwritten. This is the bottleneck: one small vector has to carry the entire past.

Attention throws that limit out. Instead of forcing every earlier word through a single memory slot, it lets each word look directly at every other word and pull exactly the parts it needs. A word at position 1 and a word at position 50 are the same distance apart, one hop, so nothing gets lost to time. Because every position is handled together rather than in a chain, the whole thing runs in parallel on a GPU (Graphics Processing Unit). That single change is what made training on internet-scale text practical, and it is why the attention mechanism replaced recurrence almost everywhere.

Q, K, V by Hand: Scaled Dot-Product Attention

Here is the analogy that makes Query, Key, and Value click. Think of a small library. Your Query is the topic you are hunting for, each book’s Key is the label printed on its spine, and each book’s Value is the actual content inside. You compare your topic against every spine, decide how much each book matters, then blend the contents of the best matches into one answer. Self-attention does this for every word at once: each word issues a Query, every word offers a Key and a Value, and the output for each word is a weighted blend of all the Values.

The diagram traces the whole pipeline as tensors. We start from the input embeddings X, multiply by three learned weight matrices to get Q, K, and V, score each Query against every Key with a dot product, scale the scores down by the square root of the head dimension so softmax stays gentle, optionally mask out future tokens, turn the scores into probabilities with softmax, and finally blend the Values. Multi-head attention just runs this same block several times side by side. Let us build it on a four-word vegetarian sentence and print every tensor along the way.

📄 scaled_dot_product.py: attention from scratch on a 4-token sentence

import torch
import torch.nn.functional as F

torch.manual_seed(0)
torch.set_printoptions(precision=3, sci_mode=False)

words = ["Aditi", "loves", "warm", "dosa"]
seq_len, d_model = 4, 4

# Toy token embeddings (normally learned)
X = torch.randn(seq_len, d_model)

# One shared set of projection weights (normally learned)
W_q = torch.randn(d_model, d_model) * 0.5
W_k = torch.randn(d_model, d_model) * 0.5
W_v = torch.randn(d_model, d_model) * 0.5

Q = X @ W_q            # Queries: "what am I looking for?"
K = X @ W_k            # Keys:    "what do I contain?"
V = X @ W_v            # Values:  "what do I offer?"

d_k = Q.shape[-1]
scores = Q @ K.T                 # every query dotted with every key
scaled = scores / (d_k ** 0.5)   # keep the numbers gentle
weights = F.softmax(scaled, dim=-1)  # each row becomes probabilities
output = weights @ V             # blend the values

print("Q = X @ W_q")
print(Q)
print("\nRaw scores = Q @ K^T, shape", tuple(scores.shape))
print(scores)
print("\nScaled scores = scores / sqrt(d_k), d_k =", d_k)
print(scaled)
print("\nAttention weights = softmax(scaled), each row sums to 1")
print(weights)
print("row sums:", weights.sum(dim=-1))
print("\nOutput = weights @ V, shape", tuple(output.shape))
print(output)

▶ Output

Q = X @ W_q
tensor([[ 0.021,  1.580, -0.306, -1.354],
        [-1.955, -2.843, -0.851,  1.841],
        [-0.251,  0.851,  0.422, -1.140],
        [ 0.537, -1.587, -0.391,  1.401]])

Raw scores = Q @ K^T, shape (4, 4)
tensor([[-1.081, -7.079,  1.543, -0.449],
        [ 1.626, 10.568, -0.826, -0.258],
        [-0.402, -3.490,  0.662, -0.255],
        [ 0.895,  5.704, -1.213,  0.265]])

Scaled scores = scores / sqrt(d_k), d_k = 4
tensor([[-0.540, -3.540,  0.772, -0.224],
        [ 0.813,  5.284, -0.413, -0.129],
        [-0.201, -1.745,  0.331, -0.128],
        [ 0.447,  2.852, -0.607,  0.132]])

Attention weights = softmax(scaled), each row sums to 1
tensor([[0.163, 0.008, 0.605, 0.224],
        [0.011, 0.981, 0.003, 0.004],
        [0.251, 0.053, 0.426, 0.270],
        [0.076, 0.842, 0.026, 0.055]])
row sums: tensor([1.000, 1.000, 1.000, 1.000])

Output = weights @ V, shape (4, 4)
tensor([[-0.064,  0.280,  0.199,  0.837],
        [ 1.308,  1.397, -1.511, -0.008],
        [ 0.113,  0.295, -0.020,  0.760],
        [ 1.163,  1.220, -1.334,  0.090]])

What happened here: Follow the shapes and it all lines up. Four tokens each become a Query, a Key, and a Value. The score matrix is 4 by 4 because every token scores itself against all four tokens, so entry (i, j) is how much token i cares about token j. Notice the raw scores reach 10.5, which is exactly the kind of spike that would make softmax collapse onto a single word.

Dividing by sqrt(d_k), here the square root of 4, pulls that 10.5 down to 5.28 and keeps the gradients healthy. After softmax every row is a clean probability distribution that sums to 1, and the output is each token’s own blend of the four Value vectors. These weights come from random untrained matrices, so treat the pattern as mechanical rather than meaningful for now.

Reading the Attention Weights

The weight matrix is the whole point of attention, so it helps to see it as a grid instead of raw numbers. Each row is one word asking “who do I pay attention to?” and the four values in that row are the percentages, adding up to 1. In real models people render this grid as a color heatmap where darker cells mean stronger attention. We can print the same thing as plain text straight from the tensor above.

📄 heatmap.py: printing the attention grid

print("Attention map (rows = query word, columns = attended word):")
header = "          " + "".join(f"{w:>8}" for w in words)
print(header)
for i, w in enumerate(words):
    row = "".join(f"{weights[i][j].item():>8.2f}" for j in range(seq_len))
    print(f"{w:>10}{row}")

▶ Output

Attention map (rows = query word, columns = attended word):
             Aditi   loves    warm    dosa
     Aditi    0.16    0.01    0.61    0.22
     loves    0.01    0.98    0.00    0.00
      warm    0.25    0.05    0.43    0.27
      dosa    0.08    0.84    0.03    0.06

What happened here: Read it row by row. The word “loves” puts 0.98 of its attention on itself, while “dosa” sends 0.84 of its attention back to “loves”. In a trained model this is where linguistic structure shows up: a verb learning to look at its object, a pronoun learning to look at the noun it refers to. When you later debug a large model that keeps confusing two names, this exact grid is the first thing you pull up. Being able to produce it from a bare tensor, with no library helper, is what the interview question is really checking.

Multi-Head Attention, Verified Against PyTorch

One attention head can only follow one kind of relationship at a time. Asking a single reader named Aviraj to track grammar, meaning, and word order all at once is a lot, and he will drop something. Multi-head attention hires a small team instead: it splits the embedding into several slices and runs a separate attention on each, so one head can chase grammar while another chases meaning. The heads are then stitched back together and passed through one final projection. The best way to prove you built it correctly is to match PyTorch’s own nn.MultiheadAttention number for number, by copying its weights into your version and comparing outputs.

📄 multi_head.py: build it, then verify against PyTorch

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

torch.manual_seed(1)

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        assert d_model % num_heads == 0
        self.h = num_heads
        self.d_k = d_model // num_heads
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)

    def forward(self, x, mask=None):
        B, T, D = x.shape
        q = self.W_q(x).view(B, T, self.h, self.d_k).transpose(1, 2)
        k = self.W_k(x).view(B, T, self.h, self.d_k).transpose(1, 2)
        v = self.W_v(x).view(B, T, self.h, self.d_k).transpose(1, 2)
        scores = (q @ k.transpose(-2, -1)) / (self.d_k ** 0.5)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float("-inf"))
        attn = F.softmax(scores, dim=-1)
        ctx = (attn @ v).transpose(1, 2).contiguous().view(B, T, D)
        return self.W_o(ctx)

d_model, heads = 8, 2
mine = MultiHeadAttention(d_model, heads)
ref = nn.MultiheadAttention(d_model, heads, batch_first=True)

# Copy PyTorch's packed Q/K/V weights into our separate layers
wq, wk, wv = ref.in_proj_weight.chunk(3, dim=0)
bq, bk, bv = ref.in_proj_bias.chunk(3, dim=0)
with torch.no_grad():
    mine.W_q.weight.copy_(wq); mine.W_q.bias.copy_(bq)
    mine.W_k.weight.copy_(wk); mine.W_k.bias.copy_(bk)
    mine.W_v.weight.copy_(wv); mine.W_v.bias.copy_(bv)
    mine.W_o.weight.copy_(ref.out_proj.weight); mine.W_o.bias.copy_(ref.out_proj.bias)

x = torch.randn(1, 5, d_model)
out_mine = mine(x)
out_ref, _ = ref(x, x, x)

print("Our output shape:      ", tuple(out_mine.shape))
print("PyTorch output shape:  ", tuple(out_ref.shape))
print("Max absolute difference:", (out_mine - out_ref).abs().max().item())
print("Outputs match (atol 1e-5):", torch.allclose(out_mine, out_ref, atol=1e-5))

▶ Output

Our output shape:       (1, 5, 8)
PyTorch output shape:   (1, 5, 8)
Max absolute difference: 5.960464477539063e-08
Outputs match (atol 1e-5): True

What happened here: The two outputs agree to about six decimal places, and the tiny gap of 6e-08 is just floating point rounding, not a real difference. That match is proof your from-scratch version is doing exactly what the battle-tested library does. The one subtlety worth remembering: PyTorch packs Q, K, and V into a single in_proj_weight tensor for speed, so we chunk it into three before copying. The .view(B, T, h, d_k).transpose(1, 2) dance is how the heads get their own slices, and .contiguous().view(...) stitches them back. This same block, with d_model in the thousands, runs inside GPT, Claude, Llama, and every other current large model (accurate at the time of writing, since architectures keep evolving).

Causal Masking for Text Generation

When a model generates text one word at a time, it must not peek at words it has not written yet, otherwise training would be cheating and generation would fall apart. Causal masking enforces this. Before softmax you set the scores for all future positions to negative infinity, so after softmax they become exactly zero. Think of writing a sentence where you are only ever allowed to look back over your shoulder at what you already put down, never forward. This one trick is the entire difference between an encoder that reads and a decoder that writes.

📄 causal_mask.py: block every token from seeing the future

import torch
import torch.nn.functional as F

torch.manual_seed(3)
torch.set_printoptions(precision=2, sci_mode=False)

T = 4
tokens = ["", "warm", "dosa", "please"]
scores = torch.randn(T, T)  # pretend these came from Q @ K^T / sqrt(d_k)

mask = torch.tril(torch.ones(T, T))          # lower triangle of ones
masked = scores.masked_fill(mask == 0, float("-inf"))
weights = F.softmax(masked, dim=-1)

print("Causal mask (1 = allowed, 0 = blocked):")
print(mask)
print("\nAttention weights after causal masking:")
print(weights)
print("\nRow by row (a token can only look left, never right):")
for i, tok in enumerate(tokens):
    seen = ", ".join(tokens[j] for j in range(T) if weights[i][j] > 0)
    print(f"  {tok:>8} can attend to: {seen}")

▶ Output

Causal mask (1 = allowed, 0 = blocked):
tensor([[1., 0., 0., 0.],
        [1., 1., 0., 0.],
        [1., 1., 1., 0.],
        [1., 1., 1., 1.]])

Attention weights after causal masking:
tensor([[1.00, 0.00, 0.00, 0.00],
        [0.72, 0.28, 0.00, 0.00],
        [0.32, 0.25, 0.43, 0.00],
        [0.36, 0.17, 0.25, 0.22]])

Row by row (a token can only look left, never right):
    can attend to: 
      warm can attend to: , warm
      dosa can attend to: , warm, dosa
    please can attend to: , warm, dosa, please

What happened here: The mask is a lower triangular matrix of ones, and everything above the diagonal is a zero that becomes negative infinity before softmax. Look at the weight grid: the whole upper triangle is exactly 0.00, so the first token attends only to itself, the second sees two tokens, and only the last row is full. This is why a decoder can be trained on a whole sentence at once yet still learn to predict each word using only the words before it. Anvi, our text model in the next section, relies on precisely this mask.

A Mini Self-Attention Model You Can Train

Time to put every piece together into something that actually learns. Below is a tiny character-level self-attention model, roughly thirty lines, trained on one short repeating phrase. It has token and position embeddings, a single masked self-attention layer, and a linear head that predicts the next character. It is a Transformer decoder shrunk to its bones, and it is enough to watch the loss fall and the model reproduce text.

📄 mini_self_attention.py: a trainable character model

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

torch.manual_seed(7)

text = "aditi likes warm dosa "
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for c, i in stoi.items()}
data = torch.tensor([stoi[c] for c in text])
V, D, T = len(chars), 24, 8   # vocab, embed dim, context length

class MiniSelfAttention(nn.Module):
    def __init__(self):
        super().__init__()
        self.tok = nn.Embedding(V, D)
        self.pos = nn.Embedding(T, D)
        self.q = nn.Linear(D, D); self.k = nn.Linear(D, D); self.v = nn.Linear(D, D)
        self.head = nn.Linear(D, V)
        self.register_buffer("mask", torch.tril(torch.ones(T, T)))

    def forward(self, idx):
        B, t = idx.shape
        x = self.tok(idx) + self.pos(torch.arange(t))
        att = (self.q(x) @ self.k(x).transpose(-2, -1)) / (D ** 0.5)
        att = att.masked_fill(self.mask[:t, :t] == 0, float("-inf"))
        x = F.softmax(att, dim=-1) @ self.v(x)
        return self.head(x)

model = MiniSelfAttention()
opt = torch.optim.Adam(model.parameters(), lr=0.01)

def batch():
    i = torch.randint(0, len(data) - T - 1, (16,))
    xb = torch.stack([data[j:j+T] for j in i])
    yb = torch.stack([data[j+1:j+T+1] for j in i])
    return xb, yb

for step in range(301):
    xb, yb = batch()
    logits = model(xb)
    loss = F.cross_entropy(logits.view(-1, V), yb.view(-1))
    opt.zero_grad(); loss.backward(); opt.step()
    if step % 100 == 0:
        print(f"step {step:>3} | loss {loss.item():.3f}")

# Greedy generation from a seed
idx = torch.tensor([[stoi[c] for c in "aditi li"]])
for _ in range(14):
    logits = model(idx[:, -T:])
    nxt = logits[:, -1, :].argmax(-1, keepdim=True)
    idx = torch.cat([idx, nxt], dim=1)
print("generated:", "".join(itos[i] for i in idx[0].tolist()))

▶ Output

step   0 | loss 2.655
step 100 | loss 0.219
step 200 | loss 0.178
step 300 | loss 0.131
generated: aditi likes warm dosaw

What happened here: The loss drops from 2.655, which is roughly random guessing across the vocabulary, down to 0.131, and the model then reproduces the training phrase almost perfectly from an eight-character seed. Nothing here is fancy: token embeddings say what a character is, position embeddings say where it sits, one masked self-attention layer lets each character look back at the ones before it, and cross entropy pushes the next-character prediction. Scale the embedding size, stack more of these blocks, add feed-forward layers and layer norm, feed it billions of tokens, and you have the architecture behind every current large language model.

Common Mistakes

Interview Corner

Q: Walk me through scaled dot-product attention in one breath.

Project the input into Query, Key, and Value with three learned matrices. Score each Query against every Key with a dot product to get a sequence-by-sequence matrix. Divide those scores by the square root of the head dimension so softmax does not saturate. Apply softmax along each row so the weights become a probability distribution. Multiply those weights by the Value matrix to get, for each token, a weighted blend of all the Values. That blend is the output, one context-aware vector per token.

More in this series:

• NLP: Word Embeddings (Word2Vec, GloVe, FastText)
• NLP: BERT, GPT, and Modern Language Models
• NLP: Text Classification with HuggingFace Transformers

Frequently Asked Questions

Interview Questions on the Attention Mechanism

Try each one aloud first. The phrasing you produce under mild pressure is what interviews measure.

Q: Your attention layer works on a 64-dim toy but the loss explodes when you scale d_k to 512. What did you most likely drop?

The sqrt(d_k) scaling. At d_k of 512 the raw dot products grow large, softmax saturates so one weight sits near 1 and the rest near 0, and the gradients through that softmax shrink toward zero. Dividing the scores by the square root of d_k before softmax keeps them in a healthy range. It is a one-line fix that people forget precisely because a small toy model seems to run fine without it.

Q: How would you prove your from-scratch multi-head attention is correct?

Compare it against a reference implementation on the same weights. Instantiate nn.MultiheadAttention, copy its packed in_proj_weight (chunked into Q, K, V) and its out_proj into your module, run the same input through both, and check the maximum absolute difference. If it is around 1e-7, your version is correct and the gap is only floating point rounding. That copy-and-compare trick is a clean, convincing test.

Q: Why must the causal mask be applied before softmax rather than after?

Because softmax normalises across the row. If you mask before softmax by setting future scores to negative infinity, those positions become exactly 0 and the remaining weights still sum to 1. If you zero them after softmax, the row no longer sums to 1, the surviving weights are wrong, and you have already let the future leak into the normalisation. Order matters here.

Q: What is the time and memory complexity of self-attention, and why does it matter?

Both are O(n squared times d) where n is the sequence length, because the score matrix is n by n. That quadratic term is why long contexts are expensive: going from 512 to 8192 tokens grows the attention matrix about 250 times. It is the reason for memory-efficient kernels like FlashAttention and for a whole family of long-context attention variants.

Q: In self-attention, where do Query, Key, and Value come from, and could they be the same?

All three are linear projections of the same input, using three different learned weight matrices W_q, W_k, and W_v. They start from the same vectors but are not the same, and that separation is what lets a token ask one thing (Query), advertise another (Key), and carry a third (Value). If you tied the weights so Q equalled K, you would lose the ability to model asymmetric relationships like a verb attending to its object but not the reverse.

What’s Next?

You have now built the attention mechanism end to end: scaled dot-product attention by hand, a readable attention heatmap, multi-head attention verified against PyTorch to seven decimals, causal masking, and a tiny self-attention model that trains and generates. That single stack of ideas is the beating heart of every modern AI system. Next, in the Transformer architecture tutorial, we wrap these attention blocks in residual connections, layer normalization, positional encoding, and feed-forward layers to assemble the full encoder and decoder.

Want to see where this fits in the bigger picture? Head back to the Python + AI/ML tutorial series home for the full roadmap from Python basics to production AI.

Go deeper: when you outgrow this post, Hugging Face documentation is the next stop.

Related Posts

Previous: Deep Learning Capstone: Image Classifier on Hugging Face

Next: Natural Language Processing (NLP): Transformer Architecture, Attention Is All You Need

Series Home: Python + AI/ML Tutorial Series