Capstone: read three equations from real papers

Scaled dot-product attention, the Adam update and layer normalization, straight from the papers that introduced them. Label every part, then read them with the seven questions.

11 min Prerequisites: The rest of the track, or as much of it as you like

Symbols in this lessonsoftmax softmax ᵀ transpose √ square root ε epsilon ⊙ Hadamard (element-wise) product η eta (learning rate)

1. The idea

The last capstone is the skill the whole track has been building: reading an equation from a paper cold. Here are three, from widely cited, openly available papers: scaled dot-product attention from the original Transformer paper, the Adam update rule, and layer normalization. The equations are the authors', written in today's usual notation; the explanations are ours.For each one, label every part with its role (output, input, learned, constant or operation) and check your answers. Then go through the seven questions on your own: what comes out, what shape everything is, what each function does, what happens at the extremes. Each uses notation you have met: transposes, softmax, square roots, ε, element-wise products, running averages.If one stops you, paste it into the notation explorer's Decode box or ask the tutor to walk you through it. Papers rarely define their symbols as carefully as these lessons have; reading them is mostly recognising conventions, and you now know the common ones.

2. Try it: label the parts

Pick an equation and choose a role for each part, then check. A part is “input” if it comes in from outside the equation (data, or something computed earlier), “learned” if training sets it, and “constant” if someone picks it by hand.

3. The math

Scaled dot-product attention

\operatorname{Attention}(Q, K, V) = \operatorname{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V

Q, K: tokens × d_k. V: tokens × d_v. QKᵀ: tokens × tokens.

Adam's running averages

m_t = \beta_1 m_{t-1} + (1 - \beta_1)\,g_t, \qquad v_t = \beta_2 v_{t-1} + (1 - \beta_2)\,g_t^2

Momentum, and an average of squared gradients.

Adam's step

\hat{m}_t = \frac{m_t}{1 - \beta_1^{\,t}}, \quad \hat{v}_t = \frac{v_t}{1 - \beta_2^{\,t}}, \qquad \theta_t = \theta_{t-1} - \eta\,\frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}

The hats correct the averages' bias towards 0 in the first steps.

Layer normalization

\operatorname{LN}(\mathbf{x}) = \boldsymbol{\gamma} \odot \frac{\mathbf{x} - \mu}{\sqrt{\sigma^2 + \epsilon}} + \boldsymbol{\beta}, \quad \mu = \frac{1}{d}\sum_{i} x_i, \quad \sigma^2 = \frac{1}{d}\sum_{i}(x_i - \mu)^2

μ and σ² are taken over the features of one token.

4. The code (python)

mm7_paper.py

import numpy as np

def attention(Q, K, V):
    s = Q @ K.T / np.sqrt(K.shape[-1])
    w = np.exp(s - s.max(axis=-1, keepdims=True))
    w = w / w.sum(axis=-1, keepdims=True)
    return w @ V, w

rng = np.random.default_rng(0)
Q, K, V = rng.normal(size=(3, 4, 8))                   # 4 tokens, d_k = d_v = 8
out, w = attention(Q, K, V)
print(out.shape, w.shape, w.sum(axis=-1).round(6).tolist())   # (4, 8) (4, 4) [1.0, 1.0, 1.0, 1.0]

def layer_norm(x, gamma=1.0, beta=0.0, eps=1e-5):
    mu = x.mean()
    var = ((x - mu) ** 2).mean()
    return gamma * (x - mu) / np.sqrt(var + eps) + beta

y = layer_norm(np.array([2.0, 4.0, 6.0, 8.0]))
print(y.round(3).tolist(), round(abs(y.mean()), 6), round(y.std(), 3))   # [-1.342, -0.447, 0.447, 1.342] 0.0 1.0

def adam_step(theta, g, m, v, t, eta=0.1, b1=0.9, b2=0.999, eps=1e-8):
    m = b1 * m + (1 - b1) * g
    v = b2 * v + (1 - b2) * g**2
    m_hat, v_hat = m / (1 - b1**t), v / (1 - b2**t)
    return theta - eta * m_hat / (np.sqrt(v_hat) + eps), m, v

theta, m, v = 3.0, 0.0, 0.0                            # minimize theta^2, whose gradient is 2 theta
for t in range(1, 201):
    theta, m, v = adam_step(theta, 2 * theta, m, v, t)
print(round(theta, 4))                                 # -0.0001

5. You’ll see this in

Modern Architectures: attention
Core Course chapter 10: optimizers Adam, built step by step
Normalization layers in the catalog
Read an equation: the seven questions

6. Check yourself

Answer all 4 questions correctly to complete the chapter · 0 / 4 done

Q1/4 Why does attention divide QKᵀ by √d_k?

Q2/4 In layer normalization, which symbols does training learn?

Q3/4 What is Adam's ε for?

Q4/4 For 4 tokens, what shape is QKᵀ?

Finished this chapter?

Weekly newsletter

What changed in ML architectures this week

New models worth knowing, benchmark moves, and one architecture explained in plain words. Free, once a week, unsubscribe anytime.

Learn machine learning from scratch