Reading the Equations of a Paper

Discover: a promise from Chamber 1

In the very first chamber you met this line, and a promise that by Chamber 9 you'd read it like an ordinary sentence and turn it into working code. Here it is again, exactly as the paper prints it:

Try reading it aloud now: “Attention of Q, K and V equals softmax of Q K transpose over the square root of d k, all times V.” You can say it, and you know what every piece of notation does. Upright words are function names (Chamber 1), the $T$ is a transpose rather than a power (Chamber 6), $d_k$ is a $d$ with a label, and a fraction divides.

But saying it isn't the same as being able to compute it. To compute it, you need to know what kind of object each letter is, and above all its shape. The paper never writes the shapes down. Before reading on, see if you can work out the first one.

Learn: a protocol for any equation

Experienced researchers don't read a hard equation once, top to bottom, and hope. They interrogate it in a fixed order. Here is that order as six steps. Each one uses skills from an earlier chamber.

1. 1

   1. Find where each symbol is defined

   Look in the notation section if there is one, then in the sentences just before the equation, in algorithm boxes and their captions, and sometimes in an appendix. Never guess a symbol's meaning from habit alone: habits are good hints (Chamber 1), not definitions.

2. 2

   2. Give each symbol a type and a shape

   Is it a number, a vector, a matrix, a set, a function or a distribution? If it's an array, what are its dimensions (Chamber 6)? Then check that every operation in the equation is legal for those shapes.

3. 3

   3. Read it aloud

   Turn the line into a sentence, and then into plain English: what does it do? If you can't say it, go back to step 1.

4. 4

   4. Work a tiny example by hand

   Choose the smallest inputs that aren't trivial: two or three numbers, $2 \times 2$ matrices. Compute the result with pen and paper.

5. 5

   5. Code it

   Translate symbol by symbol, then run it on your tiny example and check that it matches your hand calculation.

6. 6

   6. Sanity-check special cases

   What should happen when inputs are all equal, zero, huge or tiny? Do probabilities sum to 1? Do the shapes come out right? Does the formula reduce to something you already know?

It feels slow the first few times. That's normal, and it gets faster with practice. After a dozen equations the first three steps take seconds, and steps 4 to 6 are where your understanding actually comes from.

Learn: the protocol on attention

Let's run all six steps on the equation above.

Step 1: definitions. The paper has no notation section. The definitions are in the prose of Section 3.2.1: “The input consists of queries and keys of dimension $d_k$, and values of dimension $d_v$”, and in practice the queries are “packed together into a matrix $Q$”, with the keys and values packed into matrices $K$ and $V$.

Step 2: types and shapes. Say there are $n$ queries and $m$ keys, each key with its own value. Packing one vector per row:

Now follow the shapes through the equation:

Notice one thing the paper leaves unsaid: softmax is applied to each row separately, one row per query. The equation's notation doesn't say so, and the prose only implies it (“to obtain the weights on the values”). Silent conventions like this are common, and step 2 is where you catch them.

Step 3: read aloud. “For each query, score every key by its dot product with the query, scale the scores down, turn them into weights with a softmax, and return the weighted average of the values.” Attention is a soft lookup: every value contributes, weighted by how well its key matches the query.

Step 4: a tiny example. One query and two keys, with $d_k = d_v = 2$:

The scores are $QK^T = (1 \cdot 1 + 1 \cdot 1,\ 1 \cdot 1 + 1 \cdot (-1)) = (2, 0)$. Dividing by $\sqrt{2}$ gives $(1.414, 0)$. The softmax gives weights $\big(\frac{e^{1.414}}{e^{1.414} + e^{0}}, \frac{e^{0}}{e^{1.414} + e^{0}}\big) \approx (0.804, 0.196)$. The output is $0.804 \cdot (10, 0) + 0.196 \cdot (0, 10) \approx (8.04, 1.96)$. The query points the same way as the first key, so the output is mostly the first value.

Step 5: code. Symbol by symbol:

import numpy as np

def attention(Q, K, V):
    d_k = K.shape[1]
    S = Q @ K.T / np.sqrt(d_k)                 # (n, m) scores
    S = S - S.max(axis=1, keepdims=True)       # safe: softmax ignores shifts (proof below)
    W = np.exp(S) / np.exp(S).sum(axis=1, keepdims=True)   # each row sums to 1
    return W @ V                               # (n, d_v)

Q = np.array([[1.0, 1.0]])
K = np.array([[1.0, 1.0], [1.0, -1.0]])
V = np.array([[10.0, 0.0], [0.0, 10.0]])
print(attention(Q, K, V))                      # [[8.044 1.956]]

Step 6: sanity checks. If every key is identical, every score in a row is equal, the weights are all $\frac{1}{m}$, and the output is the plain average of the rows of $V$. If you multiply $Q$ by 100, the scores spread out, the softmax puts almost all its weight on the best-matching key, and attention becomes a hard lookup. And since each output row is a weighted average of rows of $V$, its entries can never leave the range of $V$'s columns. All three are quick tests of your code.

Why divide by the square root of d k? A derivation

The paper justifies the $\sqrt{d_k}$ in its footnote 4, with a probability argument you can now follow in full. Suppose the entries of a query $q$ and a key $k$ are independent random variables, each with mean 0 and variance 1. Their dot product is $q \cdot k = \sum_{i=1}^{d_k} q_i k_i$.

Each term has mean zero, because independence lets the expectation of a product split: $\mathbb{E}[q_i k_i] = \mathbb{E}[q_i]\,\mathbb{E}[k_i] = 0$. Its variance is

using the variance shortcut of Chamber 8 and independence again. The $d_k$ terms are independent, and the variances of independent variables add, so

In the paper each head has $d_k = 64$, so raw scores would have a standard deviation of 8. Scores that far apart push softmax to nearly 0 or 1, where, in the paper's words, it has “extremely small gradients”. Dividing by $\sqrt{d_k}$ keeps the scores at a comfortable scale however long the vectors are.

Learn: softmax, temperature and a proof

Chamber 4 previewed softmax. Now it's time to take it apart. Softmax turns any list of real numbers (logits) into probabilities:

Read it “softmax of z, entry i”. Exponentiating makes every entry positive, and dividing by the total makes them sum to 1. Larger logits get larger probabilities. A 2015 paper on distilling large networks into small ones adds a knob:

Here $T$ is a temperature, “normally set to 1”, as the paper says. (Notice that the paper sets $exp$ in italics, which strictly reads as $e \cdot x \cdot p$. By Chamber 1's rules it should be upright, $\exp$. Nobody is confused, because context rescues the reader: that's how small slips survive in published papers.) Here's what the temperature does to the logits $\mathbf{z} = (3, 1, 0)$:

Low temperatures sharpen the distribution towards the largest logit. High temperatures flatten it towards uniform. Distillation trains a small network to match a large network's softened outputs, which carry more information than hard labels: they say which wrong answers were nearly right.

A proof: softmax ignores shifts

Claim. For any logits $\mathbf{z} = (z_1, \ldots, z_K)$ and any real number $c$, adding $c$ to every logit leaves softmax unchanged: $\mathrm{softmax}(\mathbf{z} + c\mathbf{1}) = \mathrm{softmax}(\mathbf{z})$.

Proof. Take any entry $i$. By the definition of softmax, and then the rule $e^{a + b} = e^a e^b$ (Chamber 4),

The third step pulls the constant $e^c$ out of the sum (Chamber 5's linearity), and the last cancels it, which is allowed because $e^c > 0$. Since $i$ was arbitrary, every entry is unchanged. $\blacksquare$

The proof has a very practical use. A computer can't represent $\exp(1000)$: it overflows to infinity, and infinity divided by infinity is “not a number”. Choosing $c = -\max_j z_j$ makes the largest shifted logit exactly 0, so no exponential can overflow, and the proof guarantees the answer hasn't changed. That's the S - S.max(...) line in the code above. Problem 1 below puts it to work.

Learn: Adam's Algorithm 1, line by line

In Chamber 1 you decoded the last line of Adam and were promised the rest. Here is the loop of Algorithm 1 from Kingma and Ba's paper, with a reading and a line of NumPy for each step.

Before the loop, $m_0$, $v_0$ and $t$ start at 0. The paper's defaults are $\alpha = 0.001$, $\beta_1 = 0.9$, $\beta_2 = 0.999$ and $\epsilon = 10^{-8}$.

The caption of Algorithm 1 is a small notation section of its own, and it answers exactly the questions step 2 of the protocol would ask. It says that $g_t^2$ “indicates the elementwise square”, that “all operations on vectors are element-wise”, and that $\beta_1^t$ and $\beta_2^t$ mean “$\beta_1$ and $\beta_2$ to the power $t$”. That last one is Chamber 1's warning about powers and positions, stated by the authors themselves: the $t$ on $m_t$ is an index, but the $t$ on $\beta_1^t$ is a power.

Step 6 of the protocol, sanity checks, reveals something neat. On the very first step, $m_1 = (1 - \beta_1)\,g_1$, so $\hat{m}_1 = g_1$ exactly. Likewise $\hat{v}_1 = g_1^2$. Ignoring the tiny $\epsilon$, the first update is $\alpha \cdot g_1 / |g_1| = \pm\alpha$ for every parameter, whatever the size of its gradient. Without the hats, $m_1 = 0.1\,g_1$ would badly underestimate the gradient. How big is that bias at later steps? If the gradients are stationary (their expected value doesn't change), then

which is exactly the factor the hat divides out. The proof is a geometric series (Chamber 5), and you'll assemble it yourself in the proof puzzle below.

Learn: notation sections, and your own notation table

Step 1 of the protocol is a search, and it helps to know where definitions hide.

• A notation section or table of symbols at the front: textbooks such as Deep Learning and Mathematics for Machine Learning have one, and so do some long papers.
• A “Preliminaries” or “Background” section that fixes notation before the method.
• The prose just before an equation, like attention's Section 3.2.1.
• Algorithm captions, like Adam's.
• An appendix. The VAE paper (Kingma and Welling), which you'll decode below, uses $J$ in its Eq. (10) but only says “Let $J$ be the dimensionality of $\mathbf{z}$” in Appendix B. Its Figure 2 caption calls the same quantity $N_{\mathbf{z}}$.

When a paper doesn't give you a table, build your own as you read. It takes a few minutes and pays for itself on the second page. Here is one for the core of the VAE paper, with one row per symbol:

Section 2.1 of that paper is also a good test of Chamber 8. In a few lines it writes a marginal likelihood as an integral, $p_{\boldsymbol{\theta}}(\mathbf{x}) = \int p_{\boldsymbol{\theta}}(\mathbf{z})\, p_{\boldsymbol{\theta}}(\mathbf{x} \mid \mathbf{z})\, d\mathbf{z}$, and the true posterior with Bayes' rule, $p_{\boldsymbol{\theta}}(\mathbf{z} \mid \mathbf{x}) = p_{\boldsymbol{\theta}}(\mathbf{x} \mid \mathbf{z})\, p_{\boldsymbol{\theta}}(\mathbf{z}) / p_{\boldsymbol{\theta}}(\mathbf{x})$. You can read both now.

Papers also lean on a stock of shorthand that rarely gets defined at all:

Learn: overloaded symbols and abuse of notation

There are more ideas than letters, so symbols get reused. A single paper might use $\sigma$ for the sigmoid in one equation and a standard deviation in the next. $T$ can be a transpose, a temperature or the number of time steps. $p$ is a probability, a distribution, or dropout's keep probability, and PyTorch's nn.Dropout(p=0.5) uses p for the probability of dropping a unit, the opposite of the dropout paper's convention. $|\cdot|$ is an absolute value on a number and a size on a set (Chamber 2), and the double bar is a norm in $\Vert\mathbf{x}\Vert$ but only a separator in $D_{\mathrm{KL}}(p \,\Vert\, q)$. These are overloaded symbols. Context almost always settles them, and your notation table records which meaning this paper uses.

An abuse of notation is a deliberate shortcut: writing something that isn't strictly correct, because the correct version would be cluttered and the meaning is clear anyway. You've met several already:

• Dropping subscripts. “We drop the $\theta$ and write $p(y \mid x)$ for $p_\theta(y \mid x)$.” The model still depends on $\theta$. The notation just stops saying so.
• One letter for many functions. $p(\mathbf{x})$ and $p(\mathbf{z})$ are different distributions that share the letter $p$ (Chamber 8).
• Loose indices. The VAE paper's Eq. (2) has $\mathbf{x}^{(i)}$ on its left-hand side but writes $q_{\boldsymbol{\phi}}(\mathbf{z} \mid \mathbf{x})$ on the right, meaning the same data point. The dropout paper writes $\widetilde{\mathbf{y}}^{l}$ in one line where it means $\widetilde{\mathbf{y}}^{(l)}$.
• Unstated conventions. Attention's row-wise softmax, or $\log$ with no base (in machine learning it's almost always the natural log).

Read beyond the course

These four show the protocol from different distances: a method for whole papers, a paper turned line by line into code, a model notation table, and one more derivation read closely.

Papers and lectures

For the Transformer paper, read Section 3.2 on attention, about a page and a half with Figure 2, and use the shape table above as your notation table. Section 3.2.2 on multi-head attention is just the same function run $h = 8$ times in parallel on smaller projections. Skip positional encodings and the experiments for now.

For the VAE paper, read Section 2.1 (the problem and its notation), Section 2.2 (the variational bound, Equations 1 to 3) and Section 3 (the Gaussian example). Build your notation table as you go, or check it against the one above. The reparameterisation trick in Section 2.4 is the $\mathbf{z} = \boldsymbol{\mu} + \boldsymbol{\sigma} \odot \boldsymbol{\epsilon}$ you met in Chamber 8.

The distillation paper is short and friendly: the first three paragraphs of Section 2 surround the temperature equation above.

Watch

A visual walk through attention. Watch for the grid of dot products between queries and keys, and the moment softmax turns that grid into weights: that grid is $QK^T$. Pause and name the shape of every matrix on screen.

A researcher walking through the paper itself, page by page. Run the protocol alongside him: whenever an equation appears, pause the video and name each symbol's type and shape before he explains it.

Your turn

The last round of practice before the Sphinx. First a marathon over every symbol in the course, then two matching sets, two proofs about Adam, and three problems that implement this chamber's papers.

Proofs

The puzzle proves why Adam's bias correction divides by $1 - \beta_1^t$. It starts from the unrolled form of the moving average, which you then prove yourself by induction.

Code it up