The Math You Actually Need
You will be able to read every formula in this course, check each one with a three-number example on paper, and know why $\sqrt{d}$, logarithms and softmax keep showing up.
Open the transformer paper and you hit this on page 4: $\text{softmax}(QK^\top/\sqrt{d_k})\,V$. Four symbols you may not know, one square root nobody explains, and a superscript that looks like a typo. What do you actually need to know to read it?
Less than you fear. Language models use a small, specific slice of math: vectors, dot products, matrix multiplication, softmax, logs, one kind of derivative, and one fact about variance. This chapter covers exactly that slice, each piece with a picture and a tiny example. Nothing here is decoration; every section is something you will use within the next three chapters.
If you already know linear algebra and calculus, skim the headings and do the interactives. If you do not, go slowly and check every worked example by hand. It matters that you can.
Vectors: arrows and lists of numbers
How does a program do arithmetic on the word "cat"? It cannot. It can do arithmetic on numbers, so the first move in every language model is to represent each token as a list of numbers. That list is a vector.
A vector is two things at once. It is a list, like $x = (2, 1)$, which is how the computer stores it. And it is an arrow from the origin to the point $(2, 1)$, which is how you should picture it. Both views are correct and you will switch between them constantly.
We write the number of entries as the dimension, $d$. Real models use $d$ in the hundreds or thousands. You cannot draw a 768-dimensional arrow, but everything you learn about 2-d arrows carries over unchanged: length, angle, adding two arrows tip to tail, scaling one by a number. Trust the 2-d picture.
Length is the one formula to remember. The length (also called the norm) of $x$ squares each entry, adds them up, and takes the square root:
$$\|x\| = \sqrt{x_1^2 + x_2^2 + \cdots + x_d^2}$$So $\|(2,1)\| = \sqrt{4+1} \approx 2.24$ and $\|(1,2,2)\| = \sqrt{1+4+4} = 3$. That last one is a nice three-dimensional example to keep in your pocket; we will reuse it.
The dot product: how much do these agree?
Here is the central question of attention: given two token vectors, how related are they? We need a single number that is large when two vectors point the same way, zero when they are unrelated, and negative when they oppose. That number is the dot product.
The dot product multiplies matching entries and adds the results. Nothing more:
$$x \cdot y = x_1 y_1 + x_2 y_2 + \cdots + x_d y_d = \sum_{i=1}^{d} x_i y_i$$The $\sum$ is just "add these up for every $i$". So the dot product of two $d$-vectors costs $d$ multiplications and $d-1$ additions, and gives back one number. That collapse from two lists to one number is the whole point.
Take $a = (1, 2, 2)$ and $b = (2, 0, 1)$. Then $a \cdot b = 1{\cdot}2 + 2{\cdot}0 + 2{\cdot}1 = 2 + 0 + 2 = 4$. Lengths: $\|a\| = 3$ (from above) and $\|b\| = \sqrt{4+0+1} = \sqrt 5 \approx 2.24$. Keep these numbers; the next formula uses them.
Why does this measure agreement? Because of a second, geometric formula for the same number. The dot product is also the product of the two lengths times the cosine of the angle between them:
$$x \cdot y = \|x\|\,\|y\|\cos\theta$$So if you divide the dot product by both lengths you get $\cos\theta$, a number between $-1$ and $1$ that depends only on direction. That is cosine similarity: $+1$ means "same direction", $0$ means perpendicular (the vectors share nothing), $-1$ means opposite. In our example, $\cos\theta = 4 / (3 \times 2.24) \approx 0.60$, an angle of about 53°: the vectors lean the same way but are not aligned.
There is a third way to say the same thing, and it is the one to picture. The dot product is the length of $y$'s shadow on $x$ (its projection), times the length of $x$. Drag the vectors below and watch the dashed projection; when the angle passes 90° the shadow flips to the other side and the dot product goes negative.
Watch how the dot product depends on both lengths and the angle, while the cosine depends only on the angle.
A big dot product does not mean "similar". It can mean "long". Two vectors at 60° with lengths 10 and 10 have dot product 50; two vectors at 0° with lengths 1 and 1 have dot product 1. When you want pure similarity, use the cosine. Attention uses the raw dot product on purpose, which is why the $\sqrt{d_k}$ scaling in the attention chapter exists.
Matrices: functions that transform many vectors at once
A model has to turn a token vector into a query vector, a key vector, a bigger hidden vector, and so on. What is the simplest kind of function from vectors to vectors? One where every output entry is a weighted sum of the input entries. The weights form a grid, and that grid is a matrix.
Writing it out for a 2-d input and 2-d output, with weights $W$:
$$W x = \begin{pmatrix} w_{11} & w_{12} \\ w_{21} & w_{22} \end{pmatrix} \begin{pmatrix} x_1 \\ x_2 \end{pmatrix} = \begin{pmatrix} w_{11}x_1 + w_{12}x_2 \\ w_{21}x_1 + w_{22}x_2 \end{pmatrix}$$Look at the two output entries. Each one is a dot product of one row of $W$ with $x$. So a matrix is just a stack of dot products, one per output dimension. The first row asks "how much does $x$ agree with $(w_{11}, w_{12})$?", the second row asks about $(w_{21}, w_{22})$, and the answers become the new vector.
Let $W = \begin{pmatrix} 1 & 0 \\ 1 & 1 \end{pmatrix}$ and $x = (2, 3)$. Row one: $1{\cdot}2 + 0{\cdot}3 = 2$. Row two: $1{\cdot}2 + 1{\cdot}3 = 5$. So $Wx = (2, 5)$. The first coordinate was left alone; the second had the first added to it. That is a "shear", and every matrix is some such geometric move: rotate, stretch, shear, flatten, or a mix.
Why is this the right tool for a model? Because of the "many vectors at once" part. Put five token vectors side by side as the rows of a matrix $X$, and the single product $XW^\top$ transforms all five with one operation. On a GPU that one operation is exactly what runs fast. Almost all the arithmetic in an LLM, well over 90% of it, is multiplying a matrix of token vectors by a matrix of weights.
The shapes rule for matrix multiplication
What is the one thing that goes wrong constantly when you start writing models? Shapes. So learn the rule cold. A matrix with $n$ rows and $k$ columns is "$n \times k$". You may multiply $A$ by $B$ only if the inner numbers match, and the result takes the outer numbers:
$$(n \times k)\cdot(k \times m) = (n \times m)$$What just happened: each entry of the output is a dot product between a row of $A$ (length $k$) and a column of $B$ (also length $k$). Those lengths must agree or the dot product does not exist. There are $n$ rows to choose from and $m$ columns, hence $n \times m$ outputs.
Entry $(i, j)$ of the result, written out, is:
$$C_{ij} = \sum_{t=1}^{k} A_{it}\, B_{tj}$$Try it. Pick shapes below; the checker tells you whether the product exists and computes it live with small integers.
The inner dimensions (columns of A, rows of B) must match. When they do, each output cell is one dot product of length k.
The little $\top$ you saw in the attention formula is transpose: flip a matrix over its diagonal, so an $n \times k$ becomes $k \times n$. In $QK^\top$, both $Q$ and $K$ are $L \times d_k$ (one row per token). You cannot multiply $(L \times d_k)(L \times d_k)$, the inner numbers disagree. Transpose $K$ and you get $(L \times d_k)(d_k \times L) = L \times L$: one score for every pair of tokens. The transpose is there to make the shapes work, and the result is the "who looks at whom" table. That is the whole reason for the funny superscript.
Why (batch, seq, dim) tensors
A matrix is a 2-d grid. What if you have several sequences to process at once, each a grid of token vectors? You stack the grids into a 3-d block. Anything with more than two indices is called a tensor, and the shape you will see in nearly every line of model code is $(B, L, d)$: $B$ sequences in the batch, $L$ tokens in each, $d$ numbers per token.
Why batch at all? Because a GPU is fastest when you give it a big pile of identical work, and processing 32 sequences at once is nearly the same cost as processing one. The batch dimension is bookkeeping: the math within each sequence never looks across it. Whenever a shape confuses you, ask "which of these axes is batch, which is sequence, which is feature?" and the confusion usually dissolves. The PyTorch chapter has a broadcasting visualizer for exactly this.
Softmax: scores into probabilities
The model's last layer produces one score per vocabulary entry: $(3.1, -0.2, 1.7, \ldots)$. Scores can be negative and do not sum to anything in particular. We need probabilities: non-negative, summing to one, with bigger scores getting bigger shares. What is the cleanest way to get there?
Softmax exponentiates each score, which makes everything positive, then divides by the total so the result sums to one:
$$\text{softmax}(z)_i = \frac{e^{z_i}}{\sum_j e^{z_j}}$$What just happened: the exponential turned "score" into "unnormalized weight", and dividing by the sum turned weights into shares. Because $e^{z}$ grows fast, a score that is 1 higher gets $e \approx 2.7$ times the probability; 2 higher gets $7.4$ times. Softmax rewards the leader disproportionately, which is where the "soft max" name comes from: it is a smooth version of "pick the largest".
Scores $z = (2, 1, 0)$. Exponentials: $e^2 \approx 7.39$, $e^1 \approx 2.72$, $e^0 = 1$. Sum $\approx 11.11$. Probabilities: $(0.665, 0.245, 0.090)$. Check: they sum to $1.00$. The top score had a lead of one point and got two-thirds of the mass.
Now the knob. Divide every score by a temperature $T$ before the exponential: $\text{softmax}(z / T)$. With $T < 1$ the gaps between scores grow, so the distribution sharpens toward the leader. With $T > 1$ the gaps shrink and the distribution flattens toward uniform. Same scores, same ranking, very different confidence. In the example, $T = 0.5$ gives $(0.867, 0.117, 0.016)$ and $T = 2$ gives $(0.506, 0.307, 0.186)$. Check one yourself: at $T=2$ the scores become $(1, 0.5, 0)$.
A distribution's spread has a number: its entropy, measured in bits. Uniform over 5 options is $\log_2 5 \approx 2.32$ bits; a sure thing is 0 bits. Watch it fall as you cool the temperature.
Five scores, one temperature. The ranking never changes; only how much the leader dominates.
Softmax is not "normalize by dividing by the sum". Dividing raw scores by their sum fails when scores are negative and does not reward leaders. The exponential is doing the real work. A related trap: softmax on the same scores shifted by a constant gives the same answer ($e^{z+c}$ cancels top and bottom), which is why implementations subtract the max first for numerical safety without changing anything.
Logarithms and log-probabilities
A sentence of 1,000 tokens, each predicted with probability around 0.1, has total probability $0.1^{1000} = 10^{-1000}$. Your computer's floating point numbers bottom out around $10^{-308}$ for 64-bit and $10^{-38}$ for 32-bit. So the probability of any real paragraph is zero as far as the hardware is concerned. How do models compute with probabilities at all?
They never store probabilities. They store their logarithms. The log is the function that turns multiplication into addition:
$$\log(a \cdot b) = \log a + \log b, \qquad \log\frac{a}{b} = \log a - \log b$$What just happened: the probability of a sequence is a product of per-token probabilities, so its log is a sum of per-token log-probabilities. A sum of a thousand numbers around $-2.3$ is about $-2303$, a perfectly ordinary number. The underflow problem is gone, and sums are easier to differentiate than products, which will matter in a moment.
Five tokens with probabilities $(0.5, 0.2, 0.1, 0.4, 0.25)$. Product: $0.5 \times 0.2 \times 0.1 \times 0.4 \times 0.25 = 0.001$. Natural logs: $(-0.69, -1.61, -2.30, -0.92, -1.39)$, sum $= -6.91$. And indeed $e^{-6.91} \approx 0.001$. Same information, no tiny numbers.
Two facts to have at your fingertips. Log of 1 is 0; log of anything smaller than 1 is negative, so log-probabilities are always $\leq 0$ and "higher" means "closer to zero". And the base is a choice: natural log (base $e$) is what PyTorch uses and gives units called nats; base 2 gives bits. They differ by a constant factor ($1 \text{ nat} \approx 1.44 \text{ bits}$), so nothing conceptual depends on which you pick.
Expectation: the average you would get in the long run
When we say a model has "a loss of 3.2", over what? Not one token. It is an average over the data, weighted by how often each thing occurs. That weighted average has a name and a symbol.
The expectation of a quantity $f(x)$ under a distribution $p$ is the sum of each value times its probability:
$$\mathbb{E}_{x \sim p}[f(x)] = \sum_x p(x)\, f(x)$$Read $x \sim p$ as "$x$ drawn from $p$". What just happened: we asked what $f$ is on average if we keep sampling $x$ forever. A fair die has $\mathbb{E}[x] = \tfrac16(1+2+3+4+5+6) = 3.5$. When the distribution is "the training data" we cannot sum over all of it, so we estimate the expectation with a batch: the average over 32 sampled sequences stands in for the true average. Every loss you ever see printed is such an estimate.
Cross-entropy: the loss is surprise
The model says "mat" with probability 0.25 and the true next token was "mat". How wrong was it? Not entirely; it gave the truth a quarter of its belief. We need a single number that is 0 when the model was certain and right, and grows as it gave the truth less probability. What function does that?
Negative log-probability of the true token. Measured in bits, it is called the surprise:
$$\text{surprise} = -\log_2 p(\text{true token})$$What just happened: $p=1$ gives $-\log_2 1 = 0$ bits, no surprise. $p = 0.5$ gives 1 bit. $p = 0.25$ gives 2 bits. $p = 1/1024$ gives 10 bits. Each halving of the probability costs one more bit. This is not an arbitrary choice; Shannon showed it is the only measure of information that adds up properly over independent events, and "bits" here really are bits: a model that averages 2 bits of surprise per token could compress the text to 2 bits per token.
The cross-entropy loss is just the expected surprise over the data, in nats because PyTorch uses natural logs:
$$L = \mathbb{E}\big[-\log p_\theta(\text{true token} \mid \text{context})\big]$$The subscript $\theta$ on $p$ says the probability depends on the model's parameters. So $L$ is a function of the weights, and training means making $L$ small. Note that the loss only looks at the probability of the true token; the model is punished for probability it wasted elsewhere, regardless of where.
Three tokens, model probabilities on the true token $(0.5, 0.125, 0.25)$. Surprises in bits: $(1, 3, 2)$, average 2 bits. In nats: $(0.69, 2.08, 1.39)$, average 1.39. A useful translation: $2^{\text{bits}} = e^{\text{nats}} = 4$ is the perplexity, "the model was as unsure as if choosing uniformly among 4 options". GPT-2 small scores a perplexity of roughly 30 on WikiText-style text; you will measure it in code/lumen/eval.py.
Claude Shannon (1948), "A Mathematical Theory of Communication", defined entropy and showed that $-\log p$ is the right cost of an event. Language models are, quite literally, compressors judged by this measure.
Derivatives: which way and how much
We have a loss $L$ that depends on billions of weights. We want to make it smaller. For each weight, we need to know: if I nudge this up a little, does $L$ go up or down, and by how much per unit of nudge? That two-part question, "which way and how much", is exactly what a derivative answers.
The derivative of $f$ at $x$ is the slope of the graph there, the ratio of a tiny change in output to the tiny change in input that caused it:
$$f'(x) = \frac{df}{dx} \approx \frac{f(x + h) - f(x)}{h}\quad\text{for tiny } h$$What just happened: we replaced the curve near $x$ by its tangent line. The sign of $f'(x)$ says which way $f$ increases; the magnitude says how fast. For $f(x) = x^2$ the derivative is $2x$, so at $x = 3$ the slope is 6: nudge $x$ up by $0.01$ and $f$ rises by about $0.06$. Check: $3.01^2 = 9.0601$, up by $0.0601$.
With many inputs, the gradient $\nabla f$ is the list of all the partial derivatives, one per input, each computed by wiggling that input alone. It is a vector, the same shape as the input, and it points in the direction of steepest increase. For a model, $\nabla_\theta L$ has one entry per weight and says, for each, which way to nudge it to make the loss worse. We will go the other way.
The chain rule: derivatives through a pipeline
A model is a pipeline: the weights affect a hidden vector, which affects the scores, which affect the probability, which affects the loss. How does a change in a weight at the start propagate to the loss at the end? By multiplying the slopes along the way.
If $y$ depends on $u$ and $u$ depends on $x$, then:
$$\frac{dy}{dx} = \frac{dy}{du}\cdot\frac{du}{dx}$$What just happened: a small nudge to $x$ is scaled by $du/dx$ on its way to $u$, then scaled again by $dy/du$ on its way to $y$. Two scalings multiply. With ten stages there are ten factors. That is all backpropagation is: walk backwards through the pipeline multiplying local slopes.
Let $y = (3x + 1)^2$. Split it: $u = 3x + 1$ and $y = u^2$. Local slopes: $du/dx = 3$ and $dy/du = 2u$. Chain: $dy/dx = 2u \cdot 3 = 6(3x+1)$. At $x = 1$: $u = 4$, so $dy/dx = 24$. Numerical check: $x = 1.001$ gives $u = 4.003$ and $y = 16.024$, a rise of $0.024$ for a step of $0.001$. Slope 24. It works.
Gradient descent: walking down a bowl
You are standing on a hillside in fog. You can feel the slope under your feet but see nothing. How do you reach the bottom? Take a step downhill. Feel the slope again. Repeat. That is gradient descent, and it is the only training algorithm in this course; everything else (Adam, learning-rate schedules, momentum) is a refinement of the step.
The update moves each parameter against its gradient, scaled by a step size $\eta$ called the learning rate:
$$\theta \leftarrow \theta - \eta\, \nabla_\theta L$$What just happened: the gradient points uphill, so subtracting it goes downhill; $\eta$ decides how far. Too small and you crawl; too large and you leap over the bottom to the far slope, possibly higher than you started, and the next leap is bigger still. That failure has a name, divergence, and you will see it in the loss curves of month four. The interactive below lets you cause it on purpose.
Symbols
$\theta$ = the parameters (weights)$L(\theta)$ = the loss
$\nabla_\theta L$ = its gradient
$\eta$ = learning rate
Forward
Compute the loss $L$ at the current $\theta$ on a batch of data.Backward
Compute $\nabla_\theta L$ with the chain rule, one entry per weight.Step
$\theta \leftarrow \theta - \eta\,\nabla_\theta L$. Go to step 1.$f(x) = (x-2)^2$, whose minimum is at $x = 2$ and whose derivative is $2(x-2)$. Start at $x = 5$, so the slope is $6$. With $\eta = 0.1$: $x \leftarrow 5 - 0.6 = 4.4$, then $4.4 - 0.48 = 3.92$, then $3.54$, creeping toward 2. With $\eta = 0.5$: $x \leftarrow 5 - 3 = 2$ in one step, perfect. With $\eta = 1.1$: $x \leftarrow 5 - 6.6 = -1.6$, farther from 2 than we started, and the next step lands at $6.32$. Divergence.
The function is $f(x) = (x-2)^2 + 0.5\sin(3x)$: a bowl with ripples, so there are small local dips. Watch where each step lands relative to the slope.
Three things to try. With learning rate 0.1 from $x = 5$, count the steps to settle. Then 0.35: you get there in a couple of steps. Then 1.0: each step lands on the opposite wall, higher. Finally start at $x = -1$ with a small rate and notice the ripple traps you in a local dip: real losses have such dips too, and the fix in practice (bigger early steps, decaying later) is the learning-rate schedule in code/lumen/train.py.
Gradients of matrix operations: just get the shapes right
The chain rule was stated for scalars. A model's stages are matrix products. Do you need matrix calculus? For this course, no; PyTorch computes the gradients. But you need one piece of intuition to debug shapes and to understand LoRA later: the gradient of a loss with respect to a matrix has the same shape as that matrix, and there is only one way to arrange the transposes so the shapes work.
Take a linear layer $Y = XW$ with $X$ of shape $(n \times k)$ and $W$ of shape $(k \times m)$, so $Y$ is $(n \times m)$. Suppose you already know how the loss changes with $Y$, a matrix $G = \partial L / \partial Y$ of shape $(n \times m)$. Then:
$$\frac{\partial L}{\partial W} = X^\top G \quad (k \times n)(n \times m) = (k \times m), \qquad \frac{\partial L}{\partial X} = G\, W^\top \quad (n \times m)(m \times k) = (n \times k)$$What just happened: the weight gradient is the input, transposed, times the output gradient; the input gradient is the output gradient times the weight, transposed. You do not need to derive these. Just notice that each has exactly the shape of the thing it is the gradient of, and that the transposes are forced by the shapes rule. When something in your code has the wrong shape, this is the rule you check against.
Symbols
$X$ = input, $(n \times k)$$W$ = weights, $(k \times m)$
$G = \partial L/\partial Y$, $(n \times m)$
Forward
$Y = XW$. Keep $X$ around; backward will need it.Weight grad
$\partial L/\partial W = X^\top G$. Same shape as $W$. This is what the optimizer updates.Pass it back
$\partial L/\partial X = G W^\top$. Same shape as $X$. Becomes the $G$ of the layer before.Step 1 says backward needs the forward input $X$. So during training, every layer's input is kept in memory until the backward pass reaches it. That is why training uses several times the memory of inference, and it is the problem that gradient checkpointing and mixed precision address in the infrastructure chapter.
Variance, and why $\sqrt{d}$ keeps showing up
Back to the formula from the opening: why divide $QK^\top$ by $\sqrt{d_k}$? The answer is a fact about adding up random numbers, and the same fact explains how weights are initialized and why residual streams are scaled. It is worth twenty minutes.
Variance measures spread: it is the expected squared distance from the mean, $\text{Var}(x) = \mathbb{E}[(x - \mu)^2]$, and its square root, the standard deviation, is the typical size of a deviation. The fact we need: variances of independent quantities add.
$$\text{Var}(x_1 + x_2 + \cdots + x_d) = \text{Var}(x_1) + \cdots + \text{Var}(x_d)$$What just happened: adding $d$ independent things, each with variance 1, gives variance $d$, so a typical size of $\sqrt d$. Not $d$. Random contributions partly cancel, so the sum grows like the square root of the count. Coin flips work the same way: after 100 flips you expect about 50 heads, give or take $\sqrt{100}/2 = 5$.
Now apply it to a dot product. If $q$ and $k$ are $d$-vectors whose entries are independent with mean 0 and variance 1, then $q \cdot k = \sum_i q_i k_i$ is a sum of $d$ terms, each with variance 1 (the product of two unit-variance numbers has variance 1). So the dot product has variance $d$ and typical size $\sqrt d$. With $d = 64$ the scores are typically around $\pm 8$; feed those into softmax and it saturates, one token gets everything and gradients vanish. Dividing by $\sqrt d$ brings the typical score back to $\pm 1$. That is the whole story of the denominator.
Let every entry of $q$ and $k$ be $+1$ or $-1$ with equal chance. Each product $q_i k_i$ is $\pm 1$: mean 0, variance 1. The dot product is a sum of four such terms, so it is one of $-4, -2, 0, 2, 4$, with variance 4 and standard deviation 2 $= \sqrt 4$. Divide by $\sqrt 4 = 2$ and the typical score is 1. Try it: $q = (1,-1,1,1)$, $k = (1,1,-1,1)$ gives $1 - 1 - 1 + 1 = 0$; $k = q$ gives $4$, the maximum.
The $\sqrt{d_k}$ scaling is from Vaswani et al. (2017), "Attention Is All You Need", section 3.2.1, with exactly this variance argument in a footnote. The same reasoning gives the "Xavier" and "Kaiming" weight initializations, which pick weight variances of $1/d$ so that activations neither explode nor vanish through many layers.
Notation used in this course
Every symbol below appears repeatedly. Bookmark this table.
| Symbol | Read as | Meaning |
|---|---|---|
| $x$, $q$, $v$ | lowercase | a vector (one token's numbers). $x_i$ is its $i$-th entry. |
| $W$, $X$, $Q$, $K$ | uppercase | a matrix: a weight, or a stack of vectors one per row. |
| $W^\top$ | "W transpose" | flip rows and columns; $(n \times k)$ becomes $(k \times n)$. |
| $x \cdot y$, $XW$ | dot, matmul | dot product of vectors; matrix product (inner shapes must match). |
| $x \odot y$ | "elementwise" | multiply matching entries, keep the shape (Hadamard product). Used in gating and LayerNorm. |
| $\|x\|$ | norm | length of a vector, $\sqrt{\sum x_i^2}$. |
| $\sum_i$ | sum over i | add up the expression for every value of $i$. |
| $\mathbb{E}_{x \sim p}[f]$ | expectation | average of $f(x)$ when $x$ is drawn from $p$. |
| $x \sim p$ | "drawn from" | $x$ is a sample from distribution $p$. |
| $\nabla_\theta L$ | gradient | vector of partial derivatives of $L$ with respect to every entry of $\theta$; same shape as $\theta$. |
| $\theta$, $\eta$ | theta, eta | all the model's parameters; the learning rate. |
| $\log$ | natural log | base $e$ unless written $\log_2$. |
| $(B, L, d)$ | shape | batch, sequence length, feature width. |
| $p_\theta(y \mid x)$ | "p of y given x" | probability the model with parameters $\theta$ assigns to $y$ after seeing $x$. |
A model is matrix multiplies (many dot products at once) followed by a softmax; its loss is the surprise at the true token in log units; training walks that loss downhill using slopes multiplied through the chain rule; and $\sqrt d$ appears whenever you add $d$ random things.
Practice
Let three tokens have key vectors $k_1 = (1, 0)$, $k_2 = (0, 1)$, $k_3 = (1, 1)$ and a query $q = (2, 1)$. Compute the three dot products, divide by $\sqrt{d_k} = \sqrt 2$, and apply softmax. Which key gets the most weight, and by how much? Then redo it without the $\sqrt 2$ and compare the entropies.
Solution sketch
Dot products: $2, 1, 3$. Scaled: $1.41, 0.71, 2.12$. Exponentials: $4.11, 2.03, 8.35$, sum $14.5$, so weights $(0.28, 0.14, 0.58)$. Unscaled, exponentials $7.39, 2.72, 20.1$ give $(0.24, 0.09, 0.67)$: sharper, lower entropy. The gap grows with $d$; at $d = 64$ the difference is dramatic.
Write a Python function numgrad(f, x, h=1e-5) that returns $(f(x+h) - f(x-h)) / 2h$. Check it against the chain-rule answer for $y = (3x+1)^2$ at $x = 1$ (should be 24) and for the bumpy bowl above at $x = 5$. Then use it to implement gradient descent on the bowl and reproduce the interactive's trajectory for $\eta = 0.1$.
Solution sketch
The central difference is accurate to about $h^2$, so $10^{-5}$ gives roughly ten correct digits. For the bowl, $f'(5) = 6 + 1.5\cos 15 \approx 4.86$. Your loop: x -= 0.1 * numgrad(f, x), printing $x$ each time. Numerical gradients are how you will unit-test hand-written backward passes in code/tests/; they are too slow for training but perfect for checking.
In Python, draw two random vectors of length $d$ with entries from a standard normal (random.gauss(0, 1)), take their dot product, and repeat 1,000 times. Compute the standard deviation of the results for $d = 4, 16, 64, 256$. Plot or print it next to $\sqrt d$.
Solution sketch
You should see roughly $2, 4, 8, 16$. Then divide each dot product by $\sqrt d$ and confirm the standard deviation is about 1 for every $d$. This ten-line script is the entire justification for the scaling in code/lumen/attention.py.
Key takeaways
- A vector is a list and an arrow; its length is $\sqrt{\sum x_i^2}$.
- The dot product $\sum x_i y_i = \|x\|\|y\|\cos\theta$ measures agreement; divide by the lengths for pure direction (cosine).
- A matrix is a stack of dot products, applied to every row of the input at once; $(n \times k)(k \times m) = (n \times m)$, inner dims must match.
- Softmax turns scores into probabilities via $e^{z_i}/\sum e^{z_j}$; temperature scales the gaps.
- Logs turn products into sums; cross-entropy is the expected surprise $-\log p(\text{true})$, and $2^{\text{bits}}$ is the perplexity.
- Derivatives say which way and how much; the chain rule multiplies local slopes; gradient descent steps against the gradient by $\eta$.
- Sums of $d$ independent unit-variance terms have standard deviation $\sqrt d$: that is the attention denominator and the basis of weight initialization.
Further reading
- Parr & Howard (2018). The Matrix Calculus You Need For Deep Learning. Exactly what the title says; read it when the shapes-only intuition stops being enough.
- Goodfellow, Bengio & Courville (2016). Deep Learning, chapters 2 to 4. A careful, free treatment of linear algebra, probability and numerical computation for this field.
- Vaswani et al. (2017). Attention Is All You Need. Section 3.2.1 for the $\sqrt{d_k}$ argument in the authors' own words.
- 3Blue1Brown. Essence of Linear Algebra. Animated intuition for vectors, matrices as transformations, and dot products.