Attention and embeddings

You’ll build attention from the four steps in the Engineer’s View: compare, scale, softmax, mix. You’ll use the same GloVe word vectors as Lesson 5, so the results are real, not made up. One simplification: a real model learns separate query, key and value tables, while here each word’s own vector plays all three parts.

That’s enough to see the idea working. Use PowerShell as before; on macOS or Linux, use python3 and cp.

1Set up the folder

Make a new folder and copy the 66 MB word-vector file from Lesson 5, so you don’t download it again:

Run this:

mkdir C:\ai-lab\attention
cd C:\ai-lab\attention
copy ..\embeddings\glove-wiki-gigaword-50.gz .
notepad glove.py

Paste this whole block into glove.py, then save:

import gzip, os, urllib.request
import numpy as np

URL = ("https://github.com/RaRe-Technologies/gensim-data/releases/download/"
       "glove-wiki-gigaword-50/glove-wiki-gigaword-50.gz")
FILE = "glove-wiki-gigaword-50.gz"

if not os.path.exists(FILE):
    print("Downloading word vectors (66 MB, only the first time)...")
    urllib.request.urlretrieve(URL, FILE)

words, rows = [], []
with gzip.open(FILE, "rt", encoding="utf-8") as f:
    next(f)
    for line in f:
        parts = line.rstrip().split(" ")
        words.append(parts[0])
        rows.append(np.array(parts[1:], dtype=np.float32))
        if len(words) == 50000:
            break

vecs = np.array(rows)
index = {w: i for i, w in enumerate(words)}

This file only loads the vectors and prints nothing. The other scripts in this lesson import it.

✅ Checkpoint: Your folder has glove.py and glove-wiki-gigaword-50.gz. If you skipped Lesson 5, the file downloads on the first run.

2Watch “windows” pay attention

Run this:

notepad attention.py

Paste this whole block, then save:

import numpy as np
from glove import vecs, index           # the same 50,000 GloVe words as Lesson 5

WORD = "windows"
SENTENCE = "microsoft windows update failed on the laptop"
COMPARE = ("software", "glass")


def softmax(x):
    e = np.exp(x - x.max())
    return e / e.sum()


def cos(a, b):
    return float(a @ b / (np.linalg.norm(a) * np.linalg.norm(b)))


others = [w for w in SENTENCE.split() if w != WORD]
query = vecs[index[WORD]]                       # what WORD is looking for
keys = np.array([vecs[index[w]] for w in others])   # what each other word offers
values = keys                                   # what each word passes on

scores = keys @ query / np.sqrt(len(query))     # 1. compare, 2. scale
weights = softmax(scores)                       # 3. turn into percentages
new_vector = query + weights @ values           # 4. mix in, add to the original

print("Attention from", repr(WORD) + ":")
for w, s, p in zip(others, scores, weights):
    print(f"  {w:<10} score {s:5.2f}   weight {p:4.0%}")
print()
a, b = COMPARE
print(f"Alone:      {a} {cos(query, vecs[index[a]]):.2f}   {b} {cos(query, vecs[index[b]]):.2f}")
print(f"In context: {a} {cos(new_vector, vecs[index[a]]):.2f}   {b} {cos(new_vector, vecs[index[b]]):.2f}")

Run this:

python attention.py

The text version of the same output:

Attention from 'windows':
  microsoft  score  3.13   weight  34%
  update     score  2.00   weight  11%
  failed     score  1.32   weight   6%
  on         score  2.21   weight  14%
  the        score  1.99   weight  11%
  laptop     score  2.79   weight  24%

Alone:      software 0.65   glass 0.58
In context: software 0.77   glass 0.52

Read the weights. “windows” listens most to microsoft (34%) and laptop (24%), and least to “failed” (6%). After mixing, it’s closer to software (0.65 → 0.77) and further from glass (0.58 → 0.52).

✅ Checkpoint: microsoft gets the biggest weight, and the software score goes up from 0.65 to 0.77.

3Change the sentence and watch the meaning flip

Change only the SENTENCE line near the top of attention.py:

Change this line:

SENTENCE = "she cleaned the windows and doors of the kitchen"

Run this:

python attention.py

You should see this:

Attention from 'windows':
  she        score  1.26   weight   5%
  cleaned    score  1.68   weight   8%
  the        score  1.99   weight  11%
  and        score  1.85   weight   9%
  doors      score  3.09   weight  33%
  of         score  1.56   weight   7%
  the        score  1.99   weight  11%
  kitchen    score  2.32   weight  15%

Alone:      software 0.65   glass 0.58
In context: software 0.61   glass 0.68

Now “windows” listens most to doors (33%) and kitchen (15%), and the meaning flips: glass wins, 0.68 to 0.61. This is exactly what failed in Lesson 5’s analogy. Attention fixes it, because the word now reads its neighbours.

Try one more. Change the first three settings to:

Change these three lines:

WORD = "apple"
SENTENCE = "apple pie with banana and fresh cream"
COMPARE = ("microsoft", "fruit")

Run this:

python attention.py

You should see this:

Attention from 'apple':
  pie        score  1.88   weight  16%
  with       score  1.63   weight  12%
  banana     score  2.04   weight  19%
  and        score  1.51   weight  11%
  fresh      score  2.09   weight  20%
  cream      score  2.24   weight  23%

Alone:      microsoft 0.73   fruit 0.59
In context: microsoft 0.58   fruit 0.76

On its own, “apple” is closer to Microsoft (0.73) than to fruit (0.59), probably because the company is mentioned so often in the news text these vectors learned from. In a dessert sentence, fruit wins clearly: 0.76 against 0.58.

✅ Checkpoint: In the kitchen sentence glass beats software, and in the pie sentence fruit beats microsoft.

4See the causal mask

Now look at a whole sentence at once. Every word attends to every word allowed by the mask:

Run this:

notepad mask.py

Paste this whole block, then save:

import numpy as np
from glove import vecs, index

words = "the server crashed last night".split()
x = np.array([vecs[index[w]] for w in words])

scores = x @ x.T / np.sqrt(x.shape[1])
future = np.triu(np.ones((len(words), len(words)), dtype=bool), 1)
scores[future] = -np.inf                    # the causal mask: no peeking ahead

weights = np.exp(scores - scores.max(axis=1, keepdims=True))
weights = weights / weights.sum(axis=1, keepdims=True)

print(" " * 10 + "".join(f"{w:>9}" for w in words))
for w, row in zip(words, weights):
    print(f"{w:<10}" + "".join(f"{p:>9.0%}" for p in row))

Run this:

python mask.py

You should see this:

                the   server  crashed     last    night
the            100%       0%       0%       0%       0%
server           4%      96%       0%       0%       0%
crashed          9%       2%      89%       0%       0%
last            23%       2%      10%      65%       0%
night           13%       2%       8%      23%      53%

Each row is one word, and the row shows where its attention goes. Everything above the diagonal is 0%: “the” can only see itself, “server” can see “the” and itself, and “night” can see the whole sentence. Each row adds up to 100%.

Here every word gives most of its attention to itself, because we used the raw word vectors for queries and keys. A trained model has separate tables, so it can learn to look elsewhere.

✅ Checkpoint: Every number above the diagonal is 0%, and every row adds up to 100%.

5Measure why long prompts are slow

Run this:

notepad cost.py

Paste this whole block, then save:

import time
import numpy as np

rng = np.random.default_rng(0)
head_dim = 64

print(f"{'tokens':>7} {'scores per head':>16} {'time':>9}")
for n in (512, 1024, 2048, 4096):
    q = rng.normal(size=(n, head_dim)).astype(np.float32)
    k = rng.normal(size=(n, head_dim)).astype(np.float32)
    start = time.perf_counter()
    scores = q @ k.T                       # every token against every token
    ms = (time.perf_counter() - start) * 1000
    print(f"{n:>7,} {scores.size:>16,} {ms:>7.1f} ms")

Run this:

python cost.py

You should see something like this (your times depend on your laptop):

 tokens  scores per head      time
    512          262,144     1.0 ms
  1,024        1,048,576     2.4 ms
  2,048        4,194,304     7.6 ms
  4,096       16,777,216    25.8 ms

The middle column is exact: double the tokens and the scores go up four times. The times on your laptop will be different, but they’ll grow in roughly the same way, about three to four times for each doubling once the numbers get big.

This is for one head in one layer. Your Llama model has 32 heads in each of 16 layers, which is why Ollama’s prompt eval time from Lesson 2 jumps when you paste a long document.

✅ Checkpoint: You can explain why going from 2,048 to 4,096 tokens means four times as many scores.