Stop 08 · Deep Learning ArchitecturesTopics · 6/11

Attention Mechanism

In short

Attention lets a model decide, for every word, which other words matter most right now. It's the core idea behind transformers and every modern LLM you've used.

A technique that lets a model weigh which parts of the input matter most for each output step — the single most important idea behind modern LLMs.

The blob in an exam hall flips through a thick open textbook, with a spotlight on two pages.

Think of it like this

The open-book exam

In an open-book exam, you don't reread the whole textbook. You read the question, scan the chapter headings, and spend most of your time on the two pages that match. Attention does this: a query, matched against keys, pulling mostly from the best-matching values.

Watch it happen

Where does 'it' look?

The

animal

didn't

cross

the

street

because

1/4Sentence: 'The animal didn't cross the street because it was tired.' What does 'it' refer to?

The big ideas

What you need to know

Query, key, value

Each token makes a query (what I'm looking for), a key (what I offer) and a value (what I'll pass on).

Score, softmax, mix

Attention = softmax(Q·Kᵀ / √d) · V. Scores become weights, and weights mix the values.

Why divide by √d

Dot products grow with vector size. Scaling keeps softmax from becoming too sharp, which would stall learning.

Many heads

Multi-head attention runs several attentions in parallel. One head might track grammar, another who 'it' refers to.

See it in code

Walk through the code

python

1import torch2import torch.nn.functional as F3 4torch.manual_seed(0)5Q = torch.randn(1, 4)        # the query: what "it" is looking for6K = torch.randn(5, 4)        # keys: what each of 5 words offers7V = torch.randn(5, 4)        # values: the information each word carries8 9scores = Q @ K.T / K.shape[-1] ** 0.5     # similarity, scaled by sqrt(d)10weights = F.softmax(scores, dim=-1)       # weights that sum to 111output = weights @ V                      # weighted mix of the values12 13print(weights.round(decimals=2))14print(torch.allclose(output, F.scaled_dot_product_attention(Q[None], K[None], V[None])[0]))  # True

1/4One query and five key-value pairs, each four numbers wide. In a real model these come from learnt layers.

Scaled dot-product attention in three lines, checked against PyTorch's built-in version.

Save yourself the pain

Common mistakes

Oops

Thinking attention weights are fixed per word pair.

Do this instead

They're computed fresh for every input. The same word gets different weights in different sentences.

Oops

Reading attention weights as a perfect explanation of the model.

Do this instead

They're a useful hint about what was mixed, but not a full account of why the model decided something.

Your mission

Attention by hand

~25 min

Quick check

In attention, what decides how much a word contributes to another word's new representation?

Go deeper

1Read the original 'Attention Is All You Need' paper's intuition section (skip the math at first).
2Implement a simple attention mechanism from scratch on a toy sequence task.
3Visualize attention weights on a real example to see what the model 'focuses on.'