I want to keep the arithmetic visible here. One small image, one filter, every output, then a backward pass that changes the filter. The larger architecture grows from those operations.

[!NOTE] The prerequisites are a dot product and the chain rule. The earlier posts on linear algebra and backpropagation cover those. A dot product multiplies matching entries and adds them; the chain rule multiplies derivatives along a path and adds contributions from different paths.

What does a CNN do to an image?

Here is a small classifier I will use for the shape calculations. Its widths are chosen for this example; they are not measurements from a trained model. Both convolutions use stride one, dilation one, and a learned bias per filter.

flowchart TD
  X["RGB image: 3 × 32 × 32"] --> C1["3 × 3 convolution: 8 filters, padding 1"]
  C1 --> A1["ReLU: 8 × 32 × 32"]
  A1 --> P["2 × 2 max pooling, stride 2: 8 × 16 × 16"]
  P --> C2["3 × 3 convolution: 16 filters, padding 1"]
  C2 --> A2["ReLU: 16 × 16 × 16"]
  A2 --> G["Global average pooling: 16 numbers"]
  G --> F["Linear classifier: 10 logits"]
  F --> L["Softmax probabilities and cross-entropy loss"]

An image is a tensor, an array with several axes. I use channel-first notation: channels, height, width. A batch adds an axis in front.

RGB provides three input channels. Eight filters produce eight output channels, also called feature maps. Each filter spans all three RGB channels. The network sums those channel contributions to produce one value at each output position.¹

Pooling reduces spatial size. ReLU changes values without changing shape. The last linear layer produces one raw score, or logit, per class.

How does a convolution work?

Start with one grayscale image, one filter, stride one, and no padding. Omit the batch and channel axes. The output at row $i$, column $j$ is

The indices $u,v$ select positions inside the filter. The indices $i,j$ select its placement on the image. Move the filter, multiply matching entries, sum, and add the same bias.

One complete numerical convolution

Given a $3\times3$ input, a $2\times2$ filter, and scalar bias $b=0$:

Find the $2\times2$ output. For the top-left patch, the four products are

Add them, then add the bias:

The remaining placements use the same four weights:

Check: a two-cell-wide filter fits at two positions along each three-cell axis. The output has four values. This particular filter measures the top-left value minus the bottom-right value of each patch.

Is this convolution or cross-correlation?

The operation above is cross-correlation. Mathematical convolution flips the kernel along both spatial axes. CNN libraries commonly keep the unflipped operation and call the layer convolution.²

For the numerical kernel, a flip gives

On the top-right patch, that produces $-2+3=1$, while the unflipped filter produced $-1$. The convention changes a calculation with fixed weights. With unrestricted learned weights, either orientation can represent the same filters. I use the unflipped convention throughout.

What do stride, padding, and dilation change?

Stride $s$ is the step between filter placements. Padding $p$ adds cells on each side; here those cells are zero. Dilation $d$ is the spacing between sampled filter entries. All three affect the output size.

For one spatial axis, let $n$ be the input length and $k$ the number of filter entries. The effective filter span is

The padded input has length $n+2p$. A filter whose starting coordinate is $t$ fits if

Starts occur at $0,s,2s,\ldots$. Count those positions, including the first:

Apply this separately to height and width. This formula assumes symmetric padding, positive integer stride and dilation, and an effective kernel that fits the padded input.¹

For $n=32$, $k=3$, $p=1$, $d=1$, and $s=2$:

The final unused position is allowed. There is no requirement that the division produce an integer.

The dilated filter still has three entries per axis. It samples offsets $0,2,4$ across a five-cell span. A two-dimensional $3\times3$ dilated filter has nine spatial weights, not twenty-five.

For completeness, the channel-aware forward equation for one image is

Here $o$ selects the output channel, $c$ selects the input channel, and $\widetilde X$ is the already padded image. Subscripts $h,w$ allow separate vertical and horizontal settings. This is ordinary convolution with all input channels connected to every output channel, also called groups=1.

Why does weight sharing save parameters?

One filter has $C_{\mathrm{in}}k_hk_w$ weights and one bias. There are $C_{\mathrm{out}}$ filters, so

The image height and width are absent. A larger image uses the weights more times; it does not add weights to this layer.

For the first layer of our classifier:

It produces $8\cdot32\cdot32=8{,}192$ output values from $3\cdot32\cdot32=3{,}072$ input values. Compare layers producing that same number of outputs:

The local row counts all patch slots, including slots touching padded zeros. The dense layer can represent functions this convolution cannot. The saving comes from restricting the connections and requiring weights at different positions to match.

There is a linear-algebra view of the same restriction. Flatten a one-dimensional input $x=[x_0,x_1,x_2]^\top$ and apply weights $a,b$ with no bias:

The $2\times3$ matrix contains repeated weights and zeros. With $a=1$, $b=-1$, and $x=[1,2,4]^\top$, its rows give $y_0=1-2=-1$ and $y_1=2-4=-2$. A spatial convolution is a structured linear map. Adding a bias makes it affine.

Why add ReLU and pooling?

ReLU, the rectified linear unit, acts on each scalar:

Our numerical convolution becomes

Without nonlinear operations, a stack of affine layers remains one affine map. If $W_1,W_2$ are compatible matrices and $b_1,b_2$ their biases, expand two layers:

ReLU makes which entries pass through depend on the input. The resulting network is piecewise affine rather than one affine function everywhere.

A $2\times2$ max pool with stride two selects the largest value in each non-overlapping window. Average pooling takes the mean. Applied to the activated patch above:

Both turn this $2\times2$ map into $1\times1$ and add no learned parameters. They discard different information. Max pooling keeps the largest response; averaging keeps its contribution relative to the whole window.

Global average pooling takes the mean of an entire channel. For a map of height $h$ and width $w$, its channel-$c$ output is

In the classifier, each of the sixteen $16\times16$ maps becomes one number. The final dense layer therefore needs $10(16+1)=170$ parameters. Flattening those maps first would give $16\cdot16\cdot16=4{,}096$ inputs and $10(4{,}096+1)=40{,}970$ parameters. Averaging saves parameters but removes explicit spatial layout from the classifier input.

How does a small filter see a larger region?

A receptive field is the region of the original input that can affect a particular feature. Later filters read earlier features, so their input coverage grows with depth.³

For a sequential stack, track $r_\ell$, the receptive-field width after layer $\ell$, and $j_\ell$, the distance in original-input pixels between neighboring features. Start with $r_0=j_0=1$. With effective kernel span $k_{\mathrm{eff},\ell}$ and stride $s_\ell$:

A filter adds one previous jump for each extra sampled position when dilation is one. Dilation stretches that span. Apply the equations per axis to the example architecture:

Each feature before global pooling can depend on an $8\times8$ region. At image edges, part of that region may be padding. This is theoretical coverage; learned weights and activation gates determine which inputs influence a particular prediction.

Are CNNs translation invariant?

Translation equivariance means shifting the input shifts the feature map by the corresponding amount. Translation invariance means the output stays the same. These are different properties.

The weight-sharing equation explains equivariance. Use one dimension, stride one, and an infinite grid to avoid boundary issues. Define a shift by integer $a$ as $(T_a x)[i]=x[i-a]$. Let $F$ be cross-correlation with fixed kernel $k[u]$:

For example, with $k=[1,-1]$, a value $5$ at input position 2 and zeros elsewhere produces $-5$ at output position 1 and $5$ at position 2. Move that input value to position 3: the outputs move to positions 2 and 3. Their values do not change.

Finite-image boundaries and subsampling need more care. Here is a counterexample to the claim that max pooling guarantees invariance. Use width-two max pooling with stride two, and circularly shift a four-entry input one place right:

The pooled values changed. The shift changed which values shared a window. Zhang’s work on antialiased CNNs studies this sensitivity to downsampling.⁴

Global averaging is invariant to a permutation of an unchanged set of feature values. An image shift that introduces padding, crops content, or changes downsampled features can change that set. A CNN classifier therefore has no blanket guarantee of translation invariance.

What objective trains a CNN classifier?

For our ten-class model, global pooling produces $g\in\mathbb R^{16}$. The classifier has $V\in\mathbb R^{10\times16}$ and bias $a\in\mathbb R^{10}$:

The logit vector $\ell$ has ten entries. For $M$ mutually exclusive classes, softmax converts logits to probabilities $p_j$. If the correct class is $t$, cross-entropy loss is

Use two logits $[\log2,0]$ and correct class $t=1$ to make one prediction by hand:

The probabilities sum to one. The loss is positive because the correct class has probability below one; the decimal is rounded to six places.

To differentiate with respect to logit $\ell_j$, expand the log probability first:

The indicator $\mathbf1[j=t]$ equals one for the correct class and zero otherwise. The two-logit example gives gradients $[-1/3,1/3]$. Gradient descent raises the correct logit and lowers the other one. The cross-entropy post follows this objective in more detail.

How does backpropagation update a shared filter?

Let $\delta_{o,i,j}=\partial L/\partial Z_{o,i,j}$ be the gradient arriving at a convolution’s output before activation. For valid, stride-one cross-correlation on one image, differentiate with respect to one weight:

The same weight helped produce many output values, so it receives a contribution from each use. Biases obey the same rule:⁵

For a batch, also sum over images. If the objective averages examples, that normalization is already present in $\delta$.

The input gradient adds contributions from every output window that used the input pixel. For input row $r$, column $s$:

where out-of-range $\delta$ entries are zero. The subtraction in the indices comes from solving $r=i+u$ and $s=j+v$ for the output position.

ReLU supplies a gate before these sums. If $G=\partial L/\partial A$ is the arriving gradient, then

At zero, ReLU has no unique derivative; taking the gradient as zero is a common convention. A max-pooling backward pass routes the gradient to the selected maximum. Average pooling spreads it equally across the window. For a unique maximum in a $2\times2$ window and incoming gradient $4$, max pooling sends $4$ to that entry; average pooling sends $1$ to each of the four entries. Tied maxima need a defined selection or subgradient convention.

One full gradient-descent step

I will isolate the convolution and ReLU so every learned value fits on paper. Use the same $X$ and $K$ from the forward example, change the initial bias to $b=2$, and choose a $2\times2$ target map $Y$. This is a small regression exercise with squared error, separate from the classification objective above.

Find the gradients of all four weights and the bias, update them with learning rate $\eta=0.01$, and recompute the loss.

Add two to every entry of the earlier convolution output:

All entries are positive, so all four ReLU derivatives are one. Define a summed loss, with no division by the four outputs:

Subtract the target and square each residual:

The outer derivative of $\frac12 e^2$ with respect to residual $e$ is $e$. The derivative of $e=A-Y$ with respect to $A$ is one. Multiply those derivatives and the ReLU gate:

where $\odot$ means entry-by-entry multiplication. Each shared weight accumulates four contributions, in output order $(0,0),(0,1),(1,0),(1,1)$:

Check the gradient shapes: $\nabla_KL$ is $2\times2$, matching $K$; $\partial L/\partial b$ is a scalar. For the center input pixel, all four output windows contribute:

Some paths cancel. The filter gradients still need all four paths.

Apply $\theta_{\mathrm{new}}=\theta-\eta\nabla_\theta L$ to each learned parameter:

Recompute every output using those new values:

All outputs remain positive, so ReLU leaves them unchanged. The new residuals and loss are

One update reduced this example’s loss from $2$ to $1.06075$. That checks the direction of this step; it does not establish generalization or guarantee that every learning rate will reduce the loss.

For the full ten-class architecture, the trainable parameter count is $224+1{,}168+170=1{,}562$: two convolutions and the final classifier. The second convolution contributes $16(8\cdot3\cdot3+1)=1{,}168$. ReLU and pooling contribute zero. During training, the loss gradient passes through the classifier, the pooling operations, and both convolutions. Every filter update adds the evidence from all locations where that filter was used.

Read next: Numerical gradient checking compares these derivatives against changes in the loss from small parameter perturbations.

The paper is up. DeepSeek-AI, DeepSeek-V4: Towards Highly Efficient Million-Token Context Intelligence, arXiv 2606.19348.¹ The checkpoint I counted is DeepSeek-V4-Flash-0731. I am taking every width below from that config.json. The paper’s rounded Flash size is 284B total / 13B active; the Hub card for this file reports 304.2B. I am not reconciling those two here. Pro is 1.6T total / 49B active, same 1M context.

I cannot picture a 4096-d residual, 256 experts, or a million-token KV cache. So I locked a toy: the string deepseek as 8 raw bytes, and I painted the byte p (position 4, id 112) yellow in every figure. Hidden size 4, two residual streams, three layers so sliding-window attention, CSA, and HCA each show up once. Weights are one decimal. If a digit in a picture fights a table, the table wins. Image models smear numerals.

[!NOTE] Queries, keys, values, and the quadratic cost are in how attention works. The residual-stream arithmetic is in the transformer forward/backward pass. Byte lookup is token embeddings. Rotary positions are RoPE. The gated MLP is the SwiGLU section of the mini LLM walkthrough.

What does one Flash block look like?

Token ids, embedding, mHC expand into hc_mult = 4 residual streams, then 43 layers of attention plus DeepSeekMoE, each wrapped by mHC, then mix, RMSNorm, a main LM head for token $t+1$, and one MTP block for $t+2$.

flowchart TD
  ids["token ids"] --> emb["embedding"]
  emb --> exp["mHC expand, n_hc streams"]
  exp --> blk["block x 43"]
  blk --> mix["mHC mix + RMSNorm"]
  mix --> lm["LM head: token t+1"]
  mix --> mtp["MTP depth 1: token t+2"]
  subgraph one ["one block"]
    attn["SWA or CSA or HCA"]
    moe["DeepSeekMoE"]
    attn --> moe
  end
  blk --- one

compress_ratios in Flash-0731 starts [0, 0, 4, 128, ...]. Layers 0 and 1 are sliding-window only. After that, CSA ($m = 4$) and HCA ($m’ = 128$) alternate. The JSON lists a few trailing zeros after layer 42; num_hidden_layers is 43, so I am only reading the first 43 slots.

{
  "hidden_size": 4096,
  "num_hidden_layers": 43,
  "num_attention_heads": 64,
  "head_dim": 512,
  "num_key_value_heads": 1,
  "q_lora_rank": 1024,
  "qk_rope_head_dim": 64,
  "index_topk": 512,
  "sliding_window": 128,
  "n_routed_experts": 256,
  "num_experts_per_tok": 6,
  "n_shared_experts": 1,
  "moe_intermediate_size": 2048,
  "num_hash_layers": 3,
  "scoring_func": "sqrtsoftplus",
  "topk_method": "noaux_tc",
  "hc_mult": 4,
  "hc_sinkhorn_iters": 20,
  "num_nextn_predict_layers": 1,
  "vocab_size": 129280,
  "max_position_embeddings": 1048576,
  "swiglu_limit": 10.0
}

How does “deepseek” become an $8 \times 4$ matrix?

Raw bytes, not BPE. Ids [100, 101, 101, 112, 115, 101, 101, 107]. That is d e e p s e e k. Yellow p is ASCII 112.

I chose $E$ with one decimal. Yellow row:

Flash vocab is 129280. The toy vocab is 256 so the picture fits. Same gather: id $\to$ row of $E$. See token embeddings for the full-size version of that lookup.

How does mHC turn one stream into two?

Paper §2.2: the residual is $n_{hc}$ streams, not one vector of width $d$. Update

$A$ mixes streams into the block. $C$ writes the block output back. $B$ mixes the residual and is projected onto doubly stochastic matrices by Sinkhorn-Knopp: $M^{(0)}=\exp(\tilde B)$, then alternate row-normalize and column-normalize, Flash hc_sinkhorn_iters = 20. I froze $A,B,C$ after this figure so the later arithmetic stays traceable. Flash would recompute a tiny input-dependent term each time ($\alpha = 0.01$ in the toy).

Expand with $w = [1.0,\ 0.4]$:

Mean-pool the whole $X$ into 8 numbers, RMSNorm, then the dynamic gates. After sigmoid / $2\cdot$sigmoid:

$B$ after three Sinkhorn iterations (row sums and column sums both 1.0):

A doubly stochastic $B$ cannot amplify the residual without bound as depth goes to 43. That is the point of the manifold constraint. The 20 iterations in Flash are cheap relative to attention at 1M tokens; they are not cheap relative to a 2×2 toy.

What does layer 0 sliding-window attention do to p?

Collapse streams: $H = A X$. Yellow:

Head 1, after $W_Q,W_K,W_V$:

RoPE on the last 2 dims, position $i=3$: $\cos 3 = -0.9900$, $\sin 3 = 0.1411$.

Window 2, causal: $p$ sees the previous e and itself. Scale $1/\sqrt{4}=0.5$. Sink logit $s=-1.0$ goes in the denominator only, paper eq. (27):

Denom $2.4048$. Weights $0.4153$ and $0.4317$. Sink mass $0.1530$. The keys do not sum to 1. That leftover is the sink: the head is allowed to attend to nothing.

Head 2 is a summary row in the figure. Concat both heads, grouped $W_O$ (Flash: 64 heads in o_groups = 8):

Flash keeps a 128-token uncompressed window even after compression starts (sliding_window). The first two layers never compress. Local syntax is cheap. The 1M problem is the far tokens.

How does the residual update after attention?

Yellow stream 0:

Stream 1 lands at $[0.4287,\ -0.0915,\ 0.2592,\ 0.1877]$.

Why does layer 0 MoE hash the token id?

Config: num_hash_layers = 3. Paper §2.1: the first MoE layers replace a learned router with hash routing, expert from the token id, no dense FFN at the start the way V3 had.

Toy: expert = token_id mod 4. Second routed expert is (e+1) mod 4 so top-2 still fires. Shared expert always on. Hash has no logits, so I fixed gates $0.7$ and $0.3$.

Occupancy of the hashed primary: $[2,\ 4,\ 0,\ 2]$. Expert 2 is empty. That is why later layers switch to a learned router plus noaux_tc bias.

Input to the expert, after mixing the post-attention streams:

Expert 0 is SwiGLU, intermediate 3 (Flash 2048), output clamped to $[-10,10]$ (swiglu_limit).

$\mathrm{SiLU}(g)=g\cdot\sigma(g)$:

$y = \text{hid}\,W_{\text{down}} = [0.0062,\ 0.0032,\ 0.0013,\ 0.0044]$. Clamp does not bite. On FP8/FP4 training it does; that is why the limit exists (paper §4.2.3).

Combine: shared $+$ gated routed.

Second mHC wrap. $X$ after layer 0, yellow:

Flash activates 6 of 256 routed experts plus 1 shared. That is the 13B-active number sitting inside 284B total.

How does CSA compress 8 tokens into 4 KV entries?

Paper §2.3.1: every $m$ tokens become one KV entry by an overlapped weighted sum of two streams $C^a$ and $C^b$. Adjacent blocks share tokens. Then a lightning indexer keeps top-$k$ (Flash 512). Core attention is shared-KV MQA: the compressed entry is both $K$ and $V$. An uncompressed window of sliding_window tokens is concatenated so a query can see inside its own compression block without breaking causality.

Toy $m=2$, overlap weights $C^a=[0.5,\ 0.3]$, $C^b=[0.2]$ on the next token. Entry 2 uses tokens of entry 1. $H$ after mixing layer-0 $X$, yellow $[0.4930,\ -0.0902,\ 0.2994,\ 0.2226]$.

Entry 1 contains future token s (position 5) relative to query $p$. Causal rule: an entry is visible only if every token in it has position $\le$ the query. Entry 1 is forbidden. The indexer can still score it. Core attention must not.

[!GOTCHA] Flash CSA softmax-normalizes $2m$ overlap weights per block and uses learned $C^a$, $C^b$ plus positional biases. The toy keeps three weights that sum to 1 so the grid stays countable. Same idea, not the same tensor.

RMSNorm on the Q head and on the compressed KV head, then partial RoPE on the last 2 dims (Flash last 64). Indexer scores of $p$ against the four entries, ReLU, rank:

Top-2 by score would be 1 then 2. After the mask, only entry 0 remains. Window $n_{\text{win}}=2$ adds uncompressed tokens 3 and 4 (e and p). Keys for core attention: 1 compressed + 2 window.

Scores for $p$, sink in the denom: $[0.9980,\ 0.1329,\ 1.9828]$, weights $[0.2362,\ 0.0994,\ 0.6323]$, sink mass $0.032$. Context $[1.4511,\ 0.2906,\ -0.6491,\ -0.3722]$.

Then RoPE with position $-i$ on the last 2 dims of that output. Paper: the compressed KV already carries absolute positions, so the weighted sum would leak them unless you rotate the output by $-i$. After that, grouped $W_O$:

Learned router on the post-CSA hidden. Affinity is $\sqrt{\mathrm{softplus}(x)}$, not V3’s sigmoid (scoring_func": "sqrtsoftplus"). Bias $b=[0.0,\ 0.1,\ -0.1,\ 0.0]$ is added only for top-k selection (topk_method": "noaux_tc"). Gates are renormalized from the unbiased scores.

Top-2: experts 3 and 0. Gates $0.5307$ and $0.4693$.

What does HCA change?

Same MQA, window, RoPE, sink. Compression $m’=4$ (Flash 128), no overlap, no indexer. Dense attention over $n/m’$ entries. At toy $T=8$ that is 2 entries: tokens 0-3 and 4-7. Yellow $p$ lives in entry 0. Entry 1 still contains the future, so a causal query at $p$ only sees entry 0 plus the 2-token window.

KV entries per query at Flash context $1{,}048{,}576$:

CSA is the sparse one. HCA is the heavy shrink with no indexer. Interleaving them is the hybrid in compress_ratios. Paper §2.3.4: mixed BF16 on the RoPE dims and FP8 on the rest, FP4 in the indexer, and versus a BF16 GQA-8 head-dim-128 baseline the 1M KV cache is about 2%. Versus DeepSeek-V3.2, Pro at 1M is 27% of the single-token FLOPs and 10% of the KV; Flash is 10% of the FLOPs and 7% of the KV.

How does noaux_tc move expert bias?

No auxiliary load-balancing loss. After a batch, if expert $i$ ran hot, nudge $b_i$ down so the next batch is less likely to pick it. If it ran cold, nudge $b_i$ up. The bias does not enter the gate values, only the top-k pick. A small sequence-wise balance term stops one sequence from dumping everything on one expert (paper §2.1). Hash layers do not learn this; layers 0-2 in Flash skip the router.

On the toy hash layer, occupancy $[2,4,0,2]$ already shows the failure mode a learned bias is meant to fix: one expert empty, one stuffed with every e.

What does the main head predict for t+1?

Final mix of the two streams into $[8\times 4]$, RMSNorm, LM head. Vocab 256 in the toy; I only score four candidate bytes: e (101), k (107), space (32), s (115). True next byte after $p$ is s.

Where does the extra t+2 head get its input?

MTP depth 1, same wiring as V3 (paper §2.1: “without modification”). During training, concatenate the main hidden state of $p$ with the embedding of the true next byte s, RMSNorm, a small projection, one transformer block, a second head. Inference discards this path.

What is the loss, and who updates which weights?

V3 used $L = L_{\text{lm}} + 0.3\, L_{\text{mtp}}$. V4 keeps that MTP setup. $L_{\text{lm}}$ is cross entropy on $t+1$ (s for this $p$). $L_{\text{mtp}}$ is cross entropy on $t+2$ (e). See cross entropy if the softmax-NLL step is rusty.

Muon (paper Algorithm 1, §2.4) updates the 2D weight matrices: Newton-Schulz orthogonalization of the momentum, then a scaled step. Embeddings, RMSNorm gains, and mHC static gates stay on AdamW. Training is FP8/FP4 on clusters, $32$T+ tokens, then YaRN from 64K original_max_position_embeddings to $1{,}048{,}576$ (rope_scaling.type = yarn, factor 16).

What is stored at generate time?

Main head only. The KV cache is compressed entries, not per-token $K,V$ for every layer. After 8 toy tokens: 4 CSA entries, 2 HCA entries. The next byte grows the cache by a fraction of an entry ($1/m$ or $1/m’$), not by a full head-dim vector per layer the way a vanilla transformer does. Flash also keeps the 128-token uncompressed window.

That is the 1M trick in one picture: you are not storing a million keys per head. You are storing $n/4$ or $n/128$ mixed slots, plus 128 locals, plus (for CSA) a top-512 gather.

Can I see p’s numbers on one page?

Five equations I actually used:

Eight Flash numbers worth keeping: 43 layers, $d=4096$, $m=4$, $m’=128$, indexer 512, window 128, 256 routed / top-6 / 1 shared, context $1{,}048{,}576$.

Read next: the June V4 engineering notes for the verification checklist and the $n^2$ cost, then how attention works if the score-row still feels like a vibe. The paper HTML is arXiv 2606.19348. The file I counted is Flash-0731 config.json.

OpenAI’s own Path to Astra post (1 September 2026) confirms Astra as an upcoming model and says it meets the Critical cybersecurity threshold under their Preparedness Framework. Separate OpenAI notes on cyber-critical safeguards describe monitors that read the model’s chain of thought and can interrupt high-risk activity. What none of those posts say is “we loop a transformer stack.” The loop claim is The Information, as Raschka quotes it. I am treating the architecture as a rumor attached to a real model, and the math as the part I can actually work.

[!NOTE] The loop is ordinary transformer blocks run more than once. If queries, keys, values, and the KV cache still feel like three matrices and a vibe, start with how attention works and the full transformer forward/backward pass. I use those mechanics here.

What did the rumor claim?

Raschka reads The Information as saying two things:

1. Astra uses “recurrent depth” or “loop transformers,” and this lets the model improve answers by processing the same text multiple times.
2. Unlike current frontier models that write out how they are thinking, the technique obscures some or all of that reasoning, so humans cannot easily read the steps.

Claim 1 is a description of an architecture that already exists in public models. Claim 2 does not follow from claim 1. I want the mechanism on the page before the safety argument.

One wording fix up front: the loop does not re-tokenize the prompt and “reread the paragraph.” After the embedding lookup, you have hidden states. Those vectors go through the same weight matrices a second time. Same weights, new activations.

What is a looped transformer?

A standard decoder is a stack of distinct blocks. Layer 1 has its own $W_Q, W_K, W_V, W_O$ and MLP. Layer 2 has a different set. You apply each once, then the unembedding.

A looped transformer (also called recurrent depth, or looped depth sharing) keeps fewer unique blocks and applies that stack more than once. If $L$ unique layers run for $R$ passes, you get $LR$ block applications and $L$ sets of weights.

flowchart LR
  X["token embeddings"] --> S["22 unique blocks"]
  S -->|"pass 2, same weights"| S
  S --> Y["RMSNorm + unembedding"]

Nanbeige 4.2-3B is the public example I can measure. The config is not coy:

{
  "hidden_size": 3072,
  "num_hidden_layers": 22,
  "num_loops": 2,
  "num_attention_heads": 48,
  "num_key_value_heads": 8,
  "head_dim": 128,
  "intermediate_size": 10752,
  "vocab_size": 166144,
  "max_position_embeddings": 262144,
  "tie_word_embeddings": false
}

Twenty-two unique decoder layers. num_loops: 2. Hidden states finish the stack, then go through those same 22 weights again. Effective depth 44. Unique transformer weights stay at 22 layers’ worth.

The technical report (Nanbeige LLM Lab, Boss Zhipin) says they pretrained this from scratch on 28T tokens. The card rounds the size as 4B total parameters, 3B non-embedding.¹

[!GOTCHA] The video calls Nanbeige the first frontier-level open-weight model out of India. The PDF author line is Nanbeige LLM Lab, Boss Zhipin. I am ignoring the geography and keeping the layer counts.

How much depth do you buy without duplicating weights?

Ignore embeddings for a second, the way the report does when it talks about “capacity without adding parameters.” Each Nanbeige block is GQA attention plus a SwiGLU MLP.

Query heads are $48 \times 128 = 6144$ wide, while the residual stream is $3072$. KV is $8 \times 128 = 1024$ wide. One layer’s attention matrices:

SwiGLU is three $3072 \times 10752$ maps:

Two RMSNorm scales add $2 \times 3072 = 6144$. Per unique layer, about $143.1$ million parameters. Twenty-two unique layers:

That is the “3B non-embedding” number, rounded. Untied embeddings and lm_head are $166144 \times 3072$ twice, about $1.02$B more, which is how you get a card that says 4B total.

Now duplicate the stack into 44 unique layers instead of looping:

RAM to store the transformer weights stays near the 22-layer model. Training and decode still pay for 44 block applications: the activations still go through 44 attention-plus-MLP steps, and backprop still walks 44 applications. You saved parameters. You did not save FLOPs in the stack.

Embedding and unembedding are the large exception. They are applied once per token either way, so “almost 2× compute” is the right slogan for the blocks and the wrong slogan for the whole model.

Why train the loop from scratch?

Section 2.1 of the Nanbeige report compares two recipes.

Upcycle. Train a normal 22-layer model. After pretraining, start looping the finished stack.

From scratch. Put num_loops: 2 in the architecture on day one so every pretraining token sees two passes.

They write that from-scratch “performs significantly better.” That matches what I would expect from every other tiny architecture change. Swap GELU for SwiGLU on a finished checkpoint and the next-token distribution moves; the model spent 28T tokens learning the old map. A second pass through the same weights is a bigger change than an activation. The report does not publish the full with-or-without-loop ablation (they note the cost), so I am taking their comparison as an internal result, not a public table.

Why two passes, not four?

Same section: two passes were the best trade-off they measured. Relative to a standard transformer, two passes “retain approximately 75% of the token efficiency” while still adding capacity. Extra passes added little, slowed training, and made optimization less stable.

Token efficiency here means how far the loss drops per token seen. You pay almost another full stack of FLOPs per token. You do not get another full stack of unique features. Two was where that curve still looked worth it for them. I have not rerun the sweep.

Why keep a unique KV cache per pass?

This is the part that surprised me the first time, and it shouldn’t have.

Attention at a block is $Q K^\top$ from that block’s input, then a weighted sum of that block’s $V$. Pass 1, layer 1 sees the post-embedding residual. Pass 2, layer 1 sees the residual after 22 blocks. Same $W_K$ and $W_V$. Different $x$. So $K = x W_K$ is a different matrix.

Sharing one KV tensor across both passes is the same kind of cheat as tying two layers that are not actually seeing the same residual. Nanbeige tried it. Cutting the KV cache in half underperformed the full, non-shared cache. They kept the unique cache.

I used the same KV formula as in the DeepSeek long-context notes:

Nanbeige: $h_{kv} = 8$, $d_h = 128$, $b = 2$ (bf16). Per token, per applied layer:

At the config max $n = 262144$:

The 44 GiB figure is $176\,\text{KiB/token} \times 262144$. That is only the KV tensors, not weights or activations. Looping did not give them a free 44-layer model. It gave them 44-layer compute and, if they want the accuracy they measured, 44-layer KV.

[!IMPORTANT] Parameter tying and KV tying are different decisions. Nanbeige ties the weights across passes and refuses to tie the cache.

What does Mixture-of-Recursions add?

Nanbeige is the blunt version: every token takes both passes.

Mixture-of-Recursions (Bae, Kim, Bayat, et al., 2025) puts a router on that loop, in the same spirit as a mixture-of-experts gate, except the “experts” are extra trips through one shared stack. Easy tokens leave after one pass. Harder tokens go around again.

They train this at four sizes, 135M, 360M, 730M, and 1.7B, with three recursion steps, so the recursive variants hold about one-third as many unique layer weights as the vanilla baseline at a given FLOP budget.²

flowchart TB
  A["all tokens, pass 1<br/>shared stack"] --> R1["router"]
  R1 -->|"exit"| OUT["unembedding"]
  R1 -->|"continue"| B["subset, pass 2<br/>same weights"]
  B --> R2["router"]
  R2 -->|"exit"| OUT
  R2 -->|"continue"| C["fewer tokens, pass 3"]
  C --> OUT

Two router designs in the paper: a single router that assigns a depth, or a router at each loop that decides whether this token continues. Same job, different bookkeeping.

Attention gets messy. At pass 3, some earlier tokens already left. You can:

• attend only to tokens that are still alive at that depth (cheaper, smaller cache, less context)
• attend to some stored KV from tokens that exited (more complete, more memory)

That is the same family of trade as sliding-window or sparse attention: full context is a little more accurate and a lot more expensive. The paper also tries sharing KV across recursion depth and reports a quality drop, same direction as Nanbeige.

On their isoFLOP plots, vanilla sometimes beats a fixed recursive transformer, especially at the small end. MoR’s per-token routing is the setting where they report the recursive idea pulling ahead as scale grows. I am not copying axis values off a thumbnail. The paper is the source; the video is a walkthrough of the figures.

How far back does this go?

Raschka dates the lineage to the Universal Transformer. The year in the video comes out as 2008. The paper is Dehghani, Gouws, Vinyals, Uszkoreit, Kaiser, 2018 (ICLR 2019).³

A vanilla transformer is a fixed stack of distinct layers. A Universal Transformer revises every position’s vector with a shared self-attention plus transition, $T$ times. Depth is the number of revisions, not the number of unique blocks. They also add Adaptive Computation Time so each position can halt early. Nanbeige’s “same 22 layers, twice” is a coarse, modern-LLM version of that idea: a chunk of layers as the recurrent unit, $T = 2$, no per-token halt. MoR puts the halt back, as a learned router.

The 2025 latent-reasoning paper in the same family is Geiping, McLeish, Jain, et al., Scaling up Test-Time Compute with Latent Reasoning: A Recurrent Depth Approach. They train a 3.5B model (Huginn-0125) on 800B tokens with a recurrent block that can unroll at test time, up to a compute load they compare to a 50B dense model, without writing a long chain of thought to buy that compute.⁴ The extra trick in their diagram is injecting the original embedding back into the recurrent state (heavy residual reuse of the input), on top of looping the hidden state. Same family as the Universal Transformer: more FLOPs at a fixed parameter count.

flowchart LR
  E["input embedding e"] --> R["recurrent block"]
  H["latent state s"] --> R
  R --> H
  E -.-> R
  H --> U["unembedding"]

Does looping hide chain of thought?

This is the claim I came for, and it is the one I do not buy from the architecture alone.

Reasoning models spend extra inference compute by emitting extra tokens: scratch work in text. Deeper internal computation (more unique layers, or more loops of the same layers) spends extra compute before the next token is written. Those are different ledgers. Mixing them up is how “loop transformer” turns into “you can no longer read the model.”

A useful comparison is already inside one model family. GPT 5.6 Luna and GPT 5.6 Sol are different sizes. Sol can do more work in hidden states per token. That does not mean Sol stopped using a reasoning trace. It might mean Sol needs a shorter notepad for the same exam question. I still watch it write.

Looping is the same kind of extra internal work, purchased with reuse instead of with new matrices. A model that is better at the problem might emit fewer intermediate tokens. Hidden activations already did some work in every transformer layer that ever existed. Path to Astra, if anything, points the other way: they are investing in CoT monitors on this model, which is a strange thing to do if you believed the loop had deleted the transcript.

The math-exam version, because it is the one that stuck: if I am tired and underprepared I fill pages, strike through, restart. If I am sharp I write less, and I still write. A shorter derivation is still a derivation.

Could some other technique (latent-only recurrence pushed hard, no text trace) make monitoring worse? Yes. Geiping et al. are explicit that their point is reasoning in continuous state rather than in tokens. That is a design choice you can make. It is not what num_loops: 2 on a 22-layer Nanbeige stack does by itself. If The Information is describing a different method than “looped transformer,” the architecture name in the article is doing too much work.

Is this recursive self-improvement?

No. Recursive self-improvement is a model editing its own training, tools, or weights so the next version is stronger. A looped transformer is a for-loop around block(x, W) with $W$ frozen during inference. Same parameters, more applications. Scaling trick. Old paper.

I will start taking the Astra architecture claim seriously when it shows up in a system card, a paper, or a config. Until then the public objects are Nanbeige’s 22 × 2, MoR’s router, Huginn’s test-time unroll, and a 2018 Universal Transformer.

Watch: OpenAI Astra and Recurrent Depth / Looped Transformers (Raschka). His write-up and looped-depth gallery page have the diagrams. OpenAI’s Path to Astra is the primary source for what they will currently say about the model.

That leak is the plot. NLI is supposed to test whether a model can decide if a hypothesis follows from a premise. A lot of the famous numbers measure something cheaper: whether the hypothesis looks like the kind of sentence a worker writes for that label.

NLI is still the classification task I reach for when I want to see if a representation can do more than bag-of-words overlap. The job is a pair of sentences in, one of three labels out, and cross-entropy on the three logits. Same softmax as any classifier. The interesting part is the pair, and how easy it is to score the pair without actually comparing it.

flowchart LR
  P["premise p"] --> ENC["encoder"]
  H["hypothesis h"] --> ENC
  ENC --> Z["3 logits"]
  Z --> SM["softmax"]
  SM --> LAB["entailment / contradiction / neutral"]

What is natural language inference?

You get two texts.

• Premise $p$: the context you are allowed to trust.
• Hypothesis $h$: a claim about that context.

The model must pick one label:

These are the three-way labels used in SNLI and MultiNLI.¹ Older Recognizing Textual Entailment (RTE) datasets were often two-way: entails vs does not.² Three-way is stricter. “Not entailed” splits into “false” and “not enough information,” and that split is where models and humans both get sloppy.

[!NOTE] “A woman is playing a guitar on a sidewalk” does not prove “someone is performing music” in first-order logic. Crowd instructions ask whether a typical reader of $p$ would say $h$ is definitely true. Informal entailment, not a proof assistant.

The examples that made the task click for me are the ones on the SNLI project page, each with five Mechanical Turk labels:³

The soccer pair is easy for a reason: males $\to$ men, soccer $\to$ sport, playing stays playing. The sleeping pair shares almost no content with the premise; contradiction here is “this cannot be the same scene.” The umbrella pair is the honest one. Three workers said neutral, one said entailment, one said contradiction. Smiling vs happy is close. Fairy costume is an extra detail the premise never gave you. Neutral is the correct majority, and it is also the class with the worst agreement in SNLI (Fleiss’ $\kappa$ $0.60$ vs $0.77$ for contradiction).¹

If I cannot explain those three rows, I do not understand the task yet.

How was NLI labeled at scale?

RTE-style sets were tiny. The SNLI paper’s comparison point is “fewer than a thousand examples each.”¹ You cannot train a neural encoder on that.

Bowman’s protocol, 2015:

1. Take a Flickr30k image caption as $p$.
2. Ask a crowd worker to write three new sentences: one definitely true given $p$, one definitely false, one neither.
3. On a validation slice, collect four extra labels so each pair has five votes. Gold is a 3-out-of-5 consensus.

That produced the Stanford Natural Language Inference corpus.

Premises average 14.1 tokens; hypotheses 8.3. About 57k pairs were validated: 98% got a three-annotator consensus, 58% were unanimous, overall Fleiss’ $\kappa = 0.70$.¹ A lexicalized classifier hit 78.2% test accuracy. A 100-d LSTM encoder hit 77.6%. Humans writing short captions plus humans writing short hypotheses is a lot of signal, including the wrong kind.

MultiNLI (Williams, Nangia, Bowman, 2018) copies the protocol and changes the text.⁴ 433k pairs, 392,702 in train, ten genres of written and spoken English. Dev and test are 20k each. Five genres appear in training (matched). Five are held out (mismatched). Premise length jumps to 22.3 tokens on average. ESIM, which posted 86.7% on SNLI, drops to 72.3% / 72.1% on MultiNLI matched / mismatched when trained on MultiNLI. Caption English was doing a lot of work.

ANLI (Nie et al., 2020) flips the collection loop.⁵ A human writes a hypothesis meant to fool the current model. If the model is wrong and two verifiers agree with the writer, the pair is kept. Then they train a stronger model and do it again.

BERT trained on SNLI+MNLI scores 0.0 on A1 test. That is not a bug. A1 was collected to be the errors of that model. After throwing SNLI+MNLI+FEVER-NLI+ANLI at BERT-Large, test accuracy is 57.4 / 48.3 / 43.5 on A1 / A2 / A3. RoBERTa-Large gets to 73.8 / 48.9 / 44.4. MultiNLI at 86% did not mean the three-way decision was solved.

flowchart LR
  RTE["RTE 2005, hundreds of pairs"] --> SNLI["SNLI 2015, 570k captions"]
  SNLI --> MNLI["MultiNLI 2018, 10 genres"]
  MNLI --> ANLI["ANLI 2020, humans vs the model"]

How does a model score a pair?

Two families. I keep them separate because they fail in different ways.

Sentence-vector models encode $p$ and $h$ independently, then mix the two vectors. InferSent (Conneau et al., 2017) is the version I still sketch on paper.⁶ A shared BiLSTM-max encoder produces $u$ from the premise and $v$ from the hypothesis. Then three comparisons get concatenated:

$u$ and $v$ are the two sentence vectors. $\lvert u - v \rvert$ is a cheap “where do they disagree.” $u \odot v$ is a cheap “where do they fire together.” Feed $m$ to a small MLP, get three logits, softmax, cross-entropy. InferSent’s BiLSTM-max hits 85.0% SNLI dev and 84.5% test. The encoder then transferred to other sentence tasks, which is why people treated SNLI as ImageNet-for-sentences for a couple of years.

Independent encoding cannot let a word in $h$ attend to a word in $p$. “The man is sleeping” never gets to look at “inspects the uniform” inside the encoder. All the comparison happens after both sentences have already been crushed into one vector each.

Cross-attention / BERT-style models put both sentences in one sequence and let every token see the other side.

BERT’s pair format:⁷

Token-type IDs are 0 on the premise side and 1 on the hypothesis side. After the transformer, you take the $[\mathrm{CLS}]$ hidden state $h_{[\mathrm{CLS}]} \in \mathbb{R}^{H}$ and map it to three logits:

For BERT-base, $H = 768$, so $W$ is a $3 \times 768$ matrix. That is the entire NLI head. Fine-tuning updates $W$, $b$, and the pretrained stack. BERT-base reaches 84.6 / 83.4 on MultiNLI matched / mismatched. BERT-large reaches 86.7 / 85.9.⁷

flowchart LR
  SEQ["CLS + premise + SEP + hypothesis + SEP"] --> BERT["transformer"]
  BERT --> CLS["h_CLS"]
  CLS --> LIN["linear 3 x H"]
  LIN --> Z["z_E, z_C, z_N"]
  Z --> SM["softmax"]

Cross-attention can align soccer with sport token-by-token. It can also align not in the hypothesis with nothing in the premise and take that as a contradiction cue. Extra mixing gives the cheat a better surface unless the data punishes it.

What does a 3-class softmax look like on one example?

Same softmax + cross-entropy as any classifier. Walk one pair.

True label: entailment. Logits I will pretend a head just emitted:

Subtract the max so the exponentials stay small:

Exponentiate, one at a time:

Sum:

Probabilities:

Loss is only the true class:

Gradient of softmax-cross-entropy is $p - t$. One-hot target is $[1, 0, 0]$, so

The entailment logit gets pulled up. The other two get pushed down, in proportion to the probability they currently hold. If the model had put $p_E = 0.05$ instead, that first component would be $0.05 - 1 = -0.95$, four times steeper. Confidence on the wrong classes is the thing being taxed.

import numpy as np

z = np.array([2.0, 0.5, -0.5])
z = z - z.max()
p = np.exp(z) / np.exp(z).sum()
y = 0
loss = -np.log(p[y])
grad = p.copy()
grad[y] -= 1.0

I still run this block on random 3-d logits before I trust a training loop. The numbers have to match F.cross_entropy on the same $z$ and $y$. If they do not, the bug is in my head, not in PyTorch.

Why can a model ignore the premise?

This is the part I wish I had learned before I memorized SNLI accuracies.

Gururangan et al. (NAACL 2018) trained a classifier on the hypothesis alone. No premise.⁸ fastText gets 67.0% on SNLI vs a 34.3% majority baseline, and 53.9% / 52.3% on MultiNLI matched / mismatched vs ~35% majority. A model that cannot see $p$ is not doing inference. It is doing stylometry on $h$.

The giveaways are visible in PMI(word, class):

Workers generating a contradiction insert negation. Workers generating an entailment generalize (dog $\to$ animal, scene $\to$ outdoors) and drop specifics. Workers generating a neutral add a purpose clause or an adjective the caption never licensed. 8.8% of SNLI entailment hypotheses are a bag-of-words subset of their premise. 0.2% of neutrals and contradictions are.

[!GOTCHA] When I wrote my three hypotheses for one caption, I reproduced this table without trying. “A dog is running outdoors” is a legal entailment of a park photo caption, and it is also a sentence that means entailment even if you hide the caption.

They re-scored strong models on the subset where the hypothesis-only classifier failed (the “Hard” split). ESIM goes from 85.8% on full SNLI test to 71.3% on Hard, while Easy sits at 92.6%.⁸ The leaderboard was averaging a style detector with a harder remainder.

LLM-written NLI is not a free escape hatch. Proebsting and Poliak (2024) rebuilt a slice of SNLI with GPT-4 / Llama-2 / Mistral writing the hypotheses from the original worker instructions. BERT hypothesis-only accuracy landed between 86% and 96% on those sets. GPT-4 reused swimming in a pool in more than 10,000 contradictions.⁹ Same protocol, faster annotator, louder artifact.

So the eval I actually trust is:

1. Full pair accuracy.
2. Hypothesis-only accuracy on the same split.
3. The gap. If (1) is 84 and (2) is 67, the extra 17 points is the part that had to look at $p$. That is the number I quote, not 84.

What did later datasets change?

MultiNLI made the sentences longer and the genres messier. The artifact remains, just weaker (53% hypothesis-only, not 67%). ANLI made the label expensive: a pair only enters the set if a strong model already failed it. Hypothesis-only on ANLI-only training still lands in the 40s, so the cheat is reduced, not gone.⁵

Two other descendants I keep on the same shelf:

• XNLI (Conneau et al., 2018): MultiNLI-style pairs translated into 15 languages. Same three labels, now a transfer test. If an English NLI head collapses on Swahili, the “understanding” was lexical English, not the decision procedure.
• e-SNLI (Camburu et al., 2018): SNLI plus a written explanation and highlighted words. Accuracy stays in the same band; you get a second sequence to supervise. I treat the explanation as another string the model can overfit, not as a proof.

NLI also leaked into GLUE as QNLI: SQuAD converted to “does this paragraph sentence contain the answer?” That is binary entailment with a retrieval flavor. Same head shape, different collection story, same need for a no-context baseline.

Where does this sit next to language modeling?

Next-token prediction asks $p(x_t \mid x_{<t})$. NLI asks $p(y \mid p, h)$ with $y \in {E, C, N}$. Different $y$. Same ingredients.

• Token embeddings turn each side into vectors.
• Attention is how BERT lets $h$ tokens read $p$ tokens in one sequence.
• Cross-entropy is the loss on the three logits. There is no special NLI loss.

I got this backwards for a while. I thought NLI was a semantic capability you inspect after pretraining, like a vibe. It is a supervised 3-way classifier. The capability claim lives or dies on whether the classifier used both sentences. Hypothesis-only accuracy is the control experiment. I run it on day one now, before I get attached to a number on SNLI.

So I rebuilt the whole stack in pure NumPy — custom reverse-mode autograd, byte-BPE, training loop, KV-cache, sampler — as a tiny 3.87M-parameter decoder that only does short English small-talk. No code generation. No tools. No retrieval. Deliberately narrow, so every formula had somewhere to land.

This post is the visual architecture walkthrough of that model. If you want the older GPT-2 recipe (absolute position embeddings, MHA, GELU MLP, LayerNorm), start with that walkthrough. If you want the arithmetic of one training step past attention, keep GPT math beyond attention open beside this. Here I’m documenting what I actually shipped: RoPE + RMSNorm (Pre-LN) + GQA + QK-Norm + SwiGLU + weight-tied head.

What did I actually build?

A decoder-only language model that maps token IDs → next-token logits. Same objective as GPT. Different modern-tiny choices.

flowchart TB
    IDS["Token IDs (B, T)"] --> EMB["TiedEmbedding gather\nW ∈ R^{4096×256}"]
    EMB --> B1["Block 1: RMSNorm → Attn → + → RMSNorm → SwiGLU → +"]
    B1 --> B2["Block 2"]
    B2 --> B3["Block 3"]
    B3 --> B4["Block 4"]
    B4 --> FN["Final RMSNorm"]
    FN --> HEAD["logits = H Wᵀ"]
    HEAD --> OUT["Logits (B, T, 4096)"]

What are the exact hyperparameters?

I kept the config small enough to train on CPU and still feel like a real transformer:

Sanity check on head math: $8 \times 32 = 256 = d_{\text{model}}$, and GQA repeats each KV head $8/2 = 4$ times so query heads still match.

How do the 3.87M parameters break down?

Tied embedding (shared with the LM head):

One attention block (GQA + QK-Norm weights):

One SwiGLU:

Two RMSNorms per block ($256 + 256$), times 4 blocks, plus final RMSNorm:

That number is not marketing copy — I print model.n_params() and it returns exactly 3869184.

How does one forward pass look end to end?

Text → byte-BPE → IDs → embed → 4 Pre-LN blocks → final RMSNorm → tied projection → logits → sample.

Shapes, for a batch of size $B$ and sequence length $T$:

Next-token training uses the usual shift: predict token $t+1$ from positions $0..t$, with cross-entropy on the logits. Padding / user-turn tokens in SFT use an ignore index so only assistant tokens pull the gradient.

What happens inside one transformer block?

Pre-LN residual. Normalize before the sublayer, add the residual after.

I used to treat Post-LN vs Pre-LN as trivia. Then I tried training a tiny stack without Pre-LN and watched early layers thrash. Pre-LN is the boring choice that lets a 4-layer toy actually converge.

RMSNorm itself is LayerNorm without the mean subtraction:

No bias. One gain vector $\gamma \in \mathbb{R}^{d}$. That’s it.

How does attention work with GQA, QK-Norm, and RoPE?

This is the part that differs hardest from the GPT-2 post.

What is grouped-query attention?

Full MHA would give every head its own $K$ and $V$. GQA keeps 8 query heads but only 2 key/value heads, then repeats KV:

So each KV head is shared by four query heads. You pay less for the KV cache at inference and fewer parameters in $W_K$, $W_V$. For a chat toy with seq_len=256, that matters less than it does at 32k context — but the code path is the one I wanted to own.

Why QK-Norm before RoPE?

I RMSNorm $Q$ and $K$ per head (dim 32) before rotating. Attention scores then look like:

The scale is $1/\sqrt{32} \approx 0.1768$. Causal mask writes $-10^9$ into future positions so softmax zeros them. Same idea as causal self-attention, just with modern extras bolted on.

How does RoPE encode position?

Instead of adding a learned position vector (GPT-2’s wpe), RoPE rotates pairs of dimensions in $Q$ and $K$ by an angle that depends on token index $m$:

Relative position falls out of the dot product $q_m^\top k_n$ depending on $m-n$. No separate position table. If absolute sine/cosine encodings are what you grew up with, positional encoding explained is the classic version; RoPE is the one this model uses.

Worked fragment for $d_h=4$, position $m=1$, base $10000$:

Cosine/sine tables store $\cos(m\theta)$ and $\sin(m\theta)$ for every position up to seq_len=256. The forward applies the rotate-and-mix form $x\cos + \mathrm{rot}(x)\sin$ on the last dimension.

What is SwiGLU and why 704?

Classic GPT MLP is $d \to 4d \to d$ with GELU. SwiGLU is gated:

SiLU is $x \cdot \sigma(x)$. The $2/3$ rule shrinks the hidden width so the three matrices roughly match the parameter count of a two-matrix $4d$ MLP:

Three matrices of shape $256\times704$ / $704\times256$ dominate the block. On this model, MLP params per layer ($540{,}672$) are about $3.3\times$ the attention params ($163{,}904$). Same story as every modern decoder: the feed-forward is where most of the weight lives.

Why tie the embedding to the LM head?

One matrix $W \in \mathbb{R}^{4096 \times 256}$ does both jobs:

Without tying I’d pay another $1{,}048{,}576$ parameters for a separate lm_head. With a 4k vocab that’s ~27% of the model. GPT-2 does the same trick; see the weight-tying section in Build GPT-2 from scratch. I also wrote a dedicated note on token embeddings if the gather semantics still feel abstract.

What does the full system stack look like?

The model file is only one layer of the cake. Below it: NumPy/CuPy backend, autograd Tensor, functional ops. Above it: loss, AdamW, schedules, train/SFT scripts, KV-cache generate, chat runtime.

If you’ve built autograd from scratch, the Tensor layer will feel familiar — matmul, reshape, gather, broadcast-aware backward. I lean on numerical gradient checking before I trust any new op. Optimizer is AdamW with grad clipping; LR is warmup then cosine decay.

How was it trained?

Pretrain on everyday dialogue (DailyDialog), formatted with chat specials <bos>, <|user|>, <|assistant|>, <eos>:

• ~1.62M tokens · 2000 steps · batch 32 × block 256
• AdamW · cosine LR $3\times10^{-4} \rightarrow 3\times10^{-5}$
• Val perplexity $\approx 25$

SFT on warm supportive dialogues (EmpatheticDialogues) with loss masking — only assistant turns get supervised:

• 19.4k dialogues · 800 steps · fresh AdamW · cosine $1\times10^{-4} \rightarrow 1\times10^{-5}$
• Assistant-masked val perplexity $\approx 33$ · next-token accuracy $0.325$
• Coding-prompt refusal check: $3/3$ answered as chit-chat (no code emitted)

Tokenizer is byte-level BPE, vocab 4096, specials reserved before the byte base — same family of ideas as BPE from scratch, just smaller.

How does inference work?

Sampling supports greedy, temperature, top-$k$, top-$p$. KV-cache stores past $K,V$ so each new token is an incremental forward; I test that cached logits match a full recompute. Chat runtime renders turns, generates until <eos>, optional system persona string.

Hosting (/chat API, Docker) is still on the checklist. The learning artifact is the trainable NumPy model + SFT checkpoint, not a production companion. At 3.87M params it will mismatch emotion, forget turns, and invent nonsense. That’s expected — the point was owning every op, not beating a 7B chat model.

Can I run it?

Conceptually:

from mla.checkpoint import load_checkpoint
from mla.tokenizer import Tokenizer
from mla.chat import ChatSession

tok = Tokenizer.load("data/tokenizer/tokenizer.json")
model, _, _ = load_checkpoint("checkpoints/sft_final.npz")

chat = ChatSession(model, tok, temperature=0.8, top_k=40, top_p=0.9)
print(chat.reply("I had a rough day today."))

Backend switch for GPU is an env flag (ZYN_BACKEND=cuda → CuPy). Gradcheck and overfit runs stay on float64; GPU train/infer use float32.

Where does this sit in the series?

Read order that worked for me:

1. How GPT actually works — big picture
2. Attention: Q, K, V — intuition
3. Token embeddings · positional encoding · BPE
4. Autograd · grad checking
5. Cross-entropy · AdamW
6. Build GPT-2 in PyTorch — classic stack
7. GPT math beyond attention — one full training step
8. This post — modern-tiny NumPy companion (RoPE / GQA / SwiGLU)

I used to think “from scratch” meant reimplementing GPT-2’s 2019 choices in a new file. Rebuilding with RoPE, GQA, and SwiGLU forced me to confront what actually changed in the last few years — and how little of that change shows up if you only ever call nn.TransformerDecoder.

and make it seem like that is the final answer.

It is not.

In reality, $O$ is only one tiny intermediate tensor inside GPT. I spent months treating attention as the whole story — then I tried to trace where $W_Q$, $W_K$, and $W_V$ actually get updated during training and hit a wall. The update doesn’t happen inside attention. It happens at the end of a chain that most diagrams never draw.

Below I follow the actual mathematical flow exactly as GPT does it — one sentence, one token sequence, real shapes. If how attention works gave you the intuition for queries, keys, and values, this is the rest of the pipeline: embeddings through logits, loss, backprop, and the AdamW step that finally moves those weight matrices.

flowchart TB
    TOK["Token IDs"] --> EMB["Embedding lookup"]
    EMB --> PE["+ Positional embedding"]
    PE --> QKV["Q = XW_Q, K = XW_K, V = XW_V"]
    QKV --> MHA["Multi-head attention + mask + softmax"]
    MHA --> PROJ["Output projection W_O"]
    PROJ --> RES1["Residual + LayerNorm"]
    RES1 --> MLP["MLP 768→3072→768"]
    MLP --> RES2["Residual + LayerNorm"]
    RES2 --> BLOCKS["× n_layer blocks"]
    BLOCKS --> LN["Final LayerNorm"]
    LN --> LOG["Logits = H E^T"]
    LOG --> CE["Cross-entropy loss"]
    CE --> BP["Backprop ∂L/∂W_Q, ∂L/∂W_K, ∂L/∂W_V"]
    BP --> ADAM["AdamW update"]

Step 1: Training Data

Suppose the text is:

Tokenizer converts this into token ids:

GPT does not see words.

GPT only sees integers.

Step 2: Input and Target

For next-token prediction:

Input:

Target:

Meaning:

etc.

Every token tries to predict the next token.

Step 3: Embedding Lookup

Suppose

Each token id is converted into a vector.

For example:

becomes

with 768 numbers.

Similarly:

etc.

Stacking them:

Shape:

(See token embeddings for the lookup table mechanics.)

Step 4: Positional Embedding

Attention itself has no idea whether a word came first or last.

So GPT adds position vectors.

Position 0:

Position 1:

Position 2:

etc.

Final input becomes:

where

Now every token knows:

• what word it is
• where it is

(Positional encoding covers why we add rather than concatenate.)

Step 5: Create Q, K, V

This is where attention starts.

GPT owns three learned matrices:

Each has shape

Initially they are random.

Not identity.

Not meaningful.

Just random numbers.

For example:

etc.

Then:

Shapes:

for all three.

What Are Q, K, V Actually?

Imagine token:

Its query vector asks:

Its key vector says:

Its value vector contains:

Step 6: Multi-Head Split

Suppose:

and

Then:

GPT splits:

into 12 pieces.

Each head gets:

dimensions.

Head 1 might learn:

Head 2:

Head 3:

Head 4:

Nobody programs this.

Training discovers it automatically.

Step 7: Compute Scores

Attention score matrix:

Shape:

Suppose the row for

becomes:

This means:

strongly matches

because

is much larger than everything else.

Step 8: Scaling

Without scaling:

can become enormous.

For:

we divide by:

Therefore:

This keeps softmax stable.

Step 9: Causal Mask

GPT cannot see the future.

Suppose sequence:

The token

cannot look at

because Paris has not been generated yet.

Mask:

Then:

Future positions become impossible.

Step 10: Softmax

Suppose a row becomes:

Softmax converts it into probabilities:

Meaning:

attention goes to one token.

Step 11: Attention Output

Attention matrix:

Output:

This is the famous attention output.

For one token:

Weighted average of all value vectors.

This is where people stop.

But GPT is only beginning.

Step 12: Output Projection

All heads are concatenated.

Suppose:

heads each produce

dimensions.

Concatenation gives:

dimensions again.

Then:

mixes information from all heads.

Step 13: Residual Connection

Input before attention:

Attention output:

GPT computes:

The original information survives.

Attention only adds improvements.

Step 14: LayerNorm

For every token vector:

Compute mean:

Compute variance:

Normalize:

Then:

where

are learned.

Step 15: MLP

Now comes the feedforward network.

Suppose:

First:

Shape:

Apply GELU:

Then:

Shape:

Attention answers:

MLP answers:

Step 16: Second Residual

Again:

Now one transformer block is finished.

Step 17: Repeat Many Times

Suppose:

Then:

Each layer gradually builds better representations.

Early layers:

Middle layers:

Later layers:

Step 18: Final LayerNorm

After the last block:

Shape:

Step 19: Vocabulary Projection

Now GPT converts hidden vectors into vocabulary scores.

Using tied embedding weights:

Shape:

where

for example.

For the final position:

These are logits.

Not probabilities yet.

Step 20: Softmax

Maybe:

Step 21: Loss

Correct answer:

Probability assigned:

Cross entropy:

Small loss.

Good prediction.

If probability were:

then:

Large loss.

Bad prediction.

(See cross-entropy loss for why this is the right objective.)

Step 22: The Part Everybody Misses

Now your question:

When do $W_Q$, $W_K$, $W_V$ actually change?

Only now.

Not during attention.

Not during $O = AV$.

Not during softmax.

Only after loss exists.

Loss depends on:

which depends on

which depend on

which depends on

which depends on

which depends on

which depends on

Therefore:

Backpropagation computes:

These gradients answer:

“If I slightly change this weight, how much will the loss change?”

Step 23: AdamW Update

Suppose:

Gradient:

Learning rate:

AdamW updates:

New value:

Tiny change.

Learning happened.

The Complete GPT Learning Equation

One training step is:

The next forward pass uses these new matrices, producing different $Q$, $K$, $V$, different attention scores, different outputs, and hopefully a lower loss. That cycle repeats billions of times during GPT training.

I keep this diagram open when I’m debugging a training run — if loss isn’t moving, the bug is almost never inside the softmax row. It’s somewhere earlier in the chain, or the gradients never made it back to the weight that needed to move.

This post is the optimizer companion to cross-entropy loss. Same deal as autograd: I’ll build the update rule from SGD upward, derive bias correction, show exactly where AdamW diverges from Adam, and run one scalar parameter through three timesteps with real numbers. Defaults match what I use in AL-1: lr=3e-4, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.1.

Why isn’t plain SGD enough?

Stochastic gradient descent:

where $g_t = \nabla_\theta L$ at step $t$ and $\eta$ is learning rate.

Two problems show up immediately in deep networks:

1. Different scales per parameter. Embedding gradients and LayerNorm gradients can differ by orders of magnitude. One global $\eta$ is always wrong for someone.
2. Noisy minibatches. A single batch gradient is a rough estimate. Raw SGD jitters, especially early in training.

Momentum and adaptive methods fix these. Adam combines both.

flowchart TB
  SGD["SGD\nfixed η"] --> MOM["+ Momentum\nsmooth direction"]
  MOM --> RMS["+ RMSprop\nper-param scale"]
  RMS --> ADAM["Adam\nbias correction"]
  ADAM --> ADAMW["AdamW\ndecoupled WD"]

What does momentum add?

Exponential moving average of gradients:

With $\beta_1 = 0.9$, the update is a smoothed blend of recent gradients — dampens oscillation in ravines, builds speed along consistent directions.

Initialize $m_0 = 0$. At $t = 1$:

The first step only uses 10% of the actual gradient. That’s the cold start problem Adam fixes with bias correction.

What does RMSprop add?

Track a moving average of squared gradients:

(elementwise square)

Scale each parameter’s update inversely:

$\epsilon = 10^{-8}$ prevents division by zero. In float32 training I sometimes use $10^{-6}$; for my NumPy AL-1 build $10^{-8}$ is fine.

How does Adam combine them?

Adam (Kingma & Ba): momentum on $g$, RMSprop on $g^2$, plus bias correction.

At step $t$ (1-indexed):

Bias correction:

Update:

Why bias correction? Both $m_t$ and $v_t$ start at zero, so early $m_t$ and $v_t$ are systematically underestimated. Dividing by $1 - \beta^t$ ramps the effective average up toward the true scale. At $t = 1$, $\beta_1 = 0.9$:

First-step gradient gets full weight instead of 10%.

At $t = 10$: $1 - 0.9^{10} \approx 0.651$. Still a noticeable boost. By $t = 100$ correction is negligible for $\beta_1$, but $\beta_2 = 0.999$ needs hundreds of steps to settle — $v_t$ warms up slowly because squared gradients are tiny early on.

A numeric walk-through: one weight, three steps

Single parameter $\theta$, learning rate $\eta = 0.1$, $\beta_1 = 0.9$, $\beta_2 = 0.999$, $\epsilon = 10^{-8}$, no weight decay.

Start: $\theta_0 = 1.0$, $m_0 = 0$, $v_0 = 0$.

Step 1: $g_1 = 0.5$

Without bias correction I’d have used $m_1 = 0.05$, step size $0.01$ — ten times too small on step one.

Step 2: $g_2 = -0.3$

Gradient flipped sign; momentum ($m$ still positive) fights the flip briefly — that’s the smoothing.

Step 3: $g_3 = 0.2$ — I’ll stop writing every digit. Point is: $v_t$ grows when gradients are consistently large, shrinking effective steps on noisy axes.

Adam vs AdamW: what actually changes?

L2 regularization in Adam (the old way): add $\lambda \theta$ to the gradient before the Adam update:

Weight decay gets entangled with the adaptive scaling. Large $v_t$ shrinks not just the data gradient but also the decay term — regularization effect becomes inconsistent across parameters.

AdamW (Loshchilov & Hutter): decouple weight decay from the gradient-based update.

In code, two lines in step():

if weight_decay != 0.0:
    p.data -= lr * weight_decay * p.data
p.data -= lr * m_hat / (np.sqrt(v_hat) + eps)

Transformers almost universally use AdamW. GPT-2, LLaMA, my AL-1 pretrain — same pattern. Typical weight_decay=0.1 on weights; bias vectors and LayerNorm scales usually get weight_decay=0.

Default hyperparameters and what they mean

Learning rate is the knob I still tune. Warmup + cosine decay come later in the AL-1 checklist; the optimizer itself doesn’t care — it just consumes whatever $\eta$ I pass each step.

$\beta_2 = 0.999$ means $v_t$ is a very slow average. Early training steps have tiny $\hat{v}_t$ after correction, so effective learning rates can be large until $v$ accumulates. That’s one reason warmup helps: small $\eta$ early while $v$ calibrates.

From-scratch AdamW (minimal)

Matches my AL-1 zyn/optim.py:

class AdamW:
    def __init__(self, params, lr=3e-4, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.0):
        self.params = list(params)
        self.lr = lr
        self.beta1, self.beta2 = betas
        self.eps = eps
        self.weight_decay = weight_decay
        self.t = 0
        self.m = [np.zeros_like(p.data) for p in self.params]
        self.v = [np.zeros_like(p.data) for p in self.params]

    def zero_grad(self):
        for p in self.params:
            p.grad = np.zeros_like(p.data)

    def step(self):
        self.t += 1
        bc1 = 1.0 - self.beta1 ** self.t
        bc2 = 1.0 - self.beta2 ** self.t
        for i, p in enumerate(self.params):
            g = p.grad
            self.m[i] = self.beta1 * self.m[i] + (1.0 - self.beta1) * g
            self.v[i] = self.beta2 * self.v[i] + (1.0 - self.beta2) * (g * g)
            m_hat = self.m[i] / bc1
            v_hat = self.v[i] / bc2
            if self.weight_decay != 0.0:
                p.data -= self.lr * self.weight_decay * p.data
            p.data -= self.lr * m_hat / (np.sqrt(v_hat) + self.eps)

No gradcheck on the optimizer itself — I verify by overfitting a tiny model: one batch, loss should drop monotonically for 50 steps if CE + AdamW are wired correctly.

How does the full training step look?

sequenceDiagram
  participant Data
  participant Model
  participant Loss
  participant Autograd
  participant AdamW

  Data->>Model: batch (x, y)
  Model->>Loss: logits
  Loss->>Autograd: scalar L
  Autograd->>Model: ∂L/∂θ
  AdamW->>Model: θ ← θ − update

One step:

1. optimizer.zero_grad()
2. Forward → cross-entropy loss
3. loss.backward()
4. (Optional) clip global grad norm
5. optimizer.step()

I run gradcheck on the loss once. I sanity-check the optimizer by overfitting. Only then do I point it at a real corpus.

When AdamW is the wrong tool

Adam family excels at noisy, high-dimensional, non-stationary objectives — exactly language modeling. It can generalize slightly worse than tuned SGD on some vision tasks with careful schedules. For small convex problems SGD+momentum is simpler and sufficient.

Memory cost: two extra full-sized buffers per parameter ($m$ and $v$). For a 124M-param GPT that’s hundreds of MB of optimizer state on top of weights and activations. Tradeoff I accept.

Mistakes I watch for

Adam looked like a bag of heuristics the first time I read the paper. Momentum for direction, RMSprop for scale, bias correction for cold start, decay on the side — each piece solves a specific failure mode of SGD. AdamW is the version that plays nice with weight decay when training billion-parameter models.

I’ve got autograd, gradcheck, cross-entropy, and AdamW. The next items on my AL-1 checklist are gradient clipping, LR warmup, and a training loop that actually runs overnight without me babysitting the loss curve. The optimizer won’t save a wrong loss or a broken backward pass — but once those are solid, AdamW is what makes the update step boring in the best way.

If you’ve got working autograd and you’ve run numerical gradient checking, this is the next piece: turn model outputs into a scalar you can minimize, derive the gradient, and see the famous cancellation that makes softmax + cross-entropy so clean to implement.

Why doesn’t MSE work for classification?

Mean squared error on one-hot targets sounds reasonable until you try it.

Say three classes with logits $z = [2.0,\, 1.0,\, 0.1]$ and true label $y = 0$. One-hot target is $t = [1, 0, 0]$. Softmax probabilities:

MSE loss:

The model already assigns 66% to the correct class and MSE is tiny. MSE treats probability $0.66$ and $0.99$ as nearly equivalent errors from $1.0$ — both are “close” in L2 sense. For classification I care about confidence on the right class and suppressing wrong classes. Cross-entropy punishes low probability on the true label much harder:

Still not catastrophic, but the gradient structure is what matters. MSE + softmax gives messy, vanishing gradients when the model is confidently wrong. Cross-entropy + softmax gives a clean $p - t$ residual. That’s why every language model uses it.

flowchart LR
  LOG["Logits z"] --> SM["softmax"]
  SM --> P["probs p"]
  P --> CE["−log p_y"]
  TGT["target y"] --> CE
  CE --> L["scalar loss L"]

What is cross-entropy, intuitively?

Cross-entropy measures how surprised I am by reality, given my model’s distribution.

If my model says $p(\text{cat}) = 0.9$ and the true label is cat, surprise is low: $-\log 0.9 \approx 0.105$. If my model says $p(\text{cat}) = 0.01$ but the label is cat, surprise is huge: $-\log 0.01 \approx 4.605$.

For a single example with true class $y$ and predicted probabilities $p$:

Log does two jobs: turns multiplication of probabilities into addition (handy for long sequences), and blows up as $p_y \to 0$, which is exactly the penalty I want.

How do you derive softmax?

The model outputs raw scores $z \in \mathbb{R}^C$ — one logit per class. I need a probability vector $p$ that sums to 1.

Exponentiate (makes everything positive):

Normalize:

Numerical trick I always use: subtract the max before exponentiating so nothing overflows.

For $C = 3$ logits $z = [2.0,\, 1.0,\, 0.1]$:

How do you combine softmax and cross-entropy?

For one example with true index $y$:

Expand:

The second term is the log-sum-exp (LSE). This form is what I implement in code — no need to materialize $p$ in the forward pass unless I want probabilities for inspection.

For a batch of $N$ examples, mean loss:

Language models are the same thing with more dimensions. Logits shape $(B, T, V)$, targets shape $(B, T)$ — each position is an independent classification over vocabulary size $V$. GPT-2 small: $V = 50257$. Average the loss over all non-masked token positions.

What is the gradient? (The cancellation)

This is the part that made implementation click. Define $p = \mathrm{softmax}(z)$. For one example:

Read that: gradient is predicted probability minus one-hot target. Correct class gets $p_y - 1$ (negative when $p_y < 1$, pushing logits up). Wrong classes get $p_i$ (positive, pushing logits down).

Quick sanity check with $p = [0.659, 0.242, 0.099]$ and $y = 0$:

The true class logit should increase; the others should decrease. The gradient says exactly that.

Why does this fall out? Chain rule through $L = -z_y + \mathrm{LSE}(z)$:

No separate softmax Jacobian multiplied by cross-entropy gradient — they collapse. In autograd I can fuse the ops or implement one backward kernel. I fused them in AL-1’s cross_entropy and gradcheck passed first try.

flowchart TB
  subgraph "Backward at one logit z_i"
    PY["p_y"] --> GY["∂L/∂z_y = p_y − 1"]
    PI["p_i (i≠y)"] --> GI["∂L/∂z_i = p_i"]
  end

A full numeric walk-through (3 classes)

Logits $z = [1.0,\, 3.0,\, 0.5]$, true label $y = 1$ (middle class).

Step 1 — stable softmax

Step 2 — loss

Step 3 — gradient

Model was already pretty confident on class 1 (82%). Gradient magnitude on the true class is modest. If instead $p_1 = 0.05$:

Much steeper pull. Cross-entropy automatically scales the correction by how wrong I am.

What about perplexity?

Perplexity is just cross-entropy dressed up for language modeling:

If average CE is $2.3$ nats per token, perplexity is $e^{2.3} \approx 9.97$ — “as confused as if uniformly guessing among ~10 tokens.” GPT-2 at convergence on web text lands around 20–30 depending on dataset; random over 50k vocab would be ~50,000. Perplexity is monotonic with CE, so minimizing one minimizes the other.

From-scratch implementation sketch

Forward (mean over $N$ tokens):

import numpy as np

def cross_entropy_forward(logits, targets):
    N, V = logits.shape
    m = logits.max(axis=1, keepdims=True)
    shifted = logits - m
    log_sum_exp = m[:, 0] + np.log(np.exp(shifted).sum(axis=1))
    rows = np.arange(N)
    nll = log_sum_exp - logits[rows, targets]
    loss = nll.mean()
    return loss, shifted

Backward (given stored shifted logits from forward):

def cross_entropy_backward(shifted, targets, grad_out=1.0):
    N, V = shifted.shape
    exp = np.exp(shifted)
    probs = exp / exp.sum(axis=1, keepdims=True)
    grad = probs.copy()
    grad[np.arange(N), targets] -= 1.0
    grad *= grad_out / N
    return grad

PyTorch’s F.cross_entropy expects raw logits and integer targets — it applies log-softmax + NLL internally. Don’t softmax first; you’ll double-softmax and get garbage.

ignore_index and masked positions

Instruction fine-tuning masks loss on padding or prompt tokens. Same formula, but zero out positions where target == ignore_index and divide by the count of kept positions, not total length.

For AL-1 pretraining I use every token in the packed window. For instruction tuning I’ll mask user prompt positions and only backprop on assistant completions — same cross_entropy, different mask.

Where does this sit in the training loop?

flowchart LR
  BATCH["batch x, y"] --> GPT["GPT forward"]
  GPT --> LOG["logits (B,T,V)"]
  LOG --> CE["cross_entropy"]
  CE --> L["loss"]
  L --> BW["backward"]
  BW --> OPT["AdamW step"]

Cross-entropy is the objective. Autograd computes $\partial L / \partial \theta$. AdamW uses those gradients to update weights. Embeddings, attention, LayerNorm — none of it learns without a scalar loss telling the network what “better” means.

For language models, “better” means assign high probability to the true next token at every position. Cross-entropy is that requirement written as calculus.

Common mistakes I hit

Run gradcheck on the fused op with random logits shape $(2, 3, 7)$ before trusting the training loop. Takes five seconds, saves five hours.

I used to treat cross-entropy as a black-box API call. Implementing it once — forward, backward, numeric check — fixed that. The loss isn’t magic; it’s “make $p_y$ big” expressed as $-\log p_y$, with a gradient so clean ($p - t$) that the whole training stack downstream gets simpler. Next up in the series: actually moving the weights with AdamW.

Byte Pair Encoding is the algorithm behind GPT-2, GPT-3, LLaMA, and most modern LLM vocabularies. This post is one continuous walkthrough: why character-level and word-level tokenization fail, how byte-level BPE fixes both, the math of merge selection with every count written out, a complete training example on the tiny corpus "low low lower lowest" from pretokenization through four merges to final encodings, inference-time encoding step by step, decoding, and the Python I’m using in zyn/bpe.py. Same pipeline as numerical gradient checking for autograd — own the data path before you trust the model.

flowchart LR
  TXT["Raw text\n'def hello():'"] --> PRE["Pretokenize\nregex splits"]
  PRE --> BYTES["Bytes per pretoken\n[100, 101, 102, ...]"]
  BYTES --> MERGE["Apply learned merges\nrank-ordered"]
  MERGE --> IDS["Token IDs\n[412, 891, 37]"]
  IDS --> EMB["Embedding lookup\nin the transformer"]

Why can’t the model just read characters?

You could map each Unicode character to an integer. Shakespeare at ~1 MB becomes ~1M tokens. A 1024 context window holds roughly one screen of text. Training becomes painfully slow because every forward pass processes far more positions than necessary.

Word-level tokenization has the opposite problem. The vocabulary explodes — every typo, hashtag, and async def variant is a new symbol you’ve never seen. Out-of-vocabulary (OOV) words become <UNK> and the model loses information.

BPE sits in the middle. Start with a small base vocabulary (256 raw bytes), then learn the most frequent pairs and merge them into new symbols. Common fragments like ` the, ing, def, ()` get their own IDs. Rare strings still tokenize — they just fall back to bytes. No OOV dead ends.

What is byte pair encoding?

BPE is a compression algorithm adapted for NLP. The training loop is brutally simple:

1. Split corpus into pretokens (word-like chunks).
2. Represent each pretoken as a sequence of byte IDs.
3. Count every adjacent pair $(a, b)$ across all pretokens, weighted by word frequency.
4. Merge the most frequent pair into a new token ID.
5. Repeat until you hit your target vocabulary size.

At inference time you don’t recount anything. You greedily apply the merges you learned, in the order they were created (earliest merge = highest priority).

The original paper (Sennrich et al., 2016) operated on characters for neural machine translation. GPT-2 (Radford et al., 2019) went byte-level: each of the 256 UTF-8 bytes is a starting symbol. That single decision eliminates unknown Unicode — any string encodes to bytes, bytes decode back to UTF-8.

Why bytes, not characters?

UTF-8 is variable-width. The character é might be one code point but two bytes. Byte-level BPE sidesteps the entire Unicode normalization rabbit hole:

• Vocabulary floor is fixed: 256 bytes + a few special tokens. You always know the minimum size.
• No OOV: 🤖, Cyrillic, Zalgo text — all bytes, all representable.
• Compression where it matters: Frequent English/code fragments merge into multi-byte tokens; rare emoji stay as 4 byte tokens.

Tradeoff: byte sequences are longer than character sequences for non-Latin scripts. For English + Python code (what I’m training on), it’s a good deal.

My implementation reserves 6 special tokens before the 256 bytes:

So byte_base = 6, and raw byte $b$ maps to token ID $6 + b$.

For the letter l:

For ASCII space:

Minimum vocabulary size:

You cannot train with vocab_size < 262 — the code raises immediately.

What is pretokenization?

BPE doesn’t run on the raw string as one blob. First, a regex splits text into pretokens — chunks that look like words, numbers, or punctuation. GPT-2 uses this pattern (same in my code):

PAT = re.compile(
    r"""'s|'t|'re|'ve|'m|'ll|'d| ?[^\W\d]+| ?\d+| ?[^\s\w]+|\s+(?!\S)|\s+"""
)

What it captures:

• Contractions: 're, 've, 'll
• Words with optional leading space: ` hello, def`
• Numbers: ` 42`
• Punctuation clusters: ():
• Whitespace runs

The leading-space convention matters. In English, a space before a word is semantically attached — GPT learns ` the as one token, not ` + the. That’s why "low" and " low" are different pretokens with different merge histories.

How does training pick the next merge?

For each pretoken $w$ with frequency $\text{freq}(w)$, maintain a sequence of byte IDs. Count adjacent pairs:

Pick the pair with maximum count:

If the count is below min_freq (default 2), stop early.

Assign the merged pair a new ID next_id, store merges[(a,b)] = next_id, and set:

where $\Vert$ is byte concatenation. Then replace every occurrence of the pair $(a,b)$ with next_id in all pretoken sequences. Increment next_id. Loop until next_id == vocab_size.

flowchart TD
  A["Count all adjacent pairs\nweighted by pretoken freq"] --> B{"Best pair count\n≥ min_freq?"}
  B -->|no| STOP["Stop training"]
  B -->|yes| C["Assign new ID,\nstore merge rule"]
  C --> D["Concat vocab bytes\nvocab[new] = vocab[a]+vocab[b]"]
  D --> E["Replace pair everywhere\nin pretoken ID lists"]
  E --> F{"next_id < vocab_size?"}
  F -->|yes| A
  F -->|no| STOP

Full worked example: "low low lower lowest"

I will train BPE on the corpus "low low lower lowest" with byte_base = 6, vocab_size = 280, and min_freq = 2. Every pretokenization step, byte conversion, pair count, merge, and final encoding is written on its own line. This is the same worked example GPT tutorials use — I’m using real token IDs from my implementation, not abstract letters.

Step 1 — Pretokenize the corpus

Input string:

The regex PAT.findall splits on word boundaries and keeps leading spaces attached:

Each distinct pretoken appears once in this tiny corpus:

Note: "low" and " low" are different strings because of the leading space on the second occurrence.

Step 2 — Convert each pretoken to byte token IDs

Token ID formula: $\text{id} = 6 + \text{byte}$.

Pretoken 'low' — bytes l, o, w:

Pretoken ' low' — bytes space, l, o, w:

Pretoken ' lower' — bytes space, l, o, w, e, r:

Pretoken ' lowest' — bytes space, l, o, w, e, s, t:

Step 3 — Count all adjacent pairs (before any merges)

Pair $(114, 117)$ — the bytes lo — appears in all four pretokens:

Pair $(117, 125)$ — the bytes ow — same four pretokens:

Pair $(38, 114)$ — space followed by l:

Pair $(125, 107)$ — we — only in ` lower and lowest`:

All other pairs have count $1$. The maximum is $4$. Ties go to whichever pair the Counter returns first — here $(114, 117)$ wins.

Step 4 — Merge 1: $(114, 117) \rightarrow 262$

New token ID:

Vocabulary entry:

Replace $(114, 117)$ everywhere:

Step 5 — Count pairs after merge 1

Pair $(262, 125)$ — merged lo followed by w — in all four pretokens:

Pair $(38, 262)$ — space followed by merged lo:

Pair $(125, 107)$ — we:

Maximum count is $4$. Merge $(262, 125)$.

Step 6 — Merge 2: $(262, 125) \rightarrow 263$

Replace $(262, 125)$ everywhere:

The pretoken 'low' is now a single token.

Step 7 — Count pairs after merge 2

Pair $(38, 263)$ — space + low — in three pretokens ( low, ` lower, lowest`):

Pair $(263, 107)$ — low + e:

Maximum is $3$. Merge $(38, 263)$.

Step 8 — Merge 3: $(38, 263) \rightarrow 264$

Replace $(38, 263)$ everywhere:

Step 9 — Count pairs after merge 3

Pair $(264, 107)$ — ` low + e — in lower and lowest`:

This meets min_freq = 2. Merge it.

Step 10 — Merge 4: $(264, 107) \rightarrow 265$

Replace $(264, 107)$ everywhere:

Step 11 — Final encodings after training

| Pretoken | Token IDs | Decoded bytes | |—|—|—| | 'low' | [263] | low | | ' low' | [264] | ` low | | ’ lower’ | [265, 120] | low + e + r | | ’ lowest’ | [265, 121, 122] | lowe + s + t` |

The shared prefix ` lowe merged once into token 265. The suffix r in lower and st in lowest` stay as individual byte tokens because they appeared only once. That is the compression story — common substrings collapse, rare suffixes don’t.

Merge rank table (encoding priority):

Step 12 — Encode 'lowest' at inference time

The string 'lowest' has no leading space, so it pretokenizes to a single chunk 'lowest' (not in our training corpus, but the learned merges still apply).

Convert to byte IDs:

Greedy merge pass 1 — scan adjacent pairs, pick lowest rank:

Best pair: $(114, 117)$ with rank $0$. Apply merge $\rightarrow 262$:

Greedy merge pass 2:

All other adjacent pairs have no merge rule. Apply merge $\rightarrow 263$:

Greedy merge pass 3 — no adjacent pair exists in merges. Stop.

Final encoding:

One merged low token plus three byte tokens for est.

With BOS/EOS:

t.encode("lowest", add_bos=True, add_eos=True)
# [1, 263, 107, 121, 122, 2]

Round-trip check:

How does decoding work?

Decoding is the easy direction. Each token ID maps to a byte string in vocab. Concatenate, decode UTF-8:

For our encoding $[263, 107, 121, 122]$:

Special tokens can be skipped with skip_specials=True so <bos> doesn’t appear in generated text.

The invariant I test before wiring the dataloader:

The implementation skeleton

The full tokenizer lives in ~160 lines. Four pieces do the real work.

Pretokenize and bytes:

def _pretokens(text: str) -> list[str]:
    return PAT.findall(text)

def _word_to_ids(word: str, byte_base: int) -> list[int]:
    return [byte_base + b for b in word.encode("utf-8")]

Merge a pair in a sequence (left-to-right, non-overlapping):

def _merge_seq(ids: list[int], pair: tuple[int, int], new_id: int) -> list[int]:
    a, b = pair
    out, i, n = [], 0, len(ids)
    while i < n:
        if i < n - 1 and ids[i] == a and ids[i + 1] == b:
            out.append(new_id)
            i += 2
        else:
            out.append(ids[i])
            i += 1
    return out

Training loop — count pairs, pick best, merge everywhere:

while next_id < vocab_size:
    pair_counts = Counter()
    for w, freq in word_freq.items():
        for a, b in zip(word_ids[w], word_ids[w][1:]):
            pair_counts[(a, b)] += freq
    best, cnt = pair_counts.most_common(1)[0]
    if cnt < min_freq:
        break
    self.merges[best] = next_id
    self.vocab[next_id] = self.vocab[best[0]] + self.vocab[best[1]]
    for w in word_ids:
        word_ids[w] = _merge_seq(word_ids[w], best, next_id)
    next_id += 1
self.ranks = {pair: i for i, pair in enumerate(self.merges)}

Encoding — greedy lowest-rank merge:

while len(ids) >= 2:
    best_pair, best_rank = None, None
    for pair in zip(ids, ids[1:]):
        r = self.ranks.get(pair)
        if r is not None and (best_rank is None or r < best_rank):
            best_rank, best_pair = r, pair
    if best_pair is None:
        break
    ids = _merge_seq(ids, best_pair, self.merges[best_pair])

Save/load serializes specials and the merge list to JSON — no pickle, easy to inspect.

sequenceDiagram
  participant T as Text
  participant P as Pretokenizer
  participant B as Byte IDs
  participant M as Merge engine
  participant V as Vocab lookup

  T->>P: findall(regex)
  P->>B: UTF-8 encode each pretoken
  B->>M: greedy rank merges
  M->>V: final token IDs
  Note over V: decode = concat bytes → UTF-8

What vocabulary size should you use?

For my AL-1 code model I’m targeting 8k–16k merges on top of the 262-byte floor. Code has more distinct symbols than TinyShakespeare — operators, indentation, string delimiters. A 4k vocab works for prose; Python benefits from headroom.

Training data: I’m using a 100-file slice of the-stack-smol-xs Python split (~1 MB) for smoke tests, then scaling to a ~10 MB sample for the real merge table. You don’t need the full corpus to learn reasonable merges — frequency tails off fast.

Compare to GPT-2: 50,257 tokens. Same algorithm, bigger corpus, more merges.

Tokenizer design is also a model-efficiency trade. For vocabulary size $V$ and model width $d$, an untied input plus output embedding table costs about $2Vd$ parameters. Attention cost grows roughly with sequence length squared. A smaller vocabulary means fewer embedding parameters but longer sequences; a larger vocabulary shortens sequences but inflates the embedding matrices and creates rarer tokens that receive fewer gradient updates.

BPE vs WordPiece vs SentencePiece

For a decoder-only GPT clone, BPE is the path of least resistance — it’s what tiktoken implements for gpt2 encoding, and the merge-rank encode logic matches what you need for a from-scratch rebuild.

Common bugs I watch for

Forgetting the leading space in pretokens. "hello" and " hello" learn different merges. If your training data strips spaces but inference doesn’t, token counts drift.

Off-by-six in byte IDs. Special tokens shift everything. ord('a') is 97, but the token ID is 103 when byte_base=6. Every pair count depends on this.

Merge order mismatch. Encoding must use ranks (creation order), not frequency at encode time. Re-sorting merges breaks round-trip.

vocab_size too small. Below 262 the code refuses to start. Below ~1000 on code and you’re mostly bytes anyway — defeats the purpose.

Run round-trip tests on snippets from your actual corpus before you wire the tokenizer into the dataloader. Same advice as numerical gradient checking for autograd — catch data-pipeline bugs before they poison training.

Where this fits in the build

AL-1 pipeline so far: autograd engine → gradient checker → BPE tokenizer → token embeddings → positional encoding → transformer blocks.

The tokenizer is the bridge between raw .py files and the integer sequences the embedding layer consumes. The GPT-2 architecture post shows where those IDs go once they’re ready — wte, causal attention, lm_head with weight tying.

I built this tokenizer by hand because I wanted to stare at merge tables when generation looks wrong. When import tiktoken is one line, you lose visibility into whether the bug is the model or the bytes going in. For a toy GPT that might be overkill. For learning how these systems actually work, it’s the difference between using a library and owning the pipeline.

Train the merges. Check round-trips. Then trust your IDs.

Same three words. Same three token embeddings. Opposite meaning. Something has to tell the model where each token sits. That something is positional encoding. In GPT it’s a second lookup table living right next to the token embedding matrix. I’ll walk through why attention can’t see order on its own, what positional vectors actually are, the math of adding them, a full three-word worked example with every addition written out, learned vs sinusoidal encodings with formula steps, and where this lands in my AL-1 build.

Why is attention order-blind?

Self-attention computes a weighted sum over all tokens. The weights come from dot products between query and key vectors. A dot product is just

There is no position index $t$ inside that sum. The number $q_i \cdot k_j$ depends only on the vectors, not on whether $i$ is slot 0 or slot 99. Shuffle the tokens and you shuffle the same vectors into different slots — attention still computes the same pairwise dot products, just permuted. Without extra information, the model genuinely cannot tell position 0 from position 100.

Picture six people in a room — Alice, Bob, Charlie, David, Eva, Frank. You’re told who is present but not where they’re sitting. That’s a transformer with token embeddings alone: it knows the cast, not the seating chart. Language needs the seating chart.

flowchart LR
  subgraph blind["Token embeddings only"]
    a1["cat @ pos 0"] --> v1["vector V"]
    a2["cat @ pos 99"] --> v1
  end
  subgraph fixed["+ Positional encoding"]
    b1["cat @ pos 0"] --> w1["V + P0"]
    b2["cat @ pos 99"] --> w2["V + P99"]
  end

Same word cat, same embedding vector $V$. Only after adding a position-specific vector $P_t$ does slot 0 become $V + P_0$ and slot 99 become $V + P_{99}$ — two different inputs to attention.

What is a positional embedding?

The intuition is a second lookup table, parallel to the token table. Token embedding answers “what word is this?” Positional embedding answers “where does it sit?”

Stack all position vectors into a matrix

where $L$ is the maximum sequence length and $d$ is d_model. Row $t$ is the vector that means “I am the token at position $t$.”

GPT-2 small uses

Roughly 800k parameters whose only job is to encode order. It’s the same gather lookup as token embeddings, indexed by position ID instead of token ID.

Toy table with $L = 5$, $d = 2$:

The implementation is two lines longer than the embedding layer, because the indices are generated, not passed in:

class PositionalEmbedding:
    def __init__(self, max_seq, d_model, std=0.02):
        w = np.random.randn(max_seq, d_model) * std
        self.weight = Tensor(w)              # (max_seq, d_model), trainable

    def forward(self, seq_len):
        positions = np.arange(seq_len)       # [0, 1, 2, ..., seq_len-1]
        return self.weight.gather(positions, dim=0)   # (seq_len, d_model)

For a 5-token sentence, np.arange(5) makes [0, 1, 2, 3, 4] and gather pulls those five rows. Same $0.02$ GPT-2 initializer as the token embeddings — start small so training isn’t fighting huge activations on step one. Cost is $O(L)$ row selection, no matmul. Next to attention’s $O(L^2)$, positional lookup is rounding error.

The math: why we add, not concatenate

The transformer input is the sum of the two tables, elementwise:

where $E$ holds token embeddings and $P$ holds position embeddings, both with shape $(L, d)$.

Take a toy sentence of three tokens with $d = 2$:

Row 0 of $X$:

Row 1 of $X$:

Row 2 of $X$:

So

That tiny addition is the whole trick. After it, cat at position 0 and cat at position 99 are genuinely different vectors going into attention.

Why add two unrelated things instead of concatenating? Concatenation would stick $[1, 2]$ and $[0.1, 0.2]$ into $[1, 2, 0.1, 0.2]$, doubling $d$. Attention’s projection cost scales with $d^2$, so doubling width quadruples that cost and inflates every downstream weight. Addition keeps $d$ fixed — free position information, zero width tax. The model has 768 dimensions to play with; gradient descent learns to carve out subspaces where token identity and position don’t collide. The sum stays disentangleable enough for attention to read both.

Like token embeddings, $P$ starts as scaled noise. Backprop only updates the positions a batch actually used — a 5-token sequence sends gradients into rows 0–4 and leaves rows 5–1023 alone that step. Over training the rows specialize. Nobody hand-codes any of it; it falls out of next-token prediction, same as everything else.

Worked example: “dog bites man” vs “man bites dog”

I’ll use $d = 2$ so every addition fits on one line. Three token embeddings from the token embedding table:

Three learned position rows:

Sentence 1: “dog bites man”

Token at position 0 is dog. Combine embedding with position 0:

Token at position 1 is bites:

Token at position 2 is man:

Full input matrix for “dog bites man”:

Sentence 2: “man bites dog”

Same three token embeddings. Same three position rows. Different order — so each token meets a different position vector.

Token at position 0 is man:

Token at position 1 is bites:

Token at position 2 is dog:

Full input matrix for “man bites dog”:

What changed?

Compare row by row. Position 0 used to be dog at $[1.1, 2.0]$; now it’s man at $[0.1, 3.0]$. Position 2 used to be man at $[0.3, 3.1]$; now it’s dog at $[1.3, 2.1]$. Only position 1 (bites) lands on the same combined vector $[3.0, 1.2]$ — because bites stayed in the middle.

Without positional encoding, sentence 1 would be $[e_{\text{dog}}, e_{\text{bites}}, e_{\text{man}}]$ and sentence 2 would be $[e_{\text{man}}, e_{\text{bites}}, e_{\text{dog}}]$ — a permutation of the same three vectors. Attention is blind to that shuffle. With $P$ added, $X_{\text{dbm}} \neq X_{\text{mbd}}$ as matrices. Query-key dot products inside attention now produce different scores. The model can finally tell who did the biting.

Learned vs sinusoidal: the design fork

The original Attention Is All You Need paper used fixed sinusoidal encodings — no learned parameters, computed from a formula. For position $t$ and dimension index $i$ (with model width $d$):

Even dimensions get sines. Odd dimensions get cosines. Each dimension oscillates at a different wavelength because the divisor $10000^{2i/d}$ grows with $i$.

Worked numbers for $t = 1$, $i = 0$, $d = 4$:

For $t = 1$, $i = 1$:

So position 1 gets a unique fingerprint across all $d$ dimensions. Low-index dimensions swing fast; high-index dimensions swing slowly. Relative offsets between positions become linear functions of these encodings — that’s why sinusoids can, in principle, extrapolate to sequence lengths never seen in training. Zero parameters to learn.

GPT-2 threw that out and learned $P$ directly — the table above. Tradeoffs I weighed building AL-1:

• Learned: flexible, often marginally better on in-distribution lengths, dead simple to implement (it’s just another Embedding). Cost: hard-capped at $L = 1024$. Ask for position 1025 and there’s no row — the model can’t run.
• Sinusoidal: zero parameters, extrapolates past training length, but rigid and historically a touch worse on benchmarks.

For a toy GPT trained at fixed context length, learned wins on simplicity and I never need positions beyond $L$. I went learned. Modern models mostly moved past both to RoPE (rotates queries and keys by position inside attention rather than adding to the input) and ALiBi (distance-dependent score biases) — different mechanisms, same underlying problem.

Where this fits in the build

AL-1 pipeline: autograd → gradient checker → BPE tokenizer → token embeddings → positional encoding → attention.

The input to the first transformer block is $X = E + P$ — token identity plus position, both learned, both lookups. Every token now carries two answers at once: what am I and where am I. Strip $P$ out and the model reads “dog bites man” exactly like “man bites dog,” which is to say it reads nothing. That single 800k-parameter matrix, added at layer zero, is what lets the entire stack above it know order at all.

This is also where my two-table mental model finally locked in. Token embedding answers “what word is this?” Positional embedding answers “where does it sit?” Attention gets both in one vector and reasons over meaning and sequence simultaneously. Cheap to compute, easy to implement, and without it a transformer is a very expensive bag of words.

This post is the bridge between “text as numbers” and “text the model can reason over.” Get it wrong and token ID $824$ is just a bigger scalar than $391$, which would imply love > AI in some numeric sense — nonsense. I need each ID to become a vector I can do geometry on, where distance means something and cat lands near dog but far from car. A scalar ID cannot do that. A learned row in a matrix can.

flowchart LR
  TXT["'I love AI'"] --> TOK["BPE tokenizer"]
  TOK --> IDS["IDs [15, 824, 391]"]
  IDS --> EMB["Embedding lookup"]
  EMB --> VEC["Vectors (3, d_model)"]
  VEC --> TR["Transformer blocks"]

Why can’t the model use token IDs directly?

A token ID is a name, not a quantity. The tokenizer assigned love → 824 and AI → 391 by merge order during BPE training, not by meaning. If I feed raw integers into a matrix multiply, the network reads $824$ as “4.7× bigger than $175$” and starts fitting spurious magnitude relationships. The IDs carry zero semantic structure — cat and kitten could be $82$ and $49{,}001$, as far apart as any two random numbers.

What I want instead: each token becomes a vector in $\mathbb{R}^{d}$, where $d$ is the model width. Distances should mean something. Co-occurring tokens should end up near each other after training. You cannot get that from a scalar, so you map every ID to its own learned vector and let gradient descent arrange them.

What is a token embedding?

An embedding is a learned vector attached to each token in the vocabulary. Stack all of them and you get the embedding matrix $W$:

Here $V$ is vocabulary size and $d$ is d_model. Row $i$ is the vector for token ID $i$. In GPT-2 small:

Parameter count:

Almost 39 million parameters before attention. That is the single largest weight in a small GPT, and via weight tying it is reused as the output projection — the lm_head that maps hidden states back to vocabulary logits. Same shape $V \times d$, same meaning: “what a token means going in” tied to “how likely it is coming out.”

Take a toy vocabulary of three tokens with $d = 4$:

Row $0 \to$ dog. Row $1 \to$ cat. Row $2 \to$ car. Rows are tokens, columns are features. Nothing more.

How does the embedding lookup work?

Here is the entire layer from zyn/nn.py, give or take:

class Embedding:
    def __init__(self, vocab_size, d_model, std=0.02):
        w = np.random.randn(vocab_size, d_model) * std
        self.weight = Tensor(w)          # trainable

    def forward(self, indices):          # indices: (B, T) int array
        return self.weight.gather(indices, dim=0)   # (B, T, d_model)

gather is fancy indexing — weight[indices]. Given a tiny matrix:

And indices:

The output stacks rows $2$, then $0$, then $1$:

No multiplication. No activation. Input shape is $(B, T)$ — batch of $B$ sequences, $T$ tokens each. Output shape is $(B, T, d)$. Every integer becomes a $d$-dimensional vector, and that tensor is what the first transformer block consumes. For a real batch of shape $(32, 1024)$ with $d = 768$, the output is $(32, 1024, 768)$ floats.

Why multiply by 0.02 at initialization?

np.random.randn gives unit-variance Gaussians — values around $\pm 1$, occasionally $\pm 3$. Those are far too large to start. A row like $[1.5, -0.8, 2.1, \dots]$ fed into the residual stream blows up activations and the first few gradient steps thrash. Scaling by $\text{std} = 0.02$ shrinks everything:

A fresh row looks like $[0.012, -0.017, 0.006, \dots]$ — tiny, near the origin, all tokens starting roughly equal. That $0.02$ is the GPT-2 initializer. Training nudges rows apart slowly instead of starting from chaos.

How is lookup the same as a one-hot matrix multiply?

The equivalence that finally made embeddings click for me: a row lookup is exactly a one-hot vector times the matrix. Represent token $2$ as the one-hot row:

Multiply by a $3 \times 2$ matrix:

First product entry:

Second product entry:

Result:

That is precisely row $2$ of $W$. The zeros kill every other row; the single $1$ selects one. The embedding layer is a linear layer whose input happens to be one-hot. For $V = 50257$ you would multiply a 50k-wide vector that is all zeros except one slot. gather skips the charade and grabs the row directly. Same math, none of the floating-point busywork.

How do embeddings learn meaning?

The rows start as random noise. Meaning shows up through gradients. Say the model reads I love and predicts the next token. The loss $L$ flows back into the embedding matrix as $\frac{\partial L}{\partial W}$ — but only the rows that were actually used get a nonzero gradient. If the sequence touched token IDs $15$, $824$, and $391$, then rows $15$, $824$, and $391$ update. The other $50254$ rows sit untouched this step. That sparsity is why embedding training is cheap despite the parameter count.

Over millions of steps the geometry organizes itself. Tokens appearing in similar contexts get pulled together:

Nobody programmed dog ≈ cat. The objective — predict the next token — makes co-occurring tokens converge as a side effect.

flowchart TB
  subgraph space["Embedding space (schematic)"]
    dog["dog"] --- cat["cat"]
    cat --- lion["lion"]
    car["car"] --- truck["truck"]
  end

Why does d_model size matter?

With $d = 2$ each token has two numbers to encode everything — meaning, part of speech, tense, sentiment. Not enough room; tokens collide. Bump to $d = 768$ and each vector has 768 directions to spread information across. Wider embeddings hold more nuance, at the cost of more parameters and compute. GPT-2 small uses $768$; GPT-2 XL uses $1600$. It is a capacity dial, not a free lunch — double $d$ and you double this 39M-parameter layer.

Where does this fit in the AL-1 build?

Pipeline so far: autograd engine → gradient checker → BPE tokenizer → token embeddings → positional encoding → attention.

The tokenizer hands me IDs. The embedding layer turns them into the $(B, T, d)$ tensor that flows through every block. But this lookup is order-blind. The vector for cat is identical whether cat is the first token or the hundredth — gather does not know or care about position. A transformer with only token embeddings sees a bag of words. Positional encoding fixes that by adding a second matrix that injects “where” on top of “what.”

I implemented this by hand mostly to kill the mystery. nn.Embedding is one line and it hides the fact that the biggest weight in your model is a dictionary you are slowly teaching to mean something. Once you have watched the rows drift from noise into structure, the rest of the transformer stops feeling like magic and starts feeling like consequences.

Full worked example: I love AI with IDs [15, 824, 391]

I will walk through the entire embedding layer for I love AI — lookup, one-hot equivalence, loss gradient, and row update — with every arithmetic step on its own line. I use a toy width $d = 4$ so the numbers stay readable; GPT-2 small uses $d = 768$, but the mechanics are identical.

Step 1 — Token IDs from BPE

The BPE tokenizer maps:

Sequence:

Shape $(T) = (3)$ for one sentence. Batched with $B = 1$, shape is $(1, 3)$.

Step 2 — Rows in the embedding matrix

Pull the three relevant rows from $W \in \mathbb{R}^{50257 \times 4}$ (toy $d = 4$):

Step 3 — Embedding lookup via gather

gather stacks the rows in index order:

Shape $(T, d) = (3, 4)$. With batch dimension: $(1, 3, 4)$. Row $0$ of $X$ is the embedding for I. Row $1$ is love. Row $2$ is AI. This $(B, T, d)$ tensor is what the first transformer block receives.

Step 4 — One-hot equivalence for token 15 (I)

A one-hot vector for token ID $15$ in a vocabulary of size $V = 50257$ has a $1$ at index $15$ and $0$ everywhere else:

Multiply by $W$:

Only $j = 15$ contributes:

So:

Same result as gather. For love (ID $824$):

For AI (ID $391$):

Stacking all three one-hot rows and multiplying is the same as gather on $[15, 824, 391]$.

Step 5 — Loss gradient flowing back

Suppose the sentence passes through the transformer and the loss $L$ produces upstream gradients at each embedding position. After backprop through later layers, the gradient at position $0$ (I) is:

At position $1$ (love):

At position $2$ (AI):

Because $X[t] = W[\text{indices}[t]]$, the chain rule gives $\frac{\partial L}{\partial W[k]} = g_t$ when $\text{indices}[t] = k$, and $\mathbf{0}$ otherwise. Only three rows receive signal:

Every other row of $\frac{\partial L}{\partial W}$ is the zero vector. That is the sparsity I mentioned — one training step touches three rows out of $50257$.

Step 6 — Gradient matrix (schematic)

The full $\frac{\partial L}{\partial W}$ is mostly zeros. The three nonzero rows:

Step 7 — Row update with learning rate η = 0.1

SGD update: $W \leftarrow W - \eta \, \frac{\partial L}{\partial W}$. Only the used rows change.

Row 15 (I), component by component:

Updated row:

Row 824 (love), component by component:

Updated row:

Row 391 (AI), component by component:

Updated row:

Step 8 — What changed and what did not

After one SGD step on I love AI:

Rows $0$ through $14$, $16$ through $390$, $392$ through $823$, and $825$ through $50256$ are unchanged. Repeat this across millions of sentences and the geometry of the table drifts from random noise into structure — tokens that share contexts get pulled together, one sparse row update at a time.

If the same token ID appears twice in one sequence, its row gradient is the sum of the per-position gradients. That is how repeated tokens accumulate more signal in a single step.

Embedding closeness is useful for semantic search, retrieval, clustering, and recommendations. But a static token row cannot distinguish two meanings of the same word — bank as a river edge vs a financial institution share one row until the transformer layers above create contextual representations. That is a different problem, solved downstream. Here, at the embedding layer, the job is simple: turn each integer from BPE into a vector the rest of the model can work with. Everything else is gradient descent rearranging rows you actually used.

If your autograd is wrong, your entire GPT trains incorrectly. Numerical checking is the unit test framework for that engine. Below I’ll derive the central difference formula step by step, work a complete numeric example on $f(w) = w^2$, compare analytical vs finite-difference gradients on a tiny neural net, and show the Python gradcheck I run after every op in my autograd build.

What is a numerical gradient?

A gradient is a slope. For a function of one variable, it’s “how much does $f$ change when I nudge the input a tiny bit?”

Take the simplest function I know:

At $w = 3$, calculus gives the analytical derivative:

I don’t need calculus to estimate that slope. I can measure it by evaluating $f$ at two nearby points and asking: how much did $f$ rise per unit of horizontal step?

Pick a small step size:

Evaluate one step to the right:

Evaluate one step to the left:

The rise in $f$ across the interval:

The horizontal span of that interval:

Divide rise by run:

That matches the analytical answer exactly. This estimate — using points on both sides of $w$ — is the central difference formula.

graph LR
  subgraph "Central difference at w=3"
    L["f(3−h)\n8.9994"] --- C["w = 3\nslope ≈ 6"]
    C --- R["f(3+h)\n9.0006"]
  end
  style C fill:#d4edda,stroke:#009900,color:#111

How do you derive the central difference formula?

Start from the definition of the derivative:

That’s the forward difference — one point ahead, one at $w$. A symmetric version uses a point on each side:

Why is the symmetric version better? Taylor-expand $f(w + h)$ around $w$:

Taylor-expand $f(w - h)$:

Subtract the second from the first:

The $f(w)$ terms cancel.

The $\frac{h^2}{2} f’‘(w)$ terms also cancel — one positive, one negative.

What remains:

Divide both sides by $2h$:

The error is $O(h^2)$ — proportional to $h^2$, not $h$.

For forward difference, the analogous expansion leaves an $O(h)$ error term.

That means central difference gets accurate answers with a larger $h$, which matters when floating-point rounding noise starts fighting you.

Almost every gradient checker uses:

How does analytical vs numerical gradient compare on $f(w) = w^2$?

I’ll run the full comparison at $w = 3$ with $h = 0.0001$ and write every number on its own line.

Analytical path (calculus):

Numerical path (central difference):

Comparison:

On this simple function, the two paths agree to machine precision. That’s the pass condition I’d expect from a correct backward().

What does numerical gradient checking look like on a tiny neural net?

Calculus on $f(w) = w^2$ is a sanity check for the formula. The real use case is verifying autograd on an actual computation graph.

Consider one multiply — the smallest “network” I can write:

w = 2.0
x = 3.0

y = w * x
loss = y

What should autograd return? (analytical)

Forward pass:

Backprop from the loss:

Local derivative of multiply with respect to $w$:

Chain rule:

Autograd should set:

What does central difference return? (numerical)

Treat loss as a function of $w$ alone, holding $x$ fixed:

At $w = 2$ with $h = 0.0001$:

Analytical: $3$. Numerical: $3$. Pass.

What are we actually comparing after loss.backward()?

After loss.backward(), autograd gives me $\frac{\partial L}{\partial w}$ — say:

Numerical checking gives the same derivative from finite differences — say:

Absolute error:

Tiny. The backward implementation passes.

sequenceDiagram
  participant A as Autograd backward()
  participant N as Numerical checker
  participant C as Compare
  A->>C: w.grad = 3.000001
  N->>C: w_numeric = 2.999999
  C->>C: |error| < tolerance → pass

In practice I compare relative error, not just absolute — but the idea is the same: two independent paths to the same derivative should agree.

Why can’t you train with numerical gradients?

Because they’re insanely expensive.

Take a weight matrix W.shape = (768, 768). Parameter count:

For each parameter I need $L(\theta_i + h)$ and $L(\theta_i - h)$ — two forward passes per parameter.

For one gradient computation. On one layer.

Backprop does it in one forward + one backward. The gap isn’t a constant factor — it’s proportional to the number of parameters. I’d never finish training a GPT this way.

That’s why every framework trains with autograd and only reaches for finite differences when something smells wrong.

How do you build a gradcheck in Python?

Suppose I have a Tensor with .data and .grad. After loss.backward(), I check each element. Here is the loop I actually run — broken into the six steps I think through every time.

Step 1 — store the original value

original = w.data[i]

Step 2 — perturb up, run forward

w.data[i] = original + h
loss_plus = forward()

Step 3 — perturb down, run forward

w.data[i] = original - h
loss_minus = forward()

Step 4 — restore

w.data[i] = original

Step 5 — compute numerical gradient

grad_numeric = (loss_plus - loss_minus) / (2 * h)

Step 6 — compare

grad_auto = w.grad[i]
error = abs(grad_auto - grad_numeric)

For parameter $\theta_i$, the full comparison is:

Wrapped into a reusable function:

def gradcheck(forward_fn, params, h=1e-5, tol=1e-5):
    """Compare autograd gradients to central-difference estimates."""
    # Run forward + backward once to populate .grad
    loss = forward_fn()
    loss.backward()

    for param in params:
        flat_data = param.data.ravel()
        flat_grad = param.grad.ravel()
        for i in range(flat_data.size):
            original = flat_data[i]

            flat_data[i] = original + h
            loss_plus = forward_fn().data

            flat_data[i] = original - h
            loss_minus = forward_fn().data

            flat_data[i] = original  # always restore

            grad_numeric = (loss_plus - loss_minus) / (2 * h)
            grad_auto = flat_grad[i]

            denom = max(1.0, abs(grad_auto), abs(grad_numeric))
            rel_error = abs(grad_auto - grad_numeric) / denom

            if rel_error > tol:
                raise AssertionError(
                    f"gradcheck failed at index {i}: "
                    f"auto={grad_auto}, numeric={grad_numeric}, "
                    f"rel_error={rel_error}"
                )
    return True

flowchart TD
  S["Store θᵢ"] --> P["θᵢ ← θᵢ + h\nforward → L₊"]
  P --> M["θᵢ ← θᵢ − h\nforward → L₋"]
  M --> R["Restore θᵢ"]
  R --> N["(L₊ − L₋) / 2h"]
  N --> C{"|auto − numeric|\n< tol?"}
  C -->|yes| OK["✓ op passes"]
  C -->|no| FAIL["✗ fix backward()"]
  style OK fill:#d4edda,stroke:#009900,color:#111
  style FAIL fill:#f8d7da,stroke:#990000,color:#111

Loop over every element of every parameter tensor. Yes, it’s slow. That’s the point — I run this in tests, not in the training loop.

What should you test before you build GPT?

I’m working through AL-1 right now — a from-scratch GPT in pure NumPy. Before I touch LayerNorm, attention, softmax, cross-entropy, or AdamW, every primitive op needs to pass gradcheck(forward_fn, params):

MatMul is where I expect pain. Transpose conventions, batch dimensions, broadcasting in the backward pass — one wrong axis swap and every downstream gradient is garbage. I’d rather find that here with a 2×2 matrix than inside a 12-layer transformer.

If the checker fails on add, every later component inherits the bug. Fix the foundation first.

How do you choose epsilon and measure error?

For $f(x) = x^3$ at $x = 2$, the exact derivative is:

With $\epsilon = 0.01$:

Absolute error:

Relative error (the comparison I actually use in gradcheck):

Too-large $\epsilon$ causes truncation error because the curve is not locally linear across the interval. Too-small $\epsilon$ causes floating-point cancellation because $f(x + \epsilon)$ and $f(x - \epsilon)$ become almost equal. Around $10^{-5}$ to $10^{-4}$ is a useful starting range for float64; float32 often needs a larger step and looser tolerance.

I run checks in deterministic float64 mode, disable dropout, avoid points where ReLU or max is nondifferentiable, sample a few coordinates of huge tensors, and always restore perturbed parameters. Gradient checking validates derivative code; it does not prove that the forward equation, data, or objective is conceptually correct.

Where does this fit in the build?

This post is the sanity check that sits between understanding backprop from scratch and actually trusting my engine on real model code. I run gradcheck after implementing each op’s backward(), before moving to the next one. When all seven pass, I stop second-guessing whether my gradients are right and start worrying about things that actually matter — data pipelines, tokenizer bugs, learning rate schedules.

Numerical gradient checking won’t make my model smarter. But it will save me from the worst kind of debugging session: staring at a loss curve that looks almost right, tweaking hyperparameters that were never the problem, when a single botched backward() has been lying to me since step zero.

Build the checker first. Trust nothing until it passes.

Below is one complete training step of the original 2017 encoder-decoder transformer, solved like a single math problem. Every matrix multiply is written out. Every softmax row is summed. Every gradient lands on a number you can check with a calculator. The formulas are exactly those from Attention Is All You Need; the dimensions are a toy so you can verify every line by hand. If you’ve built autograd from scratch, this is the derivation you want open while debugging — and when something looks wrong, numerical gradient checking is how I sanity-check each layer.

[!NOTE] Conceptual attention (queries, keys, values, no arithmetic) lives in how attention works. This post is the arithmetic.

What machine am I solving?

The 2017 transformer is an encoder-decoder stack. The encoder reads the source sentence; the decoder reads a shifted target and attends to encoder memory. One layer on each side, one training step, every box below gets a number.

flowchart TB
    subgraph ENC["Encoder ×N"]
        EI["Input embed + PE"] --> ESA["Multi-head self-attention"]
        ESA --> EAN1["Add & Norm"]
        EAN1 --> EFF["Feed-forward"]
        EFF --> EAN2["Add & Norm"]
    end
    subgraph DEC["Decoder ×N"]
        DI["Output embed + PE"] --> DSA["Masked self-attention"]
        DSA --> DAN1["Add & Norm"]
        DAN1 --> DCA["Cross-attention<br/>(Q=decoder, K,V=encoder)"]
        DCA --> DAN2["Add & Norm"]
        DAN2 --> DFF["Feed-forward"]
        DFF --> DAN3["Add & Norm"]
    end
    EAN2 -->|"memory K,V"| DCA
    DAN3 --> LIN["Linear → vocab"]
    LIN --> SM["Softmax → probabilities"]
    SM --> LOSS["Cross-entropy + label smoothing"]

What toy setup am I using?

Task: translate je suis → i am. The decoder is fed the shifted target <bos> i and must predict i am.

All formulas are exact at any scale; only the numbers shrink. Vocabulary order: [<bos>, i, am, <eos>].

How do token embeddings enter the encoder?

Each token ID picks a row from the learned embedding matrix $E \in \mathbb{R}^{V \times d_{\text{model}}}$, then scales by $\sqrt{d_{\text{model}}}$.

For $d_{\text{model}} = 4$:

Encoder tokens je (position 0) and suis (position 1). After lookup and scaling, I round to clean integers for hand arithmetic (shapes stay $2 \times 4$):

Stack into the encoder input matrix $X_{\text{raw}} \in \mathbb{R}^{2 \times 4}$:

Row 0 is je. Row 1 is suis. Shape: $(2,\, 4)$.

What is positional encoding and how do I add it?

Sinusoidal positional encoding from the paper:

For $d_{\text{model}} = 4$, the frequency divisors are:

Position $pos = 0$:

Position $pos = 1$:

Add to embeddings elementwise. For this toy I round the sum to keep later arithmetic tractable (same as the original derivation):

Shape: $(2,\, 4)$. This $X$ is what enters encoder self-attention.

How does encoder self-attention compute Q, K, and V?

Multi-head attention projects $X$ into queries, keys, and values per head:

Head 1 weights ($X$ is $4 \times 2$ per row-column convention: $X \in \mathbb{R}^{2 \times 4}$, $W_Q \in \mathbb{R}^{4 \times 2}$):

Matmul: $Q = X W_Q^{(1)}$

Row 0 of $X$: $[1,\, 0,\, 1,\, 0]$.

Row 1 of $X$: $[0,\, 1,\, 0,\, 1]$.

Same $W_K^{(1)}$, so $K^{(1)} = Q^{(1)}$.

Matmul: $V = X W_V^{(1)}$

Shapes: $Q, K, V \in \mathbb{R}^{2 \times 2}$.

How do I score, scale, softmax, and mix values?

Matmul: $Q K^{\top}$

Scale by $\sqrt{d_k}$

Softmax row 0

Softmax row 1

By symmetry:

Each row sums to $1$.

Matmul: $O^{(1)} = A^{(1)} V^{(1)}$

Head 2 uses its own $W_Q^{(2)}, W_K^{(2)}, W_V^{(2)}$ and lands identically here with my chosen weights. Concatenate heads → width $4$, mix with $W_O = I_4$:

Shape: $(2,\, 4)$.

What is Add & Norm (layer norm)?

Residual connection plus layer norm:

With $\gamma = 1$, $\beta = 0$, $\epsilon = 0$.

Residual add after attention

Layer norm on row 0

Deviations from mean:

Normalized row 0:

Row 1 is antisymmetric:

How does the feed-forward network work?

Position-wise FFN:

$W_1 = [I_4 \mid I_4]$ widens $4 \to 8$. $W_2 = [I_4;\, I_4]$ projects $8 \to 4$. Biases are zero.

Row 0 through FFN

Input: $[1,\, -1,\, 1,\, -1]$.

After $W_1$:

After ReLU (negative entries → 0):

After $W_2$ (sum pairs):

Second Add & Norm (same recipe) yields encoder output — the decoder’s memory:

Shape: $(2,\, 4)$.

How does the decoder run masked self-attention?

Shifted target <bos> i, embedded and positioned (same rounding trick):

Masked self-attention sets future positions to $-\infty$ before softmax:

flowchart LR
    A["Raw scores"] --> B["future → −∞"]
    B --> C["softmax → 0 there"]
    C --> D["attend self + past only"]

Position 0 sees only itself → attention weights $[1,\, 0]$.

Position 1 sees positions 0 and 1 equally → weights $[0.5,\, 0.5]$.

After mixing, residual, and layer norm:

How does cross-attention connect encoder and decoder?

Cross-attention projections (chosen so every matmul is traceable):

Matmul: $Q = \text{DEC}_1 \, W_Q^{\text{cross}}$

Row 0 of $\text{DEC}_1$: $[1,\, 1,\, -1,\, -1]$.

Row 1: $[-1,\, -1,\, 1,\, 1]$.

Matmul: $K = \text{ENC} \, W_{K,V}^{\text{cross}}$

Row 0 of ENC: $[1,\, -1,\, 1,\, -1]$.

Row 1: $[-1,\, 1,\, -1,\, 1]$.

Matmul: $Q K^{\top}$

Scale and softmax

Row 0: $e^{1.414} \approx 4.10$, $e^{-1.414} \approx 0.243$, sum $\approx 4.34$.

Row 1 is symmetric:

Matmul: $O = A V$

FFN + second Add & Norm (same pattern as encoder). I score the final position — the one predicting am:

Shape: $(1,\, 4)$.

How do I project to vocabulary and get probabilities?

Linear layer (weight-tied to embeddings in the full model). Here $W_{\text{out}} = I_4$, bias $b = 0$:

Softmax over vocab [<bos>, i, am, <eos>]:

The correct word am (index 2) gets only $0.072$. The model is confidently wrong — good setup for a nontrivial gradient.

What is the cross-entropy loss with label smoothing?

Plain cross-entropy on am:

Label smoothing $\epsilon_{ls} = 0.1$ spreads mass onto every class:

Smoothed loss:

That is the scalar I backprop through.

How does backprop flow through the transformer?

flowchart RL
    Z["∂L/∂z = p − q"] --> WO["∂L/∂W_out"]
    Z --> H["∂L/∂h"]
    H --> LN["LayerNorm Jacobian"]
    LN --> F["FFN (ReLU gate)"]
    F --> AT["attention (softmax Jacobian)"]
    AT --> EMB["∂L/∂embeddings"]

Logits: softmax + cross-entropy collapse

The combined derivative is the elegant $p - q$:

Sum check (must be zero — softmax is shift-invariant):

Only am is negative → gradient pushes that logit up, pulls the rest down.

Output layer: $z = h W_{\text{out}}$

With $W_{\text{out}} = I$:

Layer norm backward

For $y_i = \gamma (x_i - \mu) / \sigma$ with $D$ dimensions:

The two sum terms strip gradient along the mean and variance directions — layer norm is not “just divide by std.”

FFN backward

ReLU gates the gradient to positions where pre-activation was positive:

Negative pre-activations get zero gradient. Dead ReLUs are real.

Attention backward

Forward: $O = AV$, $A = \text{softmax}(S)$, $S = QK^{\top} / \sqrt{d_k}$.

Softmax Jacobian (row $i$):

Then:

Masked positions had zero forward attention weight → zero gradient. Causality is preserved in backward pass too. Gradients from all positions and heads sum into shared weight matrices.

The one pattern underneath everything

If $Y = XW$ with $X : (N, d_{\text{in}})$, $W : (d_{\text{in}}, d_{\text{out}})$, $G = \partial L / \partial Y$:

Nearly every transformer gradient is this plus reshaping, layer-norm Jacobians, and the softmax Jacobian.

How does Adam update a weight?

Adam hyperparameters from the paper:

Learning rate schedule:

Take one scalar: $W_{\text{out}}[0,0] = 1$, gradient $g = \partial \mathcal{L} / \partial W_{\text{out}}[0,0] = 0.510$.

First moment at step $t = 1$:

Second moment:

Bias correction:

Adam step direction:

Step 1 (warmup, real $d_{\text{model}} = 512$)

Weight update:

Barely moves — warmup is conservative on purpose.

Step 4000 (end of warmup)

Now you can see the weight shift. Run millions of steps on billions of tokens and am eventually beats <bos>.

What invariants should I check while deriving?

I treat these as unit tests for the hand derivation:

• Write tensor shape beside every equation. Residual add requires identical shapes.
• Attention weights sum to $1$ across the key axis per query row.
• Masked positions get exactly zero probability forward and zero gradient backward.
• LayerNorm outputs have approximately zero mean and unit variance per token row.
• Softmax-cross-entropy gradients $\partial \mathcal{L} / \partial z$ sum to zero across the vocabulary.

Production code fuses kernels, uses mixed precision, checkpoints activations, and parallelizes — but none of that changes the calculus above. I use this derivation to debug toy implementations and to understand memory/compute flow before trusting a framework.

Where does one training step leave me?

Forward: embeddings → positional encoding → encoder attention (every matmul) → layer norm → FFN → decoder masked attention → cross-attention → vocab projection → softmax.

Loss: $\mathcal{L} = 2.53$ with label smoothing.

Backward: $\partial \mathcal{L} / \partial z = p - q$ → chain rule through every weight.

Update: Adam ticks $W_{\text{out}}[0,0]$ from $1.000000$ toward $0.999301$ by step 4000.

Next batch runs the same machine on slightly better weights. The whole thing is a matmul, a softmax, a normalize, a ReLU, an outer product, and the chain rule — stacked $N$ deep on each side, iterated over trillions of tokens.

Conceptual version: how attention works in transformers. Building blocks: token embeddings, positional encoding, autograd from scratch. Source: Attention Is All You Need (2017).

The useful part is the engineering underneath. The video stitches together several real ideas worth understanding on their own. I separate those ideas from the release story, work through the mathematics line by line, and show how I evaluate frontier-model claims before repeating them.

[!NOTE] New to how these models work under the hood? My write-up on how attention works sets up queries, keys, values, and the quadratic cost — the exact thing V4 spends most of its cleverness fighting. I’ll lean on it here.

How do I verify a DeepSeek V4 claim before citing it?

I run a short checklist before I repeat any frontier-model number:

1. Find the developer’s official announcement and paper.
2. Match the exact model name to a repository or model card.
3. Check whether parameter counts mean total or active parameters.
4. Separate maximum context support from measured retrieval quality.
5. Read benchmark methodology, baselines, sampling settings, and contamination discussion.
6. Treat third-party videos and screenshots as leads, not primary evidence.

Until those checks pass, the correct language is “the source claims,” not “the model uses” or “the paper proves.” The rest of this article follows that rule.

What does the source claim about DeepSeek V4?

Verification status: unconfirmed. The “1.6 trillion parameters,” “1 million tokens,” training-token count, architecture names, and benchmark results in this section are attributed to the third-party source. Do not cite them as official DeepSeek specifications without a primary source.

Two headline numbers, and both of them are the problem as much as the feature:

• 1.6 trillion parameters (the V4 Pro model). Parameters are the dials inside the model that store what it knows. More dials, more capacity — but also far harder to train without the whole thing falling over.
• 1 million tokens of context. That’s roughly 750,000 words — the entire Harry Potter series in the prompt, with the model still able to pull one specific number off one specific page. For agents, it means running for hours on a long task without losing the thread.

The catch is that neither of those is free. A 1.6T-parameter model wants to explode during training, and a 1M-token context wants to eat all your compute at inference. The source frames the whole design as a list of fixes for those two failure modes. I buy the engineering logic; I do not buy the release story without a primary source.

Why is a 1-million-token context so expensive?

Models don’t read like we do. For every new token, attention asks “how does this relate to every token before it?” — the mechanism from the 2017 Attention Is All You Need paper that everything since is built on.

For a 10-word sentence that’s ~10 comparisons. Nothing. But the cost grows with the square of the sequence length:

Plug in the numbers and the wall is obvious. For $n = 2{,}048$:

For $n = 100{,}000$:

For $n = 1{,}000{,}000$:

Double the context, quadruple the work. At a million tokens even high-end hardware chokes. And compute isn’t the only victim — the model has to store the intermediate keys and values for everything it’s seen, the KV cache, and that grows linearly with context until it swallows your memory budget whole.

flowchart LR
    A["New token"] --> B["Compare against<br/>every prior token"]
    B --> C["O(n²) attention cost"]
    B --> D["Store keys + values<br/>→ KV cache grows"]
    C --> E["Compute wall"]
    D --> E

So “just use a bigger context window” isn’t an engineering shrug. It’s the central bottleneck of the entire field.

How does sparse attention cut the pair count?

Dense causal attention over $n$ tokens forms roughly $n(n+1)/2$ usable query-key pairs — each query can only look at tokens at or before its position.

Start with $n = 1{,}000{,}000$:

That is about $5 \times 10^{11}$ pairs per head and layer in the dense case.

A sparse design lets each query attend to only a local window of width $w$ plus $r$ retrieved distant tokens. Pair count drops toward $n(w+r)$. Take the source’s illustrative values $w = 128$ and $r = 32$:

That is $1.6 \times 10^{8}$ pairs. Compare to the dense count:

Over 3,000 times fewer pairs on paper. This does not prove any claimed implementation achieves that speedup — routing, retrieval, memory movement, and kernel utilization add overhead. It shows why sparse patterns are attractive enough that I’d expect every long-context design to use something like them, whether or not V4 exists.

How much memory does a 1M-token KV cache actually need?

The KV cache stores one key vector and one value vector per token, per key/value head, per layer. For $L$ layers, $h_{kv}$ key/value heads, head dimension $d_h$, sequence length $n$, and $b$ bytes per stored value:

The leading factor $2$ is keys plus values. I’ll walk through a plausible toy config at $n = 1{,}000{,}000$ with FP16 ($b = 2$ bytes):

Step 1 — layers × heads:

Step 2 — multiply by head dimension:

Step 3 — multiply by sequence length (one vector per token):

Step 4 — keys + values:

Step 5 — bytes (FP16):

Step 6 — convert to GiB (divide by $1024$ three times):

One sequence. One user. No activations, no weights, no optimizer state — just the cache. Quantization, grouped-query attention, eviction, compression, or distributed storage is therefore essential at million-token scale. The math doesn’t care what the marketing slide says.

What is the claimed fix: three views of the same document?

The source says DeepSeek’s move isn’t to pick one clever attention trick — it’s to run three in parallel and let each cover a different weakness. Think of them as three simultaneous readings of the same text:

1. Compressed long-range. The distant past gets aggressively squeezed — blocks of old tokens collapsed into compact summaries so the model can still reach far back without paying full price for every token.
2. Selective retrieval. Compression alone loses precision. Ask about an exact phrase or a specific number buried half a million tokens ago and a one-block summary has thrown it away. So a second pathway selectively pulls back the specific distant tokens that matter, uncompressed.
3. Sliding-window. The most recent stretch — say the last 128 tokens — is kept word-for-word, zero compression, full fidelity. Recent context is usually the most relevant, so it gets the most exact treatment.

flowchart TB
    DOC["1M-token sequence"] --> C1["Compressed<br/>(distant past → summaries)"]
    DOC --> C2["Selective retrieval<br/>(exact distant tokens)"]
    DOC --> C3["Sliding window<br/>(last ~128 tokens, exact)"]
    C1 --> MIX["Combine per layer"]
    C2 --> MIX
    C3 --> MIX
    MIX --> OUT["Long-range recall +<br/>broad context +<br/>sharp local focus"]

By interleaving these layer by layer, the source claims V4 gets long-range retrieval, broad understanding, and sharp immediate focus at once — while sidestepping the quadratic blow-up. The payoff the video reports: against DeepSeek V3.2 (already one of the most compute-efficient models around), V4 Pro runs at 3.7× lower FLOPs — despite being bigger and having a far longer context. That’s the part I’d want to see in a reproducible benchmark before I repeat it.

How do you keep a trillion-parameter signal from exploding?

Depth is the other monster. Signals pass through hundreds of layers, and at every layer there’s a risk the values blow up or vanish. The classic fix is residual connections — bypass lanes that let the signal skip around a layer instead of being mangled by it. More recently the field moved to hyperconnections, which widen those lanes to give the signal more room.

The source’s finding: past a trillion parameters, even standard hyperconnections start throwing catastrophic spikes. The lanes aren’t enough. So they introduced manifold-constrained hyperconnections (MHC) — constrain the residual mixing onto a manifold of doubly stochastic matrices, where every row and every column sums to 1. Intuitively that means the layer can reroute and remix signal across lanes but can never amplify the total, so it can’t run away to infinity.

The problem is enforcing that constraint a trillion times. Their answer is the Sinkhorn–Knopp algorithm: before each layer, run a quick sequence of row-then-column normalizations — about 20 iterations — until the matrix satisfies the doubly-stochastic constraint.

flowchart LR
    M["Raw mixing matrix"] --> R["Normalize rows"]
    R --> C["Normalize columns"]
    C --> Q{"Doubly stochastic?<br/>(~20 iters)"}
    Q -- "no" --> R
    Q -- "yes" --> L["Layer processes signal"]

Worked example: one Sinkhorn–Knopp step on a $2 \times 2$ matrix

Start with a positive matrix:

Row sums:

Row-normalize (divide each entry by its row sum):

Column sums of $A^{(1)}$:

Columns do not sum to 1 yet. Column-normalize (divide each column by its column sum):

Row sums are now off again. Alternating row and column normalization converges toward a matrix where both sums are one. Such a constraint can control mixing magnitude, but stability claims still require ablations, training curves, and reproducible code — none of which I’ve verified for V4.

At first glance bolting a 20-step loop onto every layer of a 1.6T-parameter model sounds insane — you just spent the whole design saving compute. The source claims fused GPU kernels, selective recomputation, and other deep GPU-level tricks shrink the entire overhead to 6.7% of runtime. Plausible in principle; unverified for this specific model.

How do you train a 33T-token run without the math blowing up?

Architecture is only half the story. The source says V4 Pro was trained on 33 trillion tokens, more text than a human could read in many lifetimes.

Dump all of that on a fresh model at once and it drowns, the way handing a beginner an advanced quantum mechanics textbook teaches them nothing. So the source describes a curriculum — start small, grow the context as the model stabilizes:

flowchart LR
    S["4K tokens<br/>(grammar, syntax,<br/>local patterns)"] --> M1["16K"]
    M1 --> M2["64K"]
    M2 --> L["1M tokens<br/>(full long-range)"]

Early on the model only sees ~4,000-token windows — enough to learn local structure — then 16K, 64K, all the way to the full million. Curriculum learning at long context is a real technique; the specific schedule here is a claim.

At this scale the big threat is a loss spike: the math blows up and the run crashes. The usual industry fix is grim — stop, roll back to an earlier checkpoint, try again, burn the compute. The source says DeepSeek leaned instead on self-stabilizing training plus the Muon optimizer (the algorithm deciding how to nudge the dials when the model gets something wrong), squeezing more stable learning out of each step so the run doesn’t need babysitting.

What benchmarks does the source report?

Verification status: unconfirmed. Treat every bullet as “the video claims,” not “I measured.”

• Putnam 2025 — 120/120. A perfect score on one of the hardest undergraduate math competitions in the world.
• 1M-token retrieval beats Gemini 3.1 Pro. Pushed to the absolute limit of its context window, V4’s recall accuracy edges out Google’s latest, which also runs a million tokens.
• Outcompetes Opus 4.6 across the task suite the source reports.
• On the independent Artificial Analysis leaderboard it lands as the second-best open-source model, just under Kimi 2.6 and crowding the top closed models from Google, Anthropic, and OpenAI.

I’d want the eval scripts, the needle-in-a-haystack prompts, and the contamination discussion before I repeated any of this in a paper or a tweet.

Why is the engineering story worth understanding anyway?

What got me isn’t any single trick — it’s that there is no single trick. Interleaved sparse attention for the memory wall, manifold-constrained hyperconnections for the signal explosions, Sinkhorn–Knopp made cheap with fused kernels, a curriculum and a stable optimizer so a 33T-token run doesn’t die halfway. Dozens of small, sharp engineering decisions stacked until something impossible becomes routine. That’s the whole theme: not one breakthrough, a lot of disciplined ones.

And if the source is right about the open-weight release, they did it resource-constrained — smaller team, no top NVIDIA chips, a sliver of the compute the closed labs swing around — then put the weights on Hugging Face and wrote down the infrastructure details the closed labs treat as state secrets. For a field that’s been quietly closing up, a paper that just tells you how is the rarest thing in it. I haven’t confirmed that release happened; the engineering patterns are real either way.

Strong intuition means knowing what the equations imply and knowing what the available evidence does not establish. The sparse-attention math says million-token context is a memory and compute problem with known structure. The verification checklist says I still don’t know whether DeepSeek V4 solved it the way the video says.

Watch: The insane engineering of DeepSeek V4 — AI Search. If the attention mechanics felt fuzzy, start with how attention works in transformers.

What finally made it line up for me was 3Blue1Brown’s Attention in transformers. This is my write-up of that mechanism — one continuous story from intuition through a fully worked three-token example, then masking, the low-rank value trick, the real GPT-3 parameter counts, and multi-head attention — in the order I wish someone had handed me.

[!NOTE] If tokens and embeddings are fuzzy for you, read how GPT actually works first — it sets up why directions in vector space carry meaning. I’ll lean on that here.

What problem does attention actually solve?

Right after the embedding step in how GPT actually works, a token’s vector is just a context-free lookup. Take the word mole:

• American shrew mole — the animal
• One mole of carbon dioxide — the chemistry unit
• Take a biopsy of the mole — the skin growth

All three start with the exact same vector. That’s obviously wrong — meaning depends on what’s around it. Attention is the step where the surrounding tokens get to reach in and update the mole vector toward the right sense.

And these updates travel far. Picture feeding in a whole mystery novel that ends with:

Therefore the murderer was …

To get the next word right, that final vector — which started life just meaning was — has to have absorbed information from thousands of tokens back. Attention is the plumbing that moves that information across the whole sequence.

How do tokens ask each other questions?

The example I keep coming back to is “a fluffy blue creature.” The adjectives fluffy and blue should update the noun creature.

Imagine each noun broadcasting a question: “are there any adjectives in front of me?” And each adjective broadcasting an answer: “yes — I’m an adjective, I describe nouns.” Where a question and an answer match, information flows. Make that matching numerical and you’ve got the entire mechanism. Three players:

What is a query?

For each token, multiply its embedding by the query matrix $W_Q$ to get a query vector:

The query lives in a much smaller space than the embedding. In GPT-3 the embedding is 12,288-dimensional but the query/key space is only 128. So $W_Q$ is a $128 \times 12{,}288$ matrix that squeezes each token’s meaning down into a compact “here’s what I’m hunting for.” The query for creature roughly says “I’m a noun, any adjectives ahead of me?”

What is a key?

In parallel, multiply every embedding by the key matrix $W_K$ to get key vectors:

Same 128-dimensional space. Think of keys as potential answers to the queries — you want fluffy’s key to land close to creature’s query.

“Close” means large dot product. For every (key, query) pair you compute $\vec{K}_i \cdot \vec{Q}_j$: large positive = strong match (token $j$ attends to token $i$), near zero or negative = unrelated.

How do dot products become attention scores?

Lay all those dot products out as a grid. If you have $T$ tokens, you get a $T \times T$ matrix of raw scores. In matrix form, every key dotted with every query is one product:

Entry $(i, j)$ is key $i$ dotted with query $j$ — “how much should token $j$ listen to token $i$?”

flowchart TB
    subgraph GRID["Attention scores (dot products)"]
        direction LR
        R1["fluffy·creature = large"]
        R2["blue·creature = large"]
        R3["the·creature ≈ 0"]
    end
    GRID --> SCALE["divide by √d_k"]
    SCALE --> SM["softmax each column"]
    SM --> AP["Attention pattern<br/>(weights sum to 1 per column)"]

Why do we divide by $\sqrt{d_k}$?

Raw dot products grow with the dimension $d_k$. If each coordinate of query and key has variance near 1, their dot product sums $d_k$ terms and has variance near $d_k$. Big scores shove softmax into a nearly one-hot corner where gradients vanish. The fix is to scale each score down:

In GPT-3, $d_k = 128$, so the divisor is:

How does softmax turn scores into weights?

Raw scores can be negative and they don’t sum to 1. Softmax fixes that — applied down each column of the score matrix, so for every token $j$ the weights over all source tokens $i$ land in $[0, 1]$ and add up to 1.

For one column with scaled scores $s_1, s_2, \ldots, s_T$, the first step is exponentiate each score:

Then sum them into a normalizer:

Each weight is one exponent divided by that sum:

The full matrix of weights is the attention pattern.

What is a value?

Scores tell you how much to mix. Values are what you mix. Multiply each embedding by the value matrix to get value vectors:

Then each token’s updated output is a weighted sum of all value vectors, using the attention weights from its column:

That weighted sum is the heart of attention — you don’t pick one token; you blend them.

What is the full attention formula?

Once you’ve built it piece by piece, the original paper’s one-liner is just shorthand:

Read it left to right: dot all keys with all queries, scale, softmax each column, blend the values. That’s the whole mechanism for one head.

Can we walk through a complete three-token example?

Yes — and I think this is the step that finally made the formula feel real. I’ll use three tokens — fluffy, blue, creature — with tiny 2-dimensional query/key/value vectors so every arithmetic step fits on the page. (Real GPT-3 uses $d_k = 128$; the logic is identical.)

Step 1 — the inputs. After multiplying by $W_Q$, $W_K$, and $W_V$:

Here $d_k = 2$, so $\sqrt{d_k} = \sqrt{2} \approx 1.414$.

Step 2 — dot products for the creature column. Token creature is the one hunting for adjectives, so we build column 3 of the score matrix. Key fluffy dotted with query creature:

Key blue dotted with query creature:

Key creature dotted with query creature:

So the raw scores for creature’s column are $[3,\, 2,\, 4]$.

Step 3 — scale by $\sqrt{d_k}$. Divide each raw score by $1.414$:

Step 4 — softmax, one line per step. Exponentiate each scaled score:

Add them for the normalizer:

Divide each exponent by $Z$ to get the attention weights:

Check: $0.284 + 0.140 + 0.576 = 1.000$. Good — they sum to 1.

Step 5 — weighted sum of values. Blend the three value vectors using those weights:

First term:

Second term:

Third term:

Add them:

That’s it. Creature didn’t copy any single neighbor — it mixed information from fluffy, blue, and itself, weighted by how well each key matched its query. The biggest weight (0.576) went to creature itself, but meaningful chunks came from the adjectives. Run the same five steps for every column and you update all three tokens.

$K^{\top} Q$ decides where to read. Softmax turns scores into a distribution. Multiplying by $V$ decides what information comes back.

Why can’t the model peek at future words? (masking)

During training the model predicts the next token at every position at once — way more efficient than one prediction per pass. But that opens a cheating door: when predicting the token after position 5, it must not look at positions 6, 7, 8…

The fix is masking. Before softmax, set any score where a later token would influence an earlier one to $-\infty$:

Softmax then turns those entries to exactly zero. Each token attends only to itself and tokens before it.

flowchart LR
    A["Raw score grid"] --> B["Set 'future' entries<br/>to −∞<br/>(upper triangle)"]
    B --> C["Softmax →<br/>those become 0"]
    C --> D["Each token attends<br/>only to itself<br/>and earlier tokens"]

In GPT-style models masking is always on — it’s what makes generation strictly left-to-right.

There’s a hidden cost too. The attention pattern has one row and one column per token, so its size is the square of the context length:

Double the context, quadruple the grid. That quadratic blow-up is the real reason long context windows are hard.

What is the low-rank value trick?

Scores tell you how much to mix; values tell you what to mix. Naively, the value map goes from embedding space back to embedding space, so one full matrix would be:

~151 million for one value matrix in one head — absurdly more than the ~1.5M each for query and key. So GPT-3 factors the value map into two smaller matrices: a value-down ($12{,}288 \to 128$) and a value-up ($128 \to 12{,}288$):

In linear-algebra terms that constrains the map to a low-rank transformation. The payoff: all four maps in one head now cost about the same.

How many parameters are in one GPT-3 attention head?

GPT-3 numbers — embedding 12,288, key/query/value dim 128:

All four come out the same size — which is the whole point of that low-rank factorization.

[!NOTE] Self- vs cross-attention: everything above is self-attention — queries, keys, and values all come from one sequence. Cross-attention (translation, speech-to-text) is nearly identical; the only difference is keys and values come from a different sequence than the queries.

What is multi-head attention?

One head learns one kind of relationship — say adjective→noun. Real language is richer than that, so a transformer runs many heads in parallel, each with its own $W_Q, W_K, W_V$:

• GPT-3 uses 96 heads per block.
• 96 × ~6.3M ≈ 600 million parameters per multi-head attention block.

flowchart TB
    E["Token embeddings"] --> H1["Head 1<br/>(grammar?)"]
    E --> H2["Head 2<br/>(subject–object?)"]
    E --> H3["…"]
    E --> H96["Head 96<br/>(coreference?)"]
    H1 --> CAT["Combine all heads"]
    H2 --> CAT
    H3 --> CAT
    H96 --> CAT
    CAT --> OUT["Updated embeddings"]

Each head makes its own attention pattern and its own value mix; the results get concatenated and projected back into the embedding space. Different heads end up tracking different things — grammar, coreference (“Harry” → Harry Potter vs Prince Harry depending on nearby words), physical implications (“they crashed the car”), and so on.

For tensors: input $X$ has shape $(B, T, d)$; per head, $Q, K, V$ have shape $(B, h, T, d_h)$; scores have shape $(B, h, T, T)$; the mixed result has shape $(B, h, T, d_h)$. Concatenating all $h$ heads restores $(B, T, d)$.

How does attention fit into the full transformer?

A transformer doesn’t run attention once. Data flows through attention, then an MLP (feed-forward) block, then attention again, then another MLP — dozens of times. Each round gives the now-more-nuanced embedding another chance to be refined by its now-more-nuanced neighbors. That depth is where the model’s apparent “reasoning” quietly piles up.

flowchart LR
    A["Embeddings"] --> B["Multi-head<br/>Attention"]
    B --> C["MLP /<br/>Feed-Forward"]
    C --> D["Multi-head<br/>Attention"]
    D --> E["MLP /<br/>Feed-Forward"]
    E --> F["… ×96 layers …"]
    F --> G["Final embeddings<br/>→ predict next token"]

Attention retrieves and mixes information across tokens, but it can’t do all the computation by itself — the surrounding MLPs, residual paths, layer normalization, and depth matter just as much. Standard attention also stores a $T \times T$ score matrix, so memory grows quadratically with context length.

Where does this leave you?

The thing I didn’t expect is how small the core idea turns out to be: dot product → scale → softmax → weighted sum, repeated across heads and layers. Queries, keys, and values are just “what I’m looking for / what I offer / what I’ll hand over.” Masking keeps it causal. The low-rank value trick keeps parameter counts sane. The quadratic grid is the price you pay for letting every token talk to every other token.

After sitting with this, the original Attention Is All You Need paper finally reads like a summary of something I already understand instead of a wall of notation.

If you haven’t read the setup, start with how GPT actually works.

Watch: Attention in transformers — 3Blue1Brown · full text version on 3blue1brown.com · part of the Neural Networks series.

The thing that finally fixed that for me was 3Blue1Brown’s But what is a GPT?. So this is me writing the explanation I wish I’d had — one continuous story from raw text to a sampled next token, with every calculation step on its own line and the real GPT-3 numbers attached wherever the full model would be too big to write out.

What does a GPT actually do?

GPT = Generative Pre-trained Transformer. Generative: it writes text. Pre-trained: it learned from a huge pile of data first. Transformer: the architecture from the 2017 paper Attention Is All You Need.

But the loop underneath all of that is embarrassingly simple. Given some text, the model outputs a probability distribution over the next token. Pick one token from that distribution, glue it on, and run again:

flowchart LR
    A[Input text] --> B[Model predicts<br/>distribution over<br/>next token]
    B --> C[Sample one token<br/>from the distribution]
    C --> D[Append token<br/>to the text]
    D --> A

That predict → sample → append loop is what you watch when ChatGPT types one word at a time. Everything else in this post is just the math inside the “predict” box — and I’m going to walk it in order: tokens → embeddings → dot products → softmax → temperature, then tie it together with one complete numeric example.

How does text become tokens?

Before any number touches the transformer, the string gets chopped into tokens — sometimes whole words, often sub-word pieces. Take the sentence from the video:

To date, the cleverest thinker of all time was …

A tokenizer might split it like this:

**To

date

,

the

cle

ve

rest

thinker

of

all

time

was …**

“cleverest” becomes cle | ve | rest. The model never sees letters — only integer token IDs. In GPT-3 the vocabulary has exactly:

I’ll use clean word-sized tokens in the worked example below so the arithmetic stays readable, but the real unit is always the token ID.

How do token IDs become vectors?

Each token ID gets mapped to a vector — a long list of numbers — by looking up a column in the embedding matrix $W_E$.

In GPT-3:

Each column is one token’s vector. “Embedding token $t$” means: grab column $t$ of $W_E$. The values start random and get learned during training. That one matrix alone holds:

617 million parameters before a single attention layer — and the full model has about 175 billion. I kept a running tally while watching the video, and that’s what made the scale click for me.

What does training do to those vectors?

After training, directions in embedding space mean something. Nobody programmed this; gradient descent just found it useful.

• Gender lives in a direction. $\text{woman} - \text{man}$ is roughly $\text{queen} - \text{king}$, so $\text{king} + (\text{woman} - \text{man})$ lands near queen.
• Nationality generalizes. $\text{Italy} - \text{Germany} + \text{Hitler}$ lands near Mussolini.
• Plurality is a direction. $\text{cats} - \text{cat}$ points along “more than one.”

A raw embedding is still dumb on its own — the lookup for mole is identical whether you mean the animal, the skin mark, or the chemistry unit. The transformer’s middle layers fix that by letting each vector soak up context from its neighbors. GPT-3 can hold 2,048 tokens at once; after 96 attention + MLP blocks, the last token’s vector has absorbed information from everything before it. I’ll save how that mixing works for the attention deep dive — but the tool it uses everywhere is the dot product.

Why are dot products the similarity score?

To ask “how aligned are two vectors?” you take the dot product:

Positive → they point the same way (related). Zero → perpendicular (unrelated). Negative → opposite.

Concrete 2-D example. Let:

Multiply coordinate by coordinate:

Sum:

In GPT-3’s embedding space each vector has 12,288 coordinates, so the sum has 12,288 terms — same idea, just longer. Attention uses dot products to score “how much should token $j$ listen to token $i$?” — that’s the engine I unpack in how attention works in transformers. For predicting the next token, the model eventually takes the last context-rich vector and runs one more matrix multiply.

How does the model score every possible next token?

After all the blocks, take the final vector $\vec{h}$ for the last token — shape $12{,}288 \times 1$ in GPT-3. Multiply by the unembedding matrix $W_U$:

Each entry of $\vec{z}$ is a raw score — a logit — for one vocabulary token. You get 50,257 of them. $W_U$ is basically $W_E$ flipped, so it adds another:

Just over 1 billion parameters on the two ends alone; the other ~174 billion sit in the attention and MLP blocks in the middle.

Why only the last vector at inference? During training the model predicts the next token at every position in one pass — thousands of learning signals. At inference you only care about what comes next, so you read off the last one.

What does softmax turn those scores into?

Logits can be any real number and they don’t sum to 1. Softmax converts them into a valid probability distribution. For logits $z_1, z_2, \ldots, z_V$:

Every output lands in $(0, 1)$ and the outputs sum to 1. I’ll walk a tiny three-token case so every step is visible.

Suppose the vocabulary is [cat, dog, runs] and the logits are:

Step 1 — exponentiate each logit.

Step 2 — sum for the normalizer.

Step 3 — divide each exponential by $Z$.

Check the sum:

Greedy decoding picks cat (highest probability). Sampling rolls a weighted die with those three probabilities. In the real model, the same softmax runs over all 50,257 tokens.

What does temperature change in those probabilities?

Temperature $T$ rescales logits before softmax:

Low $T$ sharpens the distribution toward the top token — safe, repetitive, near-deterministic. High $T$ flattens it — more surprising, more creative, occasionally unhinged. Same logits $\vec{z} = [2, -1, 1]$, now with $T = 0.5$.

Step 1 — divide each logit by $T$.

Step 2 — exponentiate.

Step 3 — normalizer.

Step 4 — probabilities.

Compare to $T = 1$: cat went from 70.5% to 87.9%. The API slider I’d been blindly nudging is just this — rescale logits, then softmax.

Can we walk through one complete next-token prediction?

Yes. I’ll trace the full pipeline on a toy prompt, naming the GPT-3 dimensions wherever the real model would use them.

Prompt: The cat

Step 1 — tokenize

(GPT-3 would assign integer IDs; I’ll keep the words.)

Step 2 — embed each token

Toy embedding dimension $d = 2$ instead of 12,288. Suppose $W_E$ has learned:

In GPT-3 this lookup is a column grab from a $12{,}288 \times 50{,}257$ matrix — same operation, 12,288 numbers per token.

Step 3 — mix context (attention + MLP, ×96)

The middle of the model updates every vector so cat knows it follows The. After all blocks, the last vector might become:

This is where attention does the heavy lifting — dot products between queries and keys, softmax over scores, weighted sum of values. I’m skipping the 96-layer detail here; the output is one context-soaked vector $\vec{h}$.

Step 4 — unembed to logits

Toy vocabulary: [cat, dog, runs]. Unembedding matrix:

Row 0 → cat, row 1 → dog, row 2 → runs. Each logit is a dot product of one row with $\vec{h}$:

In GPT-3, $W_U$ has 50,257 rows and each logit is a 12,288-term dot product.

Step 5 — softmax at $T = 1$

Step 6 — sample and append

Sample runs (26% chance), append it: The cat runs. Feed the longer string back in. Repeat.

Step 7 — what training does to wrong guesses

If the correct next token were runs but the model gave $P(\text{runs}) = 0.2595$, the cross-entropy loss is:

The gradient on logits is $p - y$:

Push down cat and dog, push up runs. Backprop carries that through $W_U$, every transformer block, and $W_E$. GPT is not a database returning stored sentences — it’s a parameterized probability machine estimating $P(\text{next token} \mid \text{previous tokens})$ over and over.

What is the whole pipeline in one picture?

flowchart LR
    A["Prompt text"] --> B["Tokenize<br/>(50,257 vocab)"]
    B --> C["Embed via W_E<br/>(617M params, d=12,288)"]
    C --> D["Attention + MLP<br/>blocks ×96<br/>(dot products inside)"]
    D --> E["Take last vector"]
    E --> F["Unembed via W_U<br/>(617M params)"]
    F --> G["Logits → Softmax<br/>(÷ temperature)"]
    G --> H["Probability over<br/>50,257 tokens"]
    H --> I["Sample next token"]
    I --> A

Underneath all of it, one mental model kept the architecture from feeling like magic: parameters are weights, and nearly every step is a matrix-vector multiply with occasional nonlinearities (like softmax) wedged in so the stack doesn’t collapse into one boring linear map.

Where should I go next?

The middle of that diagram — the part where dot products decide which tokens talk to which — is attention. I wrote the full walkthrough with queries, keys, values, masking, and the real GPT-3 head counts here: how attention works in transformers →.

Watch: But what is a GPT? — 3Blue1Brown · full text version on 3blue1brown.com · part of the Neural Networks series.

That changed after I watched 3Blue1Brown’s Essence of Linear Algebra.

What is linear algebra, really?

Linear algebra is the math of lists of numbers and machines that update those lists in a predictable way. A vector can be an arrow on a plane, but in machine learning it is more often a description of something: a house might be $x=[1200,3]$ for square feet and bedrooms; a token might be a list of 768 learned features. A matrix is the machine: feed it a vector, get another vector out. Grant Sanderson’s series made that picture stick for me — a matrix is not a grid of symbols, it is a transformation of space that stretches, rotates, or shears every input direction. Once I held that image, eigenvalues, determinants, and change of basis stopped feeling like magic notation and started feeling like consequences of the same idea.

flowchart LR
  V["Input vector x"] --> M["Matrix A"]
  M --> O["Output vector Ax"]
  O --> D["Dot products measure alignment"]
  O --> E["Eigenvectors = directions that only scale"]

The rest of this post walks one small $2\times2$ example from start to finish — vectors, matrix multiply, dot product, determinant, eigenvalue — with every arithmetic step written out. I will come back to where these ideas show up in models like the ones I am building in AL-1.

What is a vector?

A vector is an ordered list of numbers. In two dimensions I write

The first entry is how far I move along the horizontal axis; the second is how far along the vertical axis. I can read the same object as an arrow from the origin to the point $(4,5)$, or as a list of features. Both views are the same data; geometry gives me intuition, lists give me code.

How does a matrix transform a vector?

A matrix packages several output directions at once. I will use the same matrix throughout this post:

Row 1 says: “to build the first output coordinate, take $2$ times the first input plus $1$ times the second.” Row 2 says: “take $0$ times the first input plus $3$ times the second.” Applying $A$ to $x=[4,5]^T$ means computing each output entry separately.

First output coordinate:

Second output coordinate:

So the transformed vector is

Geometrically, the point $(4,5)$ lands at $(13,15)$. I can also read the same calculation in column language: the columns of $A$ are where the basis vectors $\hat{\imath}=[1,0]^T$ and $\hat{\jmath}=[0,1]^T$ get sent. The input $[4,5]^T$ says “take 4 copies of the first column plus 5 copies of the second”:

Neural-network layers usually think in the row view (weighted sums into each neuron); animations usually think in the column view (where basis arrows go). Same multiplication, two pictures.

One complete example: vectors through eigenvalues

I want one thread of numbers I can follow without switching examples. I keep $A=\begin{bmatrix}2&1\0&3\end{bmatrix}$ and $x=\begin{bmatrix}4\5\end{bmatrix}$ from above, then ask what else this same matrix tells me.

Matrix multiply (recap)

Dot product on two directions from the same setup

Take the first column of $A$ as vector $a$ and a separate direction $b$:

The dot product multiplies matching entries and adds:

Length of $a$:

Length of $b$:

The geometric formula $a\cdot b = |a||b|\cos\theta$ gives:

A large positive dot product means the vectors point roughly the same way; zero would mean perpendicular; negative would mean they face opposite directions. That is why dot products score similarity in search, recommendations, embeddings, and transformer attention.

Determinant: how much area scales

Before inverses and eigenvalues, I check whether $A$ squashes space. For a $2\times2$ matrix,

A unit square becomes a parallelogram with six times the area. If $\det(A)=0$, some direction collapses to a line or point and information is lost forever — no inverse can undo that.

Eigenvalues: which directions only get stretched

An eigenvector $v$ is a direction $A$ does not rotate — it only scales by some number $\lambda$:

I find $\lambda$ by solving $\det(A - \lambda I)=0$.

Set the determinant to zero:

So one eigenvalue is

and the other is

Eigenvector for $\lambda_1 = 2$. Solve $(A - 2I)v = 0$:

The first row gives

Pick $v_1 = 1$:

Check:

Eigenvector for $\lambda_2 = 3$. Solve $(A - 3I)v = 0$:

The first row gives

Pick $v_1 = 1$:

Check:

That is the full chain on one matrix: input vector, multiply, dot product for alignment, determinant for area scaling, eigenpairs for preserved directions. When I apply $A$ many times, an eigenvector direction grows or shrinks by $\lambda^k$ — which is why eigenvalues matter for stability, PCA, and worrying about exploding or vanishing signals in deep networks.

How do dot products measure alignment?

The numeric thread above already computed $a\cdot b=6$. The extra idea is projection: how much of one vector lies along another. For $a=[2,1]$ and $b=[3,4]$ (a different pair, same mechanics),

Projection of $b$ onto $a$:

Least-squares regression is the same picture in higher dimensions: when no exact solution exists, project the target onto the space of predictions the model can actually produce.

How does matrix multiplication compose transformations?

If matrix $B$ acts on $x$ first and matrix $A$ acts on the result, the combined effect is $ABx$ — rightmost matrix hits $x$ first. Order matters: rotate then stretch is not generally the same as stretch then rotate, so $AB \neq BA$ in most cases.

Shapes encode the contract. If $B$ is $3\times2$ and $x$ is $2\times1$, then

Matrix $A$ multiplying that result must have three columns. Dimension matching is not bureaucratic bookkeeping; it says the output space of one map must equal the input space of the next.

As a tiny composition with the same $A$, let

(swap the two coordinates). Entry $(AB)_{11}$:

Entry $(AB)_{12}$:

Entry $(AB)_{21}$:

Entry $(AB)_{22}$:

So

Applying to $x=[4,5]^T$ means swap first ($Bx=[5,4]^T$), then apply $A$ — one pipeline, one matrix product.

What do determinants and inverses tell you?

I already found $\det(A)=6$. When $\det(A)\neq 0$, an inverse exists that undoes the transformation. For this $A$,

Solving $Ax=b$ is $x = A^{-1}b$ in principle. In real code I use a numerical solver instead of forming the inverse explicitly — faster and more stable — but the picture is “run the transformation backward.”

What is a change of basis?

A basis is the coordinate language I use to describe vectors. The physical arrow does not move when I change basis; only the numbers describing it change. If the columns of $P$ are new basis vectors expressed in the old coordinates, then

and

Fourier transforms pick a frequency basis; PCA picks variance-aligned axes; token embeddings learn a coordinate system where semantic relationships become geometry later layers can exploit.

Where does machine learning use linear algebra?

Dense layers are $y = Wx + b$: each row of $W$ is a dot product against the input, exactly the row view of matrix multiply. Embedding tables store one vector per token ID. Attention scores query–key alignment with dot products and computes all pairs at once with matrix multiplication. Convolutions apply the same small linear map to every local patch. Regression and PCA lean on projection and eigenvectors of covariance matrices. Backpropagation threads Jacobians — local linear approximations of how small input changes affect outputs — through the graph; that is why I sanity-check gradients with numerical checking before trusting a new backward() implementation.

The AL-1 build path I am following assumes all of this is muscle memory: autograd → gradient checker → BPE → embeddings → positional encoding → attention. Linear algebra is the floor under every step.

Who should watch the 3Blue1Brown series?

If you are heading into machine learning, computer vision, or anything drowning in matrix notation, and symbols still feel like alphabet soup, the playlist is worth the time. You do not need a PhD — just patience and willingness to pause when an animation does not land on the first pass.

I still rewatch individual chapters when a course throws something abstract at me (SVD, change of basis, dual spaces). It is faster than rereading a textbook from zero. The animations gave me a place to put $Ax=b$ and weight matrices in PyTorch; the pencil work — one $2\times2$ example at a time — made the symbols trustworthy. If you are in the same boat I was in, watch the full Essence of Linear Algebra playlist, work through a tiny matrix by hand, and then go build something that actually multiplies matrices. That combination is what finally made it click for me.

This post is one continuous walkthrough: why training needs gradients, the chain rule intuition, how expressions become graphs, a complete hand-derived backprop on $z = (x+y) \cdot y$ with $x=2$, $y=3$ where every gradient step gets its own line, then the ~80-line Python Value class that automates exactly that process. Same pipeline as numerical gradient checking — own the math before you trust the engine. Part of the ML Foundations series; if linear algebra is still fuzzy, start there.

By the end I’ll have a working autograd engine that can train a neural network on XOR.

Why do neural networks need derivatives?

Training a neural network means minimizing a loss function $L$ by nudging parameters in the direction that decreases $L$. That direction is the gradient: a vector of partial derivatives $\frac{\partial L}{\partial w_i}$ for every weight $w_i$.

For a network with millions of parameters, computing these derivatives by hand is impossible. I need a machine to do it.

Three approaches exist:

Automatic differentiation (autograd) is exact (up to floating point), fast, and works for any Python control flow — loops, conditionals, recursion. When I finish the engine below, I’ll verify every backward() with numerical gradient checking before I trust it on anything harder than XOR.

graph LR
  subgraph "Differentiation methods"
    S["Symbolic\ndiff"] -->|"Expression\nexplosion"| X1["✗"]
    N["Numerical\ndiff"] -->|"Slow +\nunstable"| X2["✗"]
    A["Automatic\ndiff"] -->|"Exact +\nfast"| OK["✓"]
  end
  style A fill:#d4edda,stroke:#009900,color:#111
  style OK fill:#d4edda,stroke:#009900,color:#111
  style X1 fill:#f8d7da,stroke:#990000,color:#111
  style X2 fill:#f8d7da,stroke:#990000,color:#111

What is the chain rule, and why is it enough?

Every neural network computation decomposes into elementary operations: add, multiply, exponentiate, ReLU, and so on. The chain rule tells me how to compose derivatives through a sequence of those operations.

If $L = f(g(x))$, then:

Read that right-to-left when doing backprop: start at the output, multiply by each local derivative as I walk toward the inputs.

More generally, for a node $v$ in a computation graph with successors $s_1, s_2, \ldots$:

where $\bar{v} = \frac{\partial L}{\partial v}$ is called the adjoint of $v$.

The sum matters. If one variable feeds two later operations — $y$ shows up in both $x+y$ and a multiply — its gradient receives contributions from both branches. I add them; I never overwrite.

That adjoint equation is the entire mathematical foundation. Everything below is implementation.

flowchart LR
  L["Output L"] -->|"local ∂"| N3["Node 3"]
  N3 -->|"local ∂"| N2["Node 2"]
  N2 -->|"local ∂"| N1["Node 1"]
  N1 -->|"local ∂"| X["Input x"]
  style L fill:#d4edda,stroke:#009900,color:#111

What is a computation graph?

Every expression I write builds a directed acyclic graph (DAG) of operations. Nodes are values; edges are dependencies. Forward pass: walk from leaves to root, computing each node’s value. Backward pass: walk from root to leaves, distributing gradients via the chain rule.

Consider a slightly larger example before the main one:

a = 2.0
b = 3.0
c = a * b   # c = 6
d = c + b   # d = 9
L = d * 2   # L = 18

The graph looks like:

graph LR
  a("a = 2.0") --> mul["× multiply"]
  b("b = 3.0") --> mul
  mul --> c("c = 6.0")
  c --> add["+ add"]
  b --> add
  add --> d("d = 9.0")
  d --> mul2["× 2"]
  mul2 --> L("L = 18.0")

  style L fill:#d4edda,stroke:#009900,color:#111
  style a fill:#fff,stroke:#dddddd,color:#111
  style b fill:#fff,stroke:#dddddd,color:#111

Leaf nodes (a, b) are inputs. The root (L) is the output. Every interior node must remember three things:

1. Its value (forward pass result)
2. Which children (parents in the dataflow sense) created it
3. How to backprop gradients to those children — the local derivative rule for that operation

Notice b appears twice in this graph. When gradients flow backward, b accumulates contributions from both the multiply that produced c and the add that produced d. That accumulation pattern is not a bug — it’s the chain rule doing its job.

How do I backpropagate $z = (x+y) \cdot y$ by hand?

This is the example I wish someone had walked me through before I wrote any code. I will not skip a step.

Setup. Let $x = 2$, $y = 3$. Define an intermediate sum and the final output:

I am treating $z$ as the scalar “loss” at the top of the graph — the thing I want derivatives of with respect to $x$ and $y$.

Forward pass — every value on its own line

Substitute $x = 2$:

Substitute $y = 3$:

Compute the sum:

Compute the product:

Forward pass done. The graph in my head:

graph LR
  x("x = 2") --> add["+ add"]
  y("y = 3") --> add
  add --> s("s = 5")
  s --> mul["× multiply"]
  y --> mul
  mul --> z("z = 15")

  style z fill:#d4edda,stroke:#009900,color:#111
  style x fill:#fff,stroke:#dddddd,color:#111
  style y fill:#fff,stroke:#dddddd,color:#111

y feeds two operations. That detail drives the backward pass.

Backward pass — every gradient on its own line

Reverse mode starts at the output. Seed the top node:

I will write $\bar{v} = \frac{\partial z}{\partial v}$ for short.

Step 1 — multiply node $z = s \cdot y$.

Local partial of $z$ with respect to $s$ (treat $y$ as constant):

Local partial of $z$ with respect to $y$ (treat $s$ as constant):

Chain rule from $z$ into $s$:

Chain rule from $z$ into $y$ through the multiply branch only:

Step 2 — add node $s = x + y$.

Local partial of $s$ with respect to $x$:

Local partial of $s$ with respect to $y$:

Chain rule from $s$ into $x$:

Chain rule from $s$ into $y$ through the add branch:

Step 3 — accumulate $\bar{y}$ from both branches.

$y$ appeared in the add and the multiply. Sum the adjoints:

Final answers

Sanity check with calculus. Expand $z = (x+y)y = xy + y^2$:

That hand walk is exactly what my Value.backward() will automate: seed the root with 1, visit nodes in reverse topological order, apply each operation’s local rule, accumulate with += when a variable appears more than once.

graph LR
  subgraph "Backward order"
    direction RL
    z2["z\n∂z/∂z = 1"] --> mul2["×\n∂z/∂s=3, ∂z/∂y|mul=5"]
    mul2 --> s2["s\n∂s/∂x=1, ∂s/∂y=1"]
    s2 --> x2["x\n∂z/∂x = 3"]
    s2 --> y2["y\n∂z/∂y = 5+3 = 8"]
    mul2 --> y2
  end
  style z2 fill:#d4edda,stroke:#009900,color:#111

How do I build the Value class?

Now I wrap every scalar in a Value object. Each Value stores its data, its gradient, its parent nodes, and a closure _backward that knows how to distribute gradients to parents — encoding the local rules I just derived by hand.

classDiagram
  class Value {
    +float data
    +float grad
    +set _prev
    +string _op
    +_backward()
    +__add__(other)
    +__mul__(other)
    +__pow__(exp)
    +relu()
    +backward()
  }
  note for Value "grad starts at 0.0\n_backward is a closure\nstored at creation time"

class Value:
    def __init__(self, data, _children=(), _op=''):
        self.data = float(data)
        self.grad = 0.0
        self._backward = lambda: None  # no-op by default
        self._prev = set(_children)
        self._op = _op  # for debugging

    def __repr__(self):
        return f"Value(data={self.data:.4f}, grad={self.grad:.4f})"

The grad field starts at zero. It accumulates contributions from all downstream nodes during the backward pass — the $\bar{y} = 5 + 3$ pattern from above.

Reproduce the hand example in code:

x = Value(2.0)
y = Value(3.0)

s = x + y       # s.data == 5
z = s * y       # z.data == 15

z.backward()

print(x.grad)   # 3.0 ✓
print(y.grad)   # 8.0 ✓  (5 from mul + 3 from add)

If either number is wrong, I do not move on. I run numerical gradient checking on that op until it matches to $10^{-5}$.

How does each operation define its local gradient?

For each operation, I do two things:

1. Compute the output value (forward pass)
2. Define a _backward closure that accumulates gradients into the inputs using the chain rule

The closures are the machine-readable version of the hand derivation above.

Addition

Mathematically: if $c = a + b$, then $\frac{\partial c}{\partial a} = 1$ and $\frac{\partial c}{\partial b} = 1$.

So when the backward pass tells me $\frac{\partial L}{\partial c}$ (stored in c.grad), I accumulate:

That is exactly the $\bar{y}_{\text{add}} = \bar{s} \cdot 1$ step from the worked example.

def __add__(self, other):
    other = other if isinstance(other, Value) else Value(other)
    out = Value(self.data + other.data, (self, other), '+')

    def _backward():
        self.grad  += 1.0 * out.grad
        other.grad += 1.0 * out.grad
    out._backward = _backward
    return out

Note the +=. A node can be used multiple times in a computation graph — gradients accumulate, they don’t overwrite.

Multiplication

If $c = a \cdot b$, then $\frac{\partial c}{\partial a} = b$ and $\frac{\partial c}{\partial b} = a$.

That is the $\bar{s} = \bar{z} \cdot y$ and $\bar{y}_{\text{mul}} = \bar{z} \cdot s$ step.

def __mul__(self, other):
    other = other if isinstance(other, Value) else Value(other)
    out = Value(self.data * other.data, (self, other), '*')

    def _backward():
        self.grad  += other.data * out.grad
        other.grad += self.data  * out.grad
    out._backward = _backward
    return out

Power

If $y = x^n$, then $\frac{\partial y}{\partial x} = n \cdot x^{n-1}$.

def __pow__(self, exponent):
    assert isinstance(exponent, (int, float))
    out = Value(self.data ** exponent, (self,), f'**{exponent}')

    def _backward():
        self.grad += exponent * (self.data ** (exponent - 1)) * out.grad
    out._backward = _backward
    return out

ReLU

def relu(self):
    out = Value(max(0, self.data), (self,), 'ReLU')

    def _backward():
        self.grad += (out.data > 0) * out.grad
    out._backward = _backward
    return out

Convenience wrappers

def __neg__(self):         return self * -1
def __sub__(self, other):  return self + (-other)
def __truediv__(self, other): return self * other**-1
def __radd__(self, other): return self + other
def __rmul__(self, other): return self * other
def __rsub__(self, other): return other + (-self)

The r variants handle expressions like 2 * v where the left operand is a Python number.

How does the backward pass visit nodes in the right order?

The hand derivation assumed I already knew $\bar{s}$ before I computed $\bar{x}$ and $\bar{y}_{\text{add}}$. That ordering is not optional.

Right order means: before processing a node, all nodes that depend on it must already have their gradients computed. That is exactly reverse topological order.

graph LR
  subgraph "Forward order (topo)"
    direction LR
    A["a"] --> C["c=a*b"] --> E["L=c+b"]
    B["b"] --> C
    B --> E
  end
  subgraph "Backward order (reversed)"
    direction RL
    E2["L"] --> C2["c"] --> A2["a"]
    E2 --> B2["b"]
    C2 --> B2
  end
  style E fill:#d4edda,stroke:#009900,color:#111
  style E2 fill:#d4edda,stroke:#009900,color:#111

For $z = (x+y) \cdot y$, backward visits: $z \to$ multiply $\to$ $s \to$ add $\to$ $x$, $y$. The multiply’s _backward runs before the add’s _backward can finish propagating into $y$ — but both write to y.grad with +=, so order among siblings does not matter; what matters is that z.grad is set before the multiply closure reads it.

I find the order with a depth-first search, then walk it in reverse:

def backward(self):
    topo = []
    visited = set()

    def build_topo(v):
        if v not in visited:
            visited.add(v)
            for child in v._prev:
                build_topo(child)
            topo.append(v)

    build_topo(self)

    # Seed: dL/dL = 1
    self.grad = 1.0
    for node in reversed(topo):
        node._backward()

The DFS appends a node after recursing into its children — so topo comes out in topological order (leaves first, root last). I reverse it so I start at the root and work toward leaves.

What does the full autograd engine look like?

All the pieces in one place — 82 lines:

class Value:
    def __init__(self, data, _children=(), _op=''):
        self.data = float(data)
        self.grad = 0.0
        self._backward = lambda: None
        self._prev = set(_children)
        self._op = _op

    def __add__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data + other.data, (self, other), '+')
        def _backward():
            self.grad  += out.grad
            other.grad += out.grad
        out._backward = _backward
        return out

    def __mul__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data * other.data, (self, other), '*')
        def _backward():
            self.grad  += other.data * out.grad
            other.grad += self.data  * out.grad
        out._backward = _backward
        return out

    def __pow__(self, exponent):
        out = Value(self.data ** exponent, (self,), f'**{exponent}')
        def _backward():
            self.grad += exponent * (self.data ** (exponent - 1)) * out.grad
        out._backward = _backward
        return out

    def relu(self):
        out = Value(max(0, self.data), (self,), 'ReLU')
        def _backward():
            self.grad += (out.data > 0) * out.grad
        out._backward = _backward
        return out

    def backward(self):
        topo, visited = [], set()
        def build_topo(v):
            if v not in visited:
                visited.add(v)
                for child in v._prev: build_topo(child)
                topo.append(v)
        build_topo(self)
        self.grad = 1.0
        for node in reversed(topo):
            node._backward()

    def __neg__(self):          return self * -1
    def __sub__(self, other):   return self + (-other)
    def __truediv__(self, other): return self * other**-1
    def __radd__(self, other):  return self + other
    def __rmul__(self, other):  return self * other
    def __repr__(self):
        return f"Value(data={self.data:.4f}, grad={self.grad:.4f})"