Attention from scratch: a step-by-step tutorial in NumPy
Watch the video: Attention, built from scratch. This article is the slow version: every number worked by hand, every line of code explained.
TL;DR
- The problem: "it" enters a model with the same numbers in every sentence. Whatever tells the model that this "it" is the cat has to come from the other words. Attention is how it looks.
- Three steps, nothing else: each word's query is compared with every word's key by a
dot product (
Q @ K.T); softmax turns each row of scores into weights that add up to 1; the weights blend the values (weights @ V). - Divide by √d before softmax. With random vectors of length 512, the top key takes about 96% of the weight on average without it, and about 32% with it.
- You can design a question by hand. In "the cat sat on the mat because it was tired", "it" looks for an animal and a noun. "Cat" is both, so it gets 92% of the attention, "mat" 5.4%.
- Attention alone ignores word order. Shuffle the words and each word's output just moves with it, so positions have to be added separately (positional encoding, the next unit).
- Run it yourself: one script at the end prints every number in this article.
Read this sentence:
The cat sat on the mat because it was tired.
You knew "it" was the cat before you reached the full stop. Not the mat: mats don't get tired.
A language model has to work that out with arithmetic, and it starts with a handicap. Each word enters the model as a list of numbers looked up from a table (its token embedding), and the list for "it" is the same in every sentence it appears in. "It was tired", "it was raining", "it compiles on my machine": the same starting numbers. Whatever tells the model that this "it" is a cat has to come from the other words.
The part of a transformer that does that is called attention. By the end of this tutorial you'll have built it in about a dozen lines of Python, worked every number by hand, and made "it" find the cat.
What you need: some Python, and NumPy arrays (import numpy as np). No machine learning, no
calculus. Every number below comes from one short script, printed in full at the end, so you can
run it and see the same digits.
Part 1: four tools you'll use
Attention is built from four small tools: a vector, the dot product, the transpose and matrix multiply. If you know them, skip ahead to Part 2.
A word is a list of numbers
Inside a model, every word becomes a vector: a list of numbers, hundreds or thousands long. Real models learn what those numbers mean, and no human can read them. In this tutorial we'll cheat and use vectors we can read. For the first example, two numbers per word:
| Word | Vector |
|---|---|
| dog | [1, 0] |
| bites | [0, 1] |
| man | [1, 1] |
The dot product: how much two vectors agree
Multiply two vectors position by position, then add the results up:
[1, 0] · [1, 0] = 1×1 + 0×0 = 1 same direction: a match
[1, 0] · [0, 1] = 1×0 + 0×1 = 0 unrelated: nothing
[1, 1] · [1, 1] = 1×1 + 1×1 = 2 same direction, and bigger: a stronger match
[1, 0] · [-1, 0] = -1 opposite directions: a negative score
That's the whole intuition you need: the dot product is big when two vectors point the same way. Attention is, at heart, a lot of dot products.
A matrix is vectors stacked up
Put the three word vectors on top of each other and you have a matrix: a grid of numbers. Call it
X. Each row is one word:
X = np.array([[1, 0], # row 0: dog
[0, 1], # row 1: bites
[1, 1]]) # row 2: man
Its shape is 3 × 2: three rows (the words), two columns (the numbers in each word).
Transpose: turn the rows into columns
X.T (the T is for transpose) flips the grid on its diagonal, so the rows become columns. Now
each column is one word:
X.T = [[1, 0, 1], dog bites man
[0, 1, 1]] ↑ ↑ ↑ each column is one word's vector
Its shape is 2 × 3. Nothing is calculated; the same six numbers are just laid out the other way round.
Matrix multiply: a table of dot products
A @ B multiplies two matrices, and the rule for it is one sentence:
The number in row i, column j of the result is row i of A dotted with column j of B.
So @ always pairs the rows of the left with the columns of the right. That's why we need the
transpose: in X the words are rows, and to pair every word with every word we need a copy where the
words are columns. X @ X.T then dots every word (a row of X) with every word (a column of X.T).
Work one cell by hand. Row 2, column 0 is "man" dotted with "dog":
row 2 of X = [1, 1] (man)
column 0 of X.T = [1, 0] (dog)
[1, 1] · [1, 0] = 1×1 + 1×0 = 1
Do that for all nine pairs and you get the whole table:
X @ X.T
# [[1, 0, 1], dog·dog, dog·bites, dog·man
# [0, 1, 1], bites·dog, bites·bites, bites·man
# [1, 1, 2]] man·dog, man·bites, man·man
The shape rule follows from the sentence above: an (n × d) matrix times a (d × m) matrix gives (n × m), and the two inner sizes (both d) have to match, because a dot product needs two vectors of the same length. Here (3 × 2) @ (2 × 3) gives 3 × 3: one score for every pair of words.
One more NumPy habit you'll see: axis=-1 means "along each row", and keepdims=True keeps the
result as a column, so it lines up when you divide every row by its own total. You'll need both for
softmax.
Part 2: the idea, query, key and value
Here's attention in one picture. Every word gets three new vectors, made from its own:
- a query: what this word is looking for;
- a key: what this word offers to anyone looking;
- a value: what this word hands over if it's picked.
Then every word shines its query at every word's key, itself included. Where a query and a key agree (a big dot product), the light is bright. Each word then collects the values of the words it lit up, in proportion to how brightly.
In the llmmart videos that's a spotlight: one word asks, light falls on every word, and the brightness is the weight. A real model does every word's asking at the same time; the spotlight shows them one at a time so you can follow.
Why three vectors and not one? Because what a word looks for is not what it is. "It" is looking for an animal, but "it" is not an animal. If a word used one vector for both jobs, "it" would mostly find other pronouns. Separate queries and keys let a word ask about one thing while advertising another. The value is separate again, because what you pass on needn't be the thing that got you noticed.
Part 3: one full pass, by hand
Let's run attention on "dog bites man" with the vectors above. To keep the first pass simple, each word's query, key and value are just its own vector. (Part 5 changes that.)
Step 1: score every pair
Each word has a query and a key. Stack them the same way as the words:
Qis all the queries, one row per word (3 × 2);Kis all the keys, one row per word (3 × 2).
A score is one word's query dotted with one word's key. We want that for every pair, and Part 1 already showed how to get every pair at once: put one set as rows and the other as columns, then multiply. So the scores are
scores = Q @ K.T
Row i of Q is word i's query; column j of K.T is word j's key. So the number in row i, column j
is "how well word i's question matches word j's answer". Read the table as row asks, column
answers.
In this first pass the queries and keys are just the word vectors themselves, so Q and K both
equal X, and Q @ K.T is exactly the X @ X.T table from Part 1. Two cells by hand, to see it:
man asks, man answers: query of man [1, 1] · key of man [1, 1] = 2
dog asks, bites answers: query of dog [1, 0] · key of bites [0, 1] = 0
The whole table:
| asks ↓ / answers → | dog | bites | man |
|---|---|---|---|
| dog | 1 | 0 | 1 |
| bites | 0 | 1 | 1 |
| man | 1 | 1 | 2 |
Step 2: divide by the square root of d
d is the length of each query, here 2, so every score is divided by √2 ≈ 1.414:
| asks ↓ / answers → | dog | bites | man |
|---|---|---|---|
| dog | 0.7071 | 0 | 0.7071 |
| bites | 0 | 0.7071 | 0.7071 |
| man | 0.7071 | 0.7071 | 1.4142 |
Why do this at all? It looks like decoration, and it isn't. Part 7 measures what happens without it.
Step 3: softmax, from scores to weights
Scores can be any size and any sign. We want weights: positive numbers that add up to 1, so each
word has one unit of attention to spend across the sentence. Softmax does it in two moves:
raise e to the power of each score, then divide each by the row's total. (e is a fixed number, about
2.718; "e to the power of x" is np.exp(x). It turns any score into a positive number, and it
stretches the gaps: a slightly bigger score becomes a noticeably bigger number.)
Take the "man" row, [0.7071, 0.7071, 1.4142]:
e^0.7071 = 2.0281 e^0.7071 = 2.0281 e^1.4142 = 4.1133
total = 8.1695
weights = 2.0281/8.1695, 2.0281/8.1695, 4.1133/8.1695 = 0.2483, 0.2483, 0.5035
"Man" spends half its attention on itself and a quarter each on "dog" and "bites". The whole table:
| asks ↓ / answers → | dog | bites | man |
|---|---|---|---|
| dog | 0.4011 | 0.1978 | 0.4011 |
| bites | 0.1978 | 0.4011 | 0.4011 |
| man | 0.2483 | 0.2483 | 0.5035 |
Every row sums to 1.
The trap. np.exp(1000) is too big for a computer number, so NumPy gives infinity, and infinity
divided by infinity is nan ("not a number"). Real scores can
get large, so every softmax first subtracts each row's biggest score. It changes nothing: softmax of
[1000, 1001, 1002] and of [0, 1, 2] are both [0.09, 0.2447, 0.6652], because subtracting the
same number from every score divides every exponential by the same factor, and the division at the
end cancels it out. PyTorch and SciPy do the same subtraction inside their softmax.
Step 4: the weighted sum of values
Each word's output is its weights times the values. For "man":
0.2483 × [1, 0] + 0.2483 × [0, 1] + 0.5035 × [1, 1] = [0.7518, 0.7518]
dog bites man
(With every decimal kept, the weights give 0.7517; the last digit is rounding.)
That's the point of the whole thing: each word's output is a blend of the words it paid
attention to. "Man" went in as [1, 1] and came out as a mix of itself and the other words. Do this
again and again, layer after layer, with learned vectors, and a word's numbers end up carrying what
its context means.
Part 4: the same thing in code
Here are the four steps as code. Each line maps to a step above.
def softmax(scores):
shifted = scores - scores.max(axis=-1, keepdims=True) # the trap: subtract each row's max
e = np.exp(shifted) # make everything positive
return e / e.sum(axis=-1, keepdims=True) # each row now sums to 1
def attention(Q, K, V, scale=True):
scores = Q @ K.T # step 1: every query · every key
if scale:
scores = scores / np.sqrt(Q.shape[-1]) # step 2: divide by √d
weights = softmax(scores) # step 3
return weights @ V, weights # step 4: blend the values
Shapes, for three words with vectors of length 2: Q and K are 3 × 2, so Q @ K.T is 3 × 3 (a
score for every pair). weights is also 3 × 3, and weights @ V is 3 × 2: one blended vector per
word, the same size as the words that went in.
What weights @ V does. V holds every word's value, one row per word. Use the rule from Part 1
on one row of the result: row "man" of weights is [0.2483, 0.2483, 0.5035], and the first column
of V is the first number of each word's value, [1, 0, 1]. Their dot product is
0.2483×1 + 0.2483×0 + 0.5035×1 = 0.7518, the first number of the blend you worked out in step 4.
The second column gives the second number. So weights @ V is step 4 for every word at once.
Q.shape[-1] is d, the length of one query: shape is (3, 2), and [-1] takes the last number.
Part 5: where queries, keys and values come from
In Part 3 every word's query, key and value were its own vector. In a real model, three matrices turn a word into its three vectors:
def self_attention(X, Wq, Wk, Wv):
return attention(X @ Wq, X @ Wk, X @ Wv)
What X @ Wq does. Each row of X is a word. By the rule from Part 1, row i of X @ Wq is word i
times Wq, so X @ Wq turns every word into its query in one go: Q, one query per row. X @ Wk
gives every key (K), and X @ Wv every value (V). The shapes have to line up: if each word has
6 numbers and each query should have 2, Wq is 6 × 2, and (words × 6) @ (6 × 2) gives (words × 2).
Part 3 was the special case where Wq changes nothing (a matrix with 1s down the diagonal and 0s
everywhere else, called the identity), so each query came out equal to its word.
"Self" means the queries, keys and values all come from the same words: the sentence attends to
itself. Wq, Wk and Wv are where the model's knowledge lives. In a trained model they're learned:
nobody writes them. Which is exactly why it's worth writing one by hand, once.
Part 6: make "it" pay attention to "cat"
Parts 3 and 4 showed how attention uses queries, keys and values. Part 5 showed how matrices make those vectors. Now we'll choose the numbers in two of those matrices ourselves.
The goal is small: give "cat" a higher score than "mat" when "it" asks. We are building a rule that prefers animal nouns, then watching attention carry out that rule.
This is a teaching example. We choose the features and the rule; a real model learns its matrices during training. This example does not discover what "it" means, and it does not use "tired" to reason about the sentence. If there were two animal nouns, this rule alone could not choose between them.
Step 1: describe each word with six yes-or-no features
Use the same sentence:
The cat sat on the mat because it was tired.
For this example, each word gets six numbers in this order:
[noun, animal, pronoun, verb, plural, tired]
A 1 means yes; a 0 means no. These are labels we supply for the example.
| Word | noun | animal | pronoun | verb | plural | tired |
|---|---|---|---|---|---|---|
| cat | 1 | 1 | 0 | 0 | 0 | 0 |
| mat | 1 | 0 | 0 | 0 | 0 | 0 |
| it | 0 | 0 | 1 | 0 | 0 | 0 |
| sat, was | 0 | 0 | 0 | 1 | 0 | 0 |
| tired | 0 | 0 | 0 | 0 | 0 | 1 |
| the, on, because | 0 | 0 | 0 | 0 | 0 | 0 |
For example, [1, 1, 0, 0, 0, 0] says that "cat" is a noun and an animal.
The "tired" feature labels the word "tired"; it does not mark the cat as tired.
Step 2: decide what the two query and key numbers mean
The input has six numbers per word. We'll turn those into two numbers per query or key. Give those two positions these jobs:
| Position | In a query: what to look for | In a key: what the word offers |
|---|---|---|
| 0, the first number | How strongly to look for an animal | Is this word an animal? |
| 1, the second number | How strongly to look for a noun | Is this word a noun? |
These meanings are our design for this example. Real learned query and key positions need not have neat names like these.
We want "it" to look strongly for both, so give it query [4, 4].
Give "cat" key [1, 1] and "mat" key [0, 1].
That already tells us who will score higher:
animal match noun match total
it asks cat: 4 × 1 + 4 × 1 = 8
it asks mat: 4 × 0 + 4 × 1 = 4
"Cat" matches both parts of the query; "mat" matches only the noun part.
Pause here: the six-number input describes the word. The two-number query describes what it looks for. The two-number key describes what can match that query. They have different jobs.
Step 3: put that design into Wq and Wk
Wq turns six input features into two query numbers. Wk turns the same six input features
into two key numbers. Each matrix therefore has six rows and two columns.
Start with zeros, then set four entries:
Wq = np.zeros((6, 2))
Wk = np.zeros((6, 2))
Wq[2, 0] = 4 # If the input is a pronoun, look for an animal with strength 4.
Wq[2, 1] = 4 # If the input is a pronoun, look for a noun with strength 4.
Wk[1, 0] = 1 # If the input is an animal, put 1 in the key's animal position.
Wk[0, 1] = 1 # If the input is a noun, put 1 in the key's noun position.
Python counts from zero. So feature 2 is pronoun, feature 1 is animal, and feature 0 is noun.
In Wq[2, 0], the 2 chooses an input feature; the 0 chooses a position in the output query.
Those indices refer to different things.
Here are the complete matrices. The labels on the left explain the rows; they aren't code:
input feature Wq Wk
animal noun animal noun
noun [0, 0] [0, 1]
animal [0, 0] [1, 0]
pronoun [4, 4] [0, 0]
verb [0, 0] [0, 0]
plural [0, 0] [0, 0]
tired [0, 0] [0, 0]
How does multiplication use these rows? Each input feature multiplies its corresponding matrix row, and the results are added. A 1 includes a row; a 0 contributes nothing.
"It" has only the pronoun feature, so its query takes only the pronoun row of Wq:
it's query = [4, 4]
"Cat" has the noun and animal features, so its key adds those two rows of Wk:
cat's key = noun row + animal row = [0, 1] + [1, 0] = [1, 1]
mat's key = noun row = [0, 1]
Every other word in this sentence has key [0, 0], because none has the noun or animal feature
in our table. X @ Wq and X @ Wk do this for every word at once.
Step 4: turn the scores into attention weights
The scores are 8 for "cat", 4 for "mat", and 0 for each of the other eight word positions. There are ten positions: "the" appears twice.
Use the same two steps as Part 3. First divide by √2, because each query has two numbers:
cat: 8 / √2 ≈ 5.657
mat: 4 / √2 ≈ 2.828
each other position: 0 / √2 = 0
Then softmax turns those scores into shares that add up to 1:
| Where "it" pays attention | Weight, approximately |
|---|---|
| cat | 0.920, or 92% |
| mat | 0.054, or 5.4% |
| Each of the other eight positions | 0.0032, or 0.32% |
The displayed numbers are rounded. A zero score still gets a small positive weight: softmax
uses exp(0) = 1, so it does not remove that word.
Step 5: collect the values
Queries and keys decided how much to take from each word. Values supply what to take.
To make the result easy to inspect, we give each word position its own ten-number marker. This is called a one-hot vector: one number is 1, all the others are 0.
the cat sat on the mat because it was tired
cat's value: [0, 1, 0, 0, 0, 0, 0, 0, 0, 0]
mat's value: [0, 0, 0, 0, 0, 1, 0, 0, 0, 0]
Multiply each marker by its attention weight and add them. "Cat" contributes 0.920 in the cat position; "mat" contributes 0.054 in the mat position. Each other position gets its small share. With these particular values, the output simply reproduces the attention weights.
That lets us see where the information came from. The 92% is an attention weight, not a 92% probability that the model has correctly understood the sentence. Real values carry learned information rather than these position markers.
Notice the three sizes: six numbers describe each input word, two numbers make each query or key, and ten numbers make each value here. Only queries and keys need the same length for their dot product; values can have a different length.
Try changing one thing
Change the numbers in Wq and watch the attention on "cat" (in the script at the end, that's
it_looks_at()):
| Change | Attention on cat | What it shows |
|---|---|---|
Original design: query [4, 4] |
92.0% | Cat matches both animal and noun. |
Remove the noun match: set Wq[2, 1] = 0 |
65.3% | Cat still wins, but it loses its extra noun score. |
| Change both query strengths from 4 to 1 | 29.1% | Cat still ranks first, but the weights are more spread out. |
| Change both strengths from 4 to 5 (the design in the video) | 96.5% | A little louder, a little more decisive. |
| Change both strengths from 4 to 8 | 99.6% | The larger score gaps make softmax concentrate on cat. |
What to take away: we wrote a preference into Wq and Wk; attention used that preference
to score the words, turn the scores into weights, and collect a weighted mixture of values.
Part 7: why divide by the square root of d?
In Part 6, changing the query from [4, 4] to [8, 8] made attention much more concentrated.
The matching words stayed the same. Only the score gaps grew.
That raises a problem: the size of the numbers can make attention concentrate even when the numbers contain no useful pattern.
Longer random vectors can produce larger score gaps
A dot product adds one product for every position in the vectors:
[a, b] · [c, d] = a×c + b×d
A query of length 512 adds 512 products. With random numbers, these products can be positive or negative, so they partly cancel. Longer vectors do not make every score larger, but they do make the scores more spread out.
For the independent random numbers used here, that spread grows roughly with √d, where d is the number of entries in a query or key. It is not the number of words.
Softmax reacts to those larger gaps
Compare these two sets of scores:
scores [1, 2, 3] → weights about [0.090, 0.245, 0.665]
scores [10, 20, 30] → weights about [0.000000002, 0.000045, 0.999955]
The last entry wins in both. In the second, it takes almost everything.
Here's a measurement, with random queries and ten random keys, over 500 trials:
| Query/key length d | Average largest weight without division | With division by √d |
|---|---|---|
| 4 | 0.501 | 0.307 |
| 64 | 0.855 | 0.317 |
| 512 | 0.958 | 0.318 |
At length 512, the winning key gets about 96% on average without scaling. These are random vectors: winning does not show that a key is meaningful. Individual trials differ.
The division compensates for vector length
Divide each dot-product score by √d before softmax:
scores = (Q @ K.T) / np.sqrt(Q.shape[-1])
This reduces the extra spread caused by having more query/key entries. In the experiment, the average largest weight stays near 0.31 across the three lengths.
Why does this help training? When softmax is extremely concentrated, small score changes often barely change its weights. Its derivatives become small, weakening the learning signal through that operation. Scaling helps avoid this problem; it still allows the model to learn strong preferences.
What to take away: divide by √d so longer vectors do not, by themselves, make attention almost all-or-nothing. This is the reasoning behind scaled dot-product attention in the original Transformer paper, section 3.2.1.
Part 8: what this attention code doesn't know yet
Our code compares word features. We haven't given it word positions.
Consider:
dog bites man
man bites dog
Those sentences mean different things. But in our code, "dog" starts with the same vector in both, and so do "bites" and "man". Each word has the same queries, keys and values available to compare, just in a different order.
So "dog" gets the same output vector in both sentences. Its output moves to a different row, but its numbers do not change. The same is true for the other words. Reordering rows alone doesn't tell this attention operation who bit whom.
We need to supply information about position as well. Positional encoding does that, and it's the next unit: Why Transformers Need Positional Encoding (Dog Bites Man). This limitation describes the unmasked, position-free attention we built here.
Before moving on, try saying the whole process without the code:
Make a query for each word. Compare it with the keys. Turn the scores into weights. Use those weights to blend the values. Supply position information if word order should matter.
Check yourself
Try each before reading the answer.
- Why does attention use a separate query and key, instead of one vector per word? Because what a word looks for isn't what it is. "It" looks for an animal without being one.
- A row of the weights table sums to what, and why? To 1: softmax divides by the row's total, so each word spends exactly one unit of attention.
- What's the shape of
Q @ K.Tfor a 10-word sentence? 10 × 10: one score for every pair of words. - What breaks if you skip "subtract the max" in softmax? Large scores overflow:
exp(1000)is infinity and the weights come out asnan. The answer doesn't change when you subtract it. - You doubled every number in
Wq. What happens to the attention on "cat", and why? It gets more decisive (closer to 1): the scores double, and bigger scores make softmax sharper.
Common mistakes
- Softmax gives
nan: you didn't subtract each row's max beforenp.exp. - NumPy says the shapes "could not be broadcast together": the row totals lost their column
shape. Add
keepdims=Trueto themaxand thesum. - "Man"'s weights come out 0.2119, 0.2119, 0.5761 instead of 0.2483, 0.2483, 0.5035: you didn't divide by √d.
- NumPy refuses the multiply in self-attention, saying the sizes differ: usually
Wq @ Xinstead ofX @ Wq. The word comes first: one row per word.
The rule to steal: attention is a dot product, a softmax and a weighted sum, and nothing else. Here we wrote the preference into
WqandWkby hand; a trained model learns its own, with features nobody labelled, but the three steps are the same.
Run every number yourself
This script prints every number in this tutorial. Paste it into a file and run it with NumPy
installed (pip install numpy, or uv run --with numpy python attention_tutorial.py):
import numpy as np
np.set_printoptions(precision=4, suppress=True)
def softmax(scores):
shifted = scores - scores.max(axis=-1, keepdims=True) # subtract each row's max: no overflow
e = np.exp(shifted) # every score becomes positive
return e / e.sum(axis=-1, keepdims=True) # each row now sums to 1
def attention(Q, K, V, scale=True):
scores = Q @ K.T # every query · every key
if scale:
scores = scores / np.sqrt(Q.shape[-1]) # divide by √d
weights = softmax(scores)
return weights @ V, weights # blend the values
def self_attention(X, Wq, Wk, Wv):
return attention(X @ Wq, X @ Wk, X @ Wv)
# Part 3: "dog bites man", each word's own vector as its query, key and value
X = np.array([[1.0, 0.0], # dog
[0.0, 1.0], # bites
[1.0, 1.0]]) # man
print("scores\n", X @ X.T)
print("divided by √2\n", X @ X.T / np.sqrt(2))
out, w = attention(X, X, X)
print("weights\n", w)
print("man's output", out[2])
# The softmax trap
with np.errstate(all="ignore"):
big = np.array([1000.0, 1001.0, 1002.0])
print("without the shift", np.exp(big) / np.exp(big).sum())
print("with the shift ", softmax(big[None, :])[0], " same as softmax([0, 1, 2]):", softmax(np.array([[0.0, 1.0, 2.0]]))[0])
# Part 6: make "it" look at "cat". Features: noun, animal, pronoun, verb, plural, tired
FEATURES = {"the": [0, 0, 0, 0, 0, 0], "cat": [1, 1, 0, 0, 0, 0], "sat": [0, 0, 0, 1, 0, 0], "on": [0, 0, 0, 0, 0, 0],
"mat": [1, 0, 0, 0, 0, 0], "because": [0, 0, 0, 0, 0, 0], "it": [0, 0, 1, 0, 0, 0], "was": [0, 0, 0, 1, 0, 0],
"tired": [0, 0, 0, 0, 0, 1]}
sentence = "the cat sat on the mat because it was tired".split()
S = np.array([FEATURES[w] for w in sentence], dtype=float)
def it_looks_at(strength=4.0, noun=True):
Wq, Wk = np.zeros((6, 2)), np.zeros((6, 2))
Wq[2, 0] = strength # a pronoun looks for an animal...
Wq[2, 1] = strength if noun else 0.0 # ...and for a noun
Wk[1, 0] = 1.0 # an animal answers in position 0
Wk[0, 1] = 1.0 # a noun answers in position 1
# each word's value is a one-hot marker of its own position, so the output shows where attention went
_, weights = attention(S @ Wq, S @ Wk, np.eye(len(sentence)))
row = weights[sentence.index("it")]
return {"cat": round(float(row[1]), 3), "mat": round(float(row[5]), 3), "each other word": round(float(row[0]), 4)}
print("design (4, 4): ", it_looks_at())
print("no noun match: ", it_looks_at(noun=False))
print("strength 1: ", it_looks_at(1.0))
print("strength 8: ", it_looks_at(8.0))
# Part 7: why divide by √d. Random queries, ten random keys, 500 tries
rng = np.random.default_rng(3)
for d in (4, 64, 512):
for scale in (False, True):
tops = [attention(rng.normal(size=(1, d)), rng.normal(size=(10, d)), np.eye(10), scale=scale)[1].max() for _ in range(500)]
print(f"d={d:3d} {'scaled ' if scale else 'unscaled'} average top weight {np.mean(tops):.3f}")
# Part 8: shuffle the words, and each word's output just moves with it
rng = np.random.default_rng(2)
Xr, Wq, Wk, Wv = rng.normal(size=(5, 8)), rng.normal(size=(8, 4)), rng.normal(size=(8, 4)), rng.normal(size=(8, 4))
out, _ = self_attention(Xr, Wq, Wk, Wv)
order = [4, 2, 0, 3, 1]
shuffled, _ = self_attention(Xr[order], Wq, Wk, Wv)
print("shuffled outputs equal the original outputs, reordered:", np.allclose(shuffled, out[order]))
Want to build it step by step? The llmmart curriculum has this unit as an exercise: the same functions with the key lines left for you to write, and checks that tell you when each step is right. Unit 08 on attention.
Comments (0)
Sign in to join the conversation.
No comments yet.