Artificial Neural Networks and Deep Learning · Insper

Attention Mechanisms

A dense layer whose weights come from the input

Two opposite choices

Convolution

  • decides in advance who talks to whom: the neighbours
  • the same weights at every position
  • fixed after training

Attention

  • computes the connection strengths from the data, at run time
  • for every pair of positions
  • different for every input

Attention is a dense layer whose weights are a function of the input. The , the masks, the heads and the cache that dominates the inference bill all come from that sentence — and from the cost of taking it literally.

Before: recurrence

  • The state is a fixed-size notepad, rewritten after every word .
  • The gradient from the last token to the first is a product of factors, one per token: if each is 0.9, .
  • The LSTM multiplies its cell by a gate instead of a matrix — better, but still a product.
  • And the steps are sequential: the GPU waits on one state at a time.

Does the LSTM solve the problem, or only push it further away?

Simulator · how far the gradient reaches

What retired recurrence

path from the last token to the first 64 tokens 256 tokens
plain RNN, weight scale 1 (measured in the simulator)
LSTM cell path, per step
attention one step one step
  • Recurrence decays exponentially with distance; the LSTM only improves the constant.
  • Attention reads every position at once: the gradient is scaled by one attention weight, not by a product with one factor per step.
  • Reach and parallelism — not an isolated gain in accuracy — is what made "Attention Is All You Need" (2017) drop the recurrence.

What a token is

  • The model reads integers, never letters. A token is a piece of text from a fixed list, the vocabulary, and its number is its id.
  • Whole words? The list never ends, and a word missing from it has no id.
  • Characters? Nothing is unknown, but "the cat climbed the" becomes 19 symbols, and attention costs grow with .
  • Subwords: frequent words whole, rarer ones in pieces, single bytes as the fallback.
text GPT-2 tokens
the cat climbed the 4
o gato subiu no 6
an English sentence of 16 words 17
its Portuguese version, 13 words 25

GPT-2: 256 bytes + 50,000 merges + 1 special token = 50,257. "o gato subiu no" = o ␣g ato ␣sub iu ␣no. The space belongs to the next word: "the" is id 1169, "␣the" is 262.

BPE on the running sentence

Count the neighbouring pairs of a corpus, merge the most frequent one, write it down, repeat. To tokenize new text, replay the merges in order.

after merge pieces how many
start: characters t h e ␣ c a t ␣ c l i m b e d ␣ t h e 19
1 t+h · 2 th+e · 3 ␣+the the ␣ c a t ␣ c l i m b e d ␣the 14
4 ␣+c · 5 a+t · 8 ␣c+at the ␣cat ␣c l i m b e d ␣the 10
14 e+d the ␣cat ␣c l i m b ed ␣the 9
18 l+i · 38–40 → ␣climb the ␣cat ␣climb ed ␣the 5

Toy corpus: 8 sentences ("the cat sat on the roof.", "the dog climbs on the sofa." …), 23 characters, 40 merges; "climbed" is not in it. The toy model keeps 12 tokens: the = 0, ␣the = 1, ␣cat = 2, ␣climb = 4, ed = 5 → [0, 2, 4, 5, 1].

Simulator · BPE, merge by merge

From id to vector

A lookup in the embedding table () plus the row of the position table : no arithmetic on the id itself.

token (id)
1 the (0) [−0.5, −0.5, 0, 1] [0, 0, 0, 0] [−0.5, −0.5, 0, 1]
2 ␣cat (2) [1, −0.5, 0, −0.5] [0.5, 0, −0.5, 0] [1.5, −0.5, −0.5, −0.5]
3 ␣climb (4) [−0.5, 1, 0, −0.5] [0, 0.5, 0, −0.5] [−0.5, 1.5, 0, −1]
4 ed (5) [0, 0.5, −0.5, 0] [−0.5, 0, 0.5, 0] [−0.5, 0.5, 0, 0]
5 ␣the (1) [−0.5, 0, −0.5, 1] [0, −0.5, 0, 0.5] [−0.5, −0.5, −0.5, 1.5]

Toy: V = 12, d = 4, 8 positions, weights picked by hand. GPT-2 small: E is 50,257 × 768 (38.6M parameters, 31% of the model), P is 1024 × 768, both learned. Llama has no P: RoPE rotates q and k instead.

The idea: every word re-reads the others

  • "I sat on the bank of the river."
  • "I went to the bank to withdraw cash."

The embedding layer gives "bank" a vector. Is it the same vector in both sentences?

  • Yes. The embedding looks only at the token, not at its neighbours.
  • Attention resolves it: "bank" builds a new vector by mixing the others — from "sat" and "river" in the first sentence, from "withdraw" and "cash" in the second.
  • The output: one vector per token, now with context.

A soft dictionary

stock = {"apple": 3, "pear": 7}
stock["pear"]   # exact match → 7

A dictionary demands an exact match and returns one value.

Attention compares the question with every key, gives each a similarity score and returns a weighted average of all the values.

The average is a smooth function of the scores, so it can be differentiated — and the network learns, by gradient, what to ask and what to advertise.

Query, Key, Value

  • Query — what this position is looking for. The question.
  • Key — what each position advertises about itself. The label.
  • Value — what each position delivers if selected. The content.

X: one row per token, n tokens of d numbers · dk: the length of each query and key · softmax: escore divided by the row's sum.

Three projections of the same input, and WQ, WK, WV are the only learned parameters. That is why a token can ask for one thing, advertise another and deliver a third — a pronoun can query for "the noun I refer to" without being a noun.

A worked example, by hand

"sat · bank · river", two dimensions picked by hand. The values read as .

token key value
sat
bank — ambiguous
river

The query of "bank" is and . Which word gets the largest weight? What comes out?

The four steps

, :

token ① score ② ③ ③ weight ④
sat 0.71 2.03 2.03 / 27.29 = 0.074
bank 2.12 8.34 0.306
river 2.83 16.92 0.620
sum 27.29 1

"bank" came in on the fence, , and leaves as the river bank. No weight was fixed by position: the choice came from how much the query resembles each key.

Simulator · one query, three keys

q — the query of "bank": [q1, q2]
d_k (only changes the divisor √d_k)
key k — what it advertises
value v — [river, money]
sat
bank
river

The shape of every matrix

tensor shape what it is
one vector per token
learned
learned
a query and a key per token
a value per token
a score for every pair — the
row : where token looks; sums to 1
a new vector, with context, per token

Lab, section 2: n = 6, d = 16, dk = 8, dv = 4 → A is 6 × 6, O is 6 × 4, 16·8 + 16·8 + 16·4 = 320 parameters — none of them depends on n. Our three lines match F.scaled_dot_product_attention to 1.2 × 10−7.

Simulator · the whole matrix

Why divide by

If and have independent components with mean 0 and variance 1:

  • At the scores have a standard deviation of ; they routinely differ by 20.
  • : one weight is 1, the rest are — a hard argmax, and no gradient.
  • Dividing by brings the variance back to 1.

Why not divide by ?

Measured in the lab

One query against 16 keys, 200 rows; median over the rows.

largest weight, with , with largest weight, without , without dead rows, without
4 0.21 0.26 0.37 0.27 0%
64 0.22 0.27 0.94 0.054 17%
1024 0.22 0.28 0.99999 0.000011 64%
  • With the division nothing moves as grows.
  • Without it, at the median row is one-hot and its gradient is times smaller.

Largest weight: the maximum of a row of A (uniform would be 1/16 ≈ 0.06). ‖∂L/∂S‖: the size of the gradient reaching the row's scores. "Dead" = ‖∂L/∂S‖ below 10−3. Notebook, section 3.

The way back through the softmax

  • is the element-wise product. If a row of is nearly one-hot, almost all of the bracket is multiplied by …
  • …and where is not tiny, the bracket itself is near zero: it subtracts the row's weighted mean (the copies it into every column), which is that very element.
  • A saturated row hands back a numerically zero gradient.
  • Check for free: every row of sums to zero — adding a constant to a row of scores does not change its softmax.

Lab, section 4: queries × 100 → four of five rows print a largest weight of 1.0, and their gradients fall from 0.15–0.51 to between 10−34 and 7 × 10−4.

Masks: what a token must not see

Added to the scores before the softmax, as : .

position sees must predict
1 the ␣cat
2 the␣cat ␣climb
3 the␣cat␣climb ed
4 the␣cat␣climbed ␣the
5 the␣cat␣climbed␣the ␣roof
  • The causal mask is the entire difference between an encoder and a decoder.
  • One forward pass gives n training examples instead of one. Without the mask, position 1 would see ␣cat and learn to copy.

Four masks, one line each

mask forbids where
causal looking at the future every autoregressive decoder
padding attending to <pad> filler any batch of variable-length sequences
sliding window looking further back than Mistral, Gemma, hybrid stacks
document crossing into another packed document all modern pretraining

A zero vector is not "nothing": its score is 0, not . In the lab, padding without a mask moved the real tokens' outputs by 0.67; with the mask, by . Forgetting the document mask is the silent version of the same bug.

Where position comes from

Attention is permutation-equivariant: shuffle the tokens and the outputs shuffle with them. "dog bites man" = "man bites dog" (in the lab: a difference of ).

scheme mechanism verdict
sinusoidal fixed sin/cos added to the input historical
learned absolute one trainable vector per position cannot pass the trained length
ALiBi a linear distance penalty on the scores largely superseded
RoPE rotate Q and K by an angle ∝ position universal
NoPE / hybrid no signal in some layers growing in long-context stacks

RoPE: a clock per pair of dimensions

  • Query at position , key at position . Each pair of dimensions is a clock hand, turned (: the pair's speed). Rotations add angles, so the score depends only on .
  • Fast hands tell neighbours apart; slow hands, long distances.
  • Rescale the frequencies and the same weights address a longer context (position interpolation, NTK, YaRN).

01234567891011121314152θkqq at 3, k at 101234567891011121314152θkqq at 10, k at 8same angle between the hands → same score

Lab, section 6: the same q and k at positions (3, 1), (10, 8) and (1000, 998) give the same score, +0.033671; offsets 1 and 3 give −0.81 and +1.18.

Multi-head is a reshape

step shape
, one product for all heads
split into heads
scores, one per head
concatenate, then

Four matrices whatever is. Lab: → parameters (with the biases) for = 1, 2, 8 or 16. Heads are free parallelism, not extra capacity.

B: sequences in the batch · n: tokens · h: heads, each with d/h dimensions.

Simulator · multi-head, line by line

The KV cache

Generating token needs attention over all the earlier tokens. What do you need to keep from each of them?

  • Not the query: it only served its own token, back then. The key and the value are read by every later token.
  • Computed once, kept: without the cache step reruns the network over all tokens — generating tokens costs token passes instead of .
  • The price: the memory of K and V for every token, in every layer.

How big the cache is

term what it is Llama 3 8B
K and V 2
layers 32
key/value heads 8
dimensions per head 128
bytes per number 16 bits 2
per token bytes 128 KiB
tokens × sequences: 128k = 131,072 tokens, one user 16 GiB

Simulator · the cache on one accelerator

is the term you can shrink

scheme KV heads cache quality 2026
MHA 1× baseline historical at long context
MQA 1 measurably worse where memory dominates
GQA (≈ 8) ≈ MHA the default
MLA — ~1/57 reported ≥ MHA DeepSeek-V2/V3

Decoding reads the whole cache to produce one token: it is bound by memory bandwidth, not FLOPs. GQA, quantized caches, PagedAttention, bigger batches — all attacks on bytes moved.

Making affordable

  • Same math, better execution — FlashAttention never builds ; exact up to rounding. Use F.scaled_dot_product_attention.
  • Attend to less — sliding windows, attention sinks, interleaved global layers.
  • Change the math — linear attention, SSMs (Mamba): , but weaker exact recall. Hybrids win.
  • Learn what to skip — trainable sparse attention (NSA, MoBA).

Lab, section 8, one layer with 8 heads: S takes 8 · 4096² · 4 bytes = 512 MiB at n = 4096 and the work grows 4× per doubling; on the lab's CPU the naive version took 239 ms and the fused kernel 74 ms.

Measured in a real model

SmolLM2-135M (Llama architecture: RoPE, GQA): , , , .

  • Cache per token in float32 (4 bytes): 46,080 bytes — measured and predicted. With MHA, × more.
  • 80 tokens, greedy: the same tokens with and without the cache; without it, a step went from 55 to 95 ms (the prefix grows); with it, it stayed near 23 ms.
  • A trained previous-token head (layer 16, head 0): 0.78 of its weight on token .
  • The first token gets, on average, 0.45 of every row — sharing equally among the visible tokens would give 0.10: an attention sink.

Lab · the mechanism, computed and measured

Colab notebook, ~2 h, CPU. Each section ends in a checkpoint with a check value.

  1. the soft dictionary, by hand
  2. every query at once; against PyTorch
  3. why , measured
  4. the backward pass, by hand
  5. masks: causal and padding
  1. no order — and RoPE
  2. multi-head and GQA
  3. the bill
  4. the KV cache in SmolLM2-135M
  5. teach attention where to look

open in Colab →

Takeaways

  1. Attention is a dense layer with weights computed from the input — a soft dictionary. It beat recurrence on reach and parallelism.
  2. Q, K, V are three projections of the same input: what to ask, what to advertise, what to deliver.
  3. keeps the score variance at 1. Without it the softmax saturates and the gradient vanishes.
  4. The mask is one added matrix: the encoder/decoder distinction, padding, documents.
  5. Attention does not see order. RoPE: absolute rotations, relative offsets.
  6. Heads are a reshape. The KV cache dominates inference memory; is the term to shrink — GQA.
  7. FlashAttention is exact and free; the other answers to trade something.

Outline: (1) what recurrence could not do; (2) the idea: a soft dictionary, Q/K/V, one sum by hand; (3) the details that make it work: √d_k, masks, position, heads; (4) the cost: the KV cache and the O(n²). Closes with the Colab lab.

Defaults: 64 tokens, scale 1. RNN 5e-4, LSTM cell 4e-2. Drag to 256: 1e-14 and 2e-6. Both are straight lines on a log plot — exponential decay. The LSTM bought a better constant. Attention is flat at 1/n: small, but the same for every position.

Ask first: "how many tokens is 'the cat climbed the'?" Most say 4 words = 4 tokens; GPT-2 agrees here. Then the Portuguese line: 6 for 4 words, 25 for 13 words in the long sentence (1.92 per word against 1.06): more compute, less context, a bigger bill for the same content.

The unseen word is the point: "climbed" was never in the corpus and BPE writes it with two known pieces, ␣climb (from climb, climbs) and ed (from jumped, walked). Only 11 of the 40 merges touch the sentence. GPT-2's 50,000 merges include " climbed" as one token (19952).

Next through the merges: 19 characters, 5 pieces, ids [0, 2, 4, 5, 1]. Then the cases: "… roof." adds ␣roof and "."; "the bird climbed the roof." writes "bird" letter by letter (␣ b i r d, no id among the 12); the Portuguese sentence through the English tokenizer: 21 pieces.

X (5 x 4) is what attention receives. The same id always gets the same row: "bank" gets one vector in both sentences of the next slide. Making the vector depend on the neighbours is attention's job.

It starts at the numbers of the previous slide. Then: (a) value of river → [0, 1]: the verdict flips and no weight moves; (b) query → [1, 0]: the weight leaves river, 0.40/0.40/0.20; (c) query → 0: uniform, a plain average.

Read by row: where each token looks; every row sums to 1. The first three heads are patterns trained models really have (content, previous token, distance). Then: untick "divide by √d_k" and drag d_k up — the matrix goes binary. Then tick the causal mask: the upper triangle goes dark. Come back to the mask in four slides.

Variance 1/d_k: every score near zero, a near-uniform softmax that picks no one.

Lab, section 5: change the last token; without the mask all five rows change, with it only the last.

The lab's code (section 7, with the attention function of section 2), three tokens, d = 4, h = 2. Head 1 is the bank of slides 13-15: its row ends at [0.85, 0.15]; head 2 makes bank look at sat (0.77). Step line by line, then: "with the causal mask" - bank loses river and is back on the fence, [0.60, 0.40]; sat sees only itself. "a single head (h = 1)": same numbers, √4 = 2, bank splits 0.38/0.38/0.23 - one row cannot point at both. "all keys equal": every weight 1/3, a plain average.

Start at 8B, GQA, 32k, 8 requests: 34.4 GB of cache + 16 GB of weights, it fits. Switch to MHA: the cache grows 4× (137 GB) and does not fit. Back to GQA and up to 128k: 137 GB again. 70B or 405B: the weights alone (140 GB, 810 GB in 16 bits) do not fit on one 80 GB card.