Notebook 1: Attention from Scratch

Words, tokens, and embeddings. In a transformer, input text is split into tokens — usually words or subword pieces. Each token is represented by an embedding: a numerical vector that encodes its meaning. In our example, the words "The", "cat", "sat", "down" are the tokens, and each row of $X$ is that token's 3-dimensional embedding. Real models use hundreds or thousands of dimensions, but the mechanics are identical.

Dot-product similarity. Each token's embedding is a row in $X$ and a column in $X^\top$. Entry $(i, j)$ of the similarity matrix is the dot product of row $i$ of $X$ with column $j$ of $X^\top$:

$$(\text{X} \cdot \text{X}^\top)_{ij} = \mathbf{x}_i \cdot \mathbf{x}_j = \|\mathbf{x}_i\| \, \|\mathbf{x}_j\| \cos\theta_{ij}$$

You can trace this visually: pick a row of $X$ on the left and the corresponding column of $X^\top$, dot them together, and you get the corresponding entry in the similarity matrix. The white lines help you track which row feeds into each entry.

Observe: The diagonal is largest — each token is most similar to itself. "cat" and "sat" have high mutual similarity (1.47) because their embeddings both have large values in $d_1$ and $d_2$, pointing in similar directions. "The" and "sat" are slightly anti-correlated (−0.10) because "The" points along $d_0$ while "sat" points away from it.

In attention, we use this to ask: "How relevant is each key to this query?"

From similarity to attention weights. Raw dot products can be any real number, but we need a probability distribution over keys for each query. The softmax function does this:

$$\text{attention}_{ij} = \text{softmax}\!\left(\frac{s_{ij}}{\tau}\right) = \frac{e^{\, s_{ij}/\tau}}{\displaystyle\sum_k e^{\, s_{ik}/\tau}}$$

where $s_{ij}$ is the raw similarity and $\tau$ is the temperature. This maps each row from $\mathbb{R}^n$ onto the probability simplex $\Delta^{n-1}$ (all entries non-negative, summing to 1).

Queries and keys. The axes on the attention heatmap are labeled "Queries" (rows) and "Keys" (columns). Think of the query as the token asking "who should I pay attention to?" and the keys as the candidates being considered. Each row of the attention matrix answers that question for one query token: it's a probability distribution over all the keys, saying how much each key matters to that query. We'll formalize this with separate $Q$/$K$/$V$ projections in Section 3, but the row-vs-column structure is already visible here.

Temperature controls how peaked the distribution is:

• Low temperature ($\tau \to 0$) → attention concentrates on the single most similar key (approaches one-hot). The model becomes very selective.
• High temperature ($\tau \to \infty$) → attention spreads uniformly across all keys (approaches $1/n$). The model treats all keys equally.
• $\tau = 1$ is the "unscaled" baseline shown by default. In transformers, the standard scaling uses $\tau = \sqrt{d_k}$, which for our 3D embeddings is $\sqrt{3} \approx 1.73$.

Try dragging the slider: at $\tau = 0.1$ the attention map is nearly binary; at $\tau = 5.0$ it's nearly uniform.

Why isn't the attention matrix symmetric? The similarity matrix $X X^\top$ is symmetric ($s_{ij} = s_{ji}$), but the attention matrix is not, because softmax normalizes each row independently. Entry $\text{attention}_{ij}$ depends not just on $s_{ij}$ but on all the scores in row $i$. Different rows have different normalization constants, so even though $s_{ij} = s_{ji}$, we generally have $\text{attention}_{ij} \neq \text{attention}_{ji}$.

This asymmetry is a fundamental feature of attention: each token gets its own perspective on the sequence. Note that the specific patterns here have no semantic significance — they're an artifact of our hand-chosen embeddings in a 3-dimensional space. In a trained model, the embeddings and separate $Q$/$K$ projections are learned so that attention patterns capture meaningful relationships.

Self-attention with $Q$, $K$, $V$ projections. This is called self-attention because queries, keys, and values all come from the same input:

$$Q = X W_Q, \quad K = X W_K, \quad V = X W_V$$

In cross-attention (which appears in encoder–decoder models), $K$ and $V$ come from a different sequence — but the mechanics are otherwise identical.

Why separate projections? In Section 1, the raw embedding $X$ served as both the question and the answer: we computed $X X^\top$ directly. This forces each token to use the same vector for "What am I looking for?" and "What do I advertise to others?" Separate projections decouple these roles:

• $Q$ (query): "What am I looking for?"
• $K$ (key): "What do I advertise to others?"
• $V$ (value): "What information do I provide when attended to?"

The same token can present a completely different face depending on which role it's playing.

Reading the diagram. Follow the flow from left to right:

1. The input $X$ branches three ways through learned weight matrices $W_Q$, $W_K$, $W_V$ to produce $Q$, $K$, $V$.
2. $Q$ and $K^\top$ multiply — $Q$ on the left, $K^\top$ below — to produce raw scores, scaled by $\sqrt{d_k}$ (the same temperature scaling from Section 2, with $\tau = \sqrt{d_k}$).
3. Softmax converts each row of scores into a probability distribution: the attention weights.
4. The attention weights multiply $V$, producing the output. Each row of the output is a weighted average of $V$'s rows, where the weights come from how well that token's query matched each key.

What the output means. Token $i$'s output vector isn't its own value — it's a blend of everyone's values, weighted by relevance. A token that strongly attends to one neighbor will have an output close to that neighbor's value vector. A token that spreads attention uniformly gets something like the average of all value vectors. This is how information flows between positions in a transformer.

What projections buy you. In Section 1, each token had a single embedding that determined all of its relationships. Projections break this constraint:

• Identity: $W_Q = W_K = I$ reproduces raw similarity exactly — projections don't have to change anything.
• Dimension focus: Zeroing out dimensions lets the model attend based on specific features. Focusing on $d_0$ makes "The" and "down" the dominant pair, completely changing the pattern from raw similarity where "cat"–"sat" dominated.
• Flipped similarity: $W_K = -I$ inverts all relationships. Tokens that were similar now repel. After softmax, each token attends most to its least similar neighbor.
• Random learned: A trained model discovers $W_Q$, $W_K$ that create useful patterns for the task — patterns that may bear no resemblance to raw similarity.

The key takeaway: the $\otimes$ symbols in the Section 3 diagram aren't just mechanical steps. They're where the model learns what to pay attention to, independently of what information to extract (which is $W_V$'s job).