Attention: the intuition

import math

the, dry, jacket = [0, 0], [1, 1], [2, 0]
words = [the, dry, jacket]

# Step 1, score: how relevant is each word to "dry"?
scores = [dry[0] * w[0] + dry[1] * w[1] for w in words]
print("scores: ", scores)

# Step 2, normalize: softmax turns the scores into weights
exps = [math.exp(s) for s in scores]
weights = [e / sum(exps) for e in exps]
print("weights:", [round(w, 2) for w in weights])

# Step 3, mix: the weighted average of the word vectors
new_dry = [sum(weights[j] * words[j][k] for j in range(3)) for k in range(2)]
print("new dry:", [round(v, 2) for v in new_dry])

Same word, new sentence

Now swap "jacket" for "toast", which is [0, 2]: pure food. This time the code uses NumPy and does every word at once.

X holds one row per word. X @ X.T is matrix multiplication, from Phase 2. It works out the dot product of every pair of words in one go. Row $i$, column $j$ of the result is the score $s_{ij}$. The next line is the softmax for each row. W @ X then does the mixing step for every word.

import numpy as np

X = np.array([[0.0, 0.0],    # the
              [1.0, 1.0],    # dry
              [0.0, 2.0]])   # toast
scores = X @ X.T                                  # every word scores every word
W = np.exp(scores) / np.exp(scores).sum(axis=1, keepdims=True)
new = W @ X                                       # every word's mixed vector
print("weights for dry:", W[1].round(2))
print("new dry:", new[1].round(2))

The weights for "dry" come out the same as before: 0.06, 0.47, 0.47. But the vectors it blends are different. So "dry" ends up at [0.47, 1.4], leaning toward food. Same word, same starting vector, a different meaning in a different sentence. That's attention's whole point.

import numpy as np
scores = np.array([1.0, 1.0, 1.0])
w = np.exp(scores) / np.exp(scores).sum()
print(np.round(w, 2))

Seeing it

Now try real words. The code below embeds each word of a sentence with the practice embedding model. Its 24 dimensions are topics, much like "rain gear" and "food" above. Then every word attends to every word.

Two details in the code are new.

• * 4 sharpens the scores. Multiplying every score by 4 widens the gaps between them. So softmax hands more weight to the top few words and less to the rest.
• plt.imshow draws a heatmap. It's a grid of colored squares. Each square shows one weight. Brighter means more attention.

import numpy as np
import requests
import matplotlib.pyplot as plt

words = "the waterproof jacket kept me dry on a wet muddy trail run".split()
data = requests.post("https://api.gradient.test/v1/embeddings", json={"input": words}).json()["data"]
X = np.array([d["embedding"] for d in data])          # 12 words × 24 dimensions

scores = X @ X.T * 4                                  # similarity of every pair (× 4 to sharpen)
W = np.exp(scores - scores.max(axis=1, keepdims=True))
W /= W.sum(axis=1, keepdims=True)                     # softmax: each row sums to 1

for w in ["jacket", "trail"]:
    i = words.index(w)
    top = np.argsort(-W[i])[:4]
    print(f"{w!r} attends to: " + ", ".join(f"{words[j]} {W[i, j]:.2f}" for j in top))

new = W @ X                                           # each word's context-mixed vector
print("\nshape before and after attention:", X.shape, new.shape)

plt.imshow(W, cmap="viridis")
plt.xticks(range(len(words)), words, rotation=60, fontsize=8)
plt.yticks(range(len(words)), words, fontsize=8)
plt.xlabel("attends to")
plt.ylabel("word")

Each row of the heatmap is one word's attention. The rain-gear words (waterproof, jacket, dry, wet) attend to each other. "Trail" and "run" attend to each other. The filler words, which this practice model can't tell apart, spread their attention among themselves.

The shape line shows that attention doesn't change the size of anything. There are 12 words with 24 numbers each, before and after. What changes is what's in those numbers. The vector for "dry" now carries information from "waterproof", "jacket" and "wet". It means "dry" in the context of rain gear.

What real attention adds

This toy uses raw embeddings and plain similarity. So it can only link words about the same topic. Real transformer attention, coming up over the next lessons, adds four things.

• It learns what to look for. It uses three learned projections, called queries, keys and values. A projection is a learned matrix that turns each vector into a new one. With them, "it" can look for a noun, a verb can look for its subject, and a closing bracket can look for the opening one. Those links aren't simple similarity.
• It has several heads. Each head is a separate attention pattern, and they run side by side. Each one is free to learn a different kind of link.
• It can't look ahead in a language model. The model's job is to predict the next token, so it must not peek at it. A causal mask hides every later token.
• It knows word order. Attention on its own can't tell which token came first. Positional information, added to the embeddings, tells it.

Why it won

• Any distance, one step. Token 1 and token 1,000 are linked directly. In a recurrent network, information passes through a long chain of hidden states and fades along the way. So do the gradients that train the model.
• All at once. Every position is computed at the same time. That suits GPUs, which are built for parallel work. A recurrent network has to finish each step before it starts the next.
• Driven by content. What gets blended depends on the input. So one mechanism handles countless kinds of links.

The price

Every token scores every other token. With 10 tokens, that's 10 × 10 = 100 scores. With 1,000 tokens, it's 1,000 × 1,000, a million scores. Double the text and the work goes up four times. In other words, the work grows with the square of the sequence length. That's the problem the long-context lesson tackles.

For marketers: say a model follows an instruction from the top of a long brief while it writes the last paragraph. Attention is what carries that instruction across the gap.