§3.3Simple Self-Attention Without Trainable Weights

This section computes attention by hand, with zero trainable weights — the purest form of the mechanism. Once the moving parts are visible, §3.4 will add the trainable matrices, and nothing conceptual will change.

Recall the “self” from the previous section: the mechanism relates positions within one sequence to each other — every word of a sentence asks “how relevant is every other word to me?” — in contrast to traditional attention, which connects two different sequences (like the German input and English output of a translator).

3.3.1 One query, step by step

The running example is the six-token sentence “Your journey starts with one step”, where each token already has a 3-dimensional embedding from chapter 2’s pipeline (the book uses 3 dimensions so every number fits on screen; GPT-2 uses 768):

import torch
inputs = torch.tensor(
[[0.43, 0.15, 0.89], # Your     (x^1)
 [0.55, 0.87, 0.66], # journey  (x^2)
 [0.57, 0.85, 0.64], # starts   (x^3)
 [0.22, 0.58, 0.33], # with     (x^4)
 [0.77, 0.25, 0.10], # one      (x^5)
 [0.05, 0.80, 0.55]] # step     (x^6)
)

The goal: for the second token, $x^{(2)}$ (“journey”), compute a context vector context vector The weighted sum over all tokens' (value) vectors — an enriched representation of one token that incorporates information from the whole sequence. defined in ch. 3 — open in glossary $z^{(2)}$ — an enriched embedding of “journey” that also contains information from $x^{(1)}$ through $x^{(T)}$, weighted by relevance. This is what self-attention is for: building, for every element of a sequence, a representation that incorporates information from all the other elements. Later, trainable weights will let the LLM learn to build context vectors that are useful for generating the next token.

Step 1 — attention scores via dot products

The query is $x^{(2)}$. Its attention score attention score ω — the raw query·key dot product between two tokens, before normalization. defined in ch. 3 — open in glossary toward each token $x^{(i)}$ is simply the dot product dot product Element-wise multiplication of two vectors followed by summing — a measure of how aligned (similar) the vectors are. defined in ch. 3 — open in glossary :

$\omega_{2i} = x^{(2)} \cdot x^{(i)} = \sum_{d=1}^{3} x^{(2)}_d \, x^{(i)}_d$

Worked example for $\omega_{21}$ (query “journey” · token “Your”): $0.55{\cdot}0.43 + 0.87{\cdot}0.15 + 0.66{\cdot}0.89 = 0.2365 + 0.1305 + 0.5874 = 0.9544$.

query = inputs[1]  #A
attn_scores_2 = torch.empty(inputs.shape[0])
for i, x_i in enumerate(inputs):
  attn_scores_2[i] = torch.dot(x_i, query)
print(attn_scores_2)
# tensor([0.9544, 1.4950, 1.4754, 0.8434, 0.7070, 1.0865])

The same six numbers, computed as one matrix product — hover each result cell to watch its row·column dot product happen:

Step 2 — normalize the scores into attention weights

Raw scores have an arbitrary sum. To use them as mixing proportions we normalize them into attention weights attention weight α — the softmax-normalized attention scores; each query's row of weights sums to 1. defined in ch. 3 — open in glossary $\alpha$ (alpha) that sum to 1. The obvious way — divide by the sum — works, but in practice the softmax softmax Exponentiate-and-normalize: turns a vector of scores into positive probabilities that sum to 1. defined in ch. 3 — open in glossary function is used:

$\alpha_{2i} = \mathrm{softmax}(\omega_{2i}) = \frac{e^{\omega_{2i}}}{\sum_{j=1}^{T} e^{\omega_{2j}}}$

Compare the two normalizations on the real numbers — hover a token to track it across the charts:

attn_weights_2_tmp = attn_scores_2 / attn_scores_2.sum()
print("Attention weights:", attn_weights_2_tmp)
# Attention weights: tensor([0.1455, 0.2278, 0.2249, 0.1285, 0.1077, 0.1656])

def softmax_naive(x):
  return torch.exp(x) / torch.exp(x).sum(dim=0)  #A
attn_weights_2_naive = softmax_naive(attn_scores_2)
print("Attention weights:", attn_weights_2_naive)
# Attention weights: tensor([0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581])

attn_weights_2 = torch.softmax(attn_scores_2, dim=0)  #B
print("Attention weights:", attn_weights_2)
# Attention weights: tensor([0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581])

Step 3 — the context vector is a weighted sum

Multiply each input vector by its attention weight, then add everything up:

$z^{(2)} = \sum_{i=1}^{T} \alpha_{2i}\, x^{(i)}$

query = inputs[1]  # 2nd input token is the query
context_vec_2 = torch.zeros(query.shape)
for i, x_i in enumerate(inputs):
  context_vec_2 += attn_weights_2[i] * x_i
print(context_vec_2)
# tensor([0.4419, 0.6515, 0.5683])

3.3.2 All queries at once — three matmuls, no loops

Nothing about “journey” was special: every token wants its own context vector. Instead of looping, compute all pairwise scores in one matrix multiplication — inputs @ inputs.T puts the dot product $x^{(i)} \cdot x^{(j)}$ in cell $[i,j]$ of a 6×6 matrix. Hover the result to see each cell’s row·column origin:

Then softmax each row (dim=1 — normalize across the last dimension, so every query’s weights sum to 1 independently). Toggle between the raw scores and the normalized weights:

attn_scores = inputs @ inputs.T
attn_weights = torch.softmax(attn_scores, dim=1)  #A
all_context_vecs = attn_weights @ inputs  #B
print(all_context_vecs)
# tensor([[0.4421, 0.5931, 0.5790],
#         [0.4419, 0.6515, 0.5683],
#         [0.4431, 0.6496, 0.5671],
#         [0.4304, 0.6298, 0.5510],
#         [0.4671, 0.5910, 0.5266],
#         [0.4177, 0.6503, 0.5645]])

The final matmul is the “weighted sum” of step 3, done for all queries at once — hover row 1 (the “journey” row) and confirm it rebuilds $z^{(2)} = [0.4419, 0.6515, 0.5683]$: