🧮 Autoregressive Transformer: Encoder + Decoder
This notebook implements the original Transformer (Vaswani et al., 2017, Attention Is All You Need) in pure NumPy — the full encoder–decoder architecture, with every intermediate matrix printed and every formula written out.
The task here is autoregressive next-token prediction:
- the encoder reads a prefix of the sentence (9 tokens);
- the decoder predicts the continuation of that same sentence (6 tokens), one token at a time, each prediction conditioned on the prefix and on everything it has already generated.
This is the "prefix-LM" arrangement: same architecture as translation, different task. Nothing about the encoder, the decoder, or the cross-attention changes — only what the two sequences contain.
The chain is carried all the way to the end:
tokens → IDs → embeddings → positional encoding → encoder (MHA + FFN) → decoder (masked MHA + cross-attention + FFN) → linear projection → logits → softmax → next token
The Autoregressive Factorization
"Autoregressive" means the model factorizes the probability of the whole continuation into a product of next-token probabilities, each conditioned only on what came before:
$$P(y_1, y_2, \dots, y_n \mid x) \;=\; \prod_{t=1}^{n} P(y_t \mid y_1, \dots, y_{t-1},\, x)$$
Two architectural components enforce the "only on what came before" part, and both are implemented below:
- Shift right — at position $t$ the decoder is fed $y_{t-1}$, never $y_t$.
- Causal mask — position $t$ cannot attend to positions $> t$.
The conditioning on $x$ (the prefix) enters through cross-attention, which is the one place the encoder and decoder touch:
$$Q \text{ from the \textbf{decoder}}, \qquad K, V \text{ from the \textbf{encoder}}$$
Honest caveat, stated once: the weights here are random, not trained. Every number is computed correctly, but the predicted token is meaningless. The purpose is the mechanics of the forward pass, not linguistic quality.
NumPy: 2.4.4
Display helpers ready: show_matmul, show_elementwise, show_masked, show_weights_masked, expand.
1. Task & Hyperparameters
The sentence is split into two parts:
| content | length | fed to | |
|---|---|---|---|
| prefix $x$ | the cat sat on the mat and was looking |
9 | encoder |
| continuation $y$ | at a tiny mouse running by |
6 | decoder (autoregressively) |
The decoder must produce at a tiny mouse running by <eos>, one token at a time.
The two lengths are deliberately different (9 vs 6). This matters: it forces the cross-attention matrix to be rectangular ($n_{tgt} \times n_{src}$), which is easy to get wrong when both lengths are equal, and it makes the shapes unambiguous when reading the printouts.
Note that both halves share one vocabulary — they are two parts of the same English sentence. In the translation version of this architecture the two halves would be two different languages; the machinery is identical either way.
| Symbol | Meaning | Value |
|---|---|---|
| $n_{src}$ | prefix length | 9 |
| $n_{tgt}$ | decoder input length | 7 (<sos> + 6 words) |
| $d_{model}$ | embedding / model width | 8 |
| $h$ | number of attention heads | 2 |
| $d_k = d_v = d_{model}/h$ | width of one head | 4 |
| $d_{ff}$ | inner width of the FFN | 32 ($=4 d_{model}$) |
| $V$ | vocabulary size | 500 |
Full sentence (15 tokens): ['the', 'cat', 'sat', 'on', 'the', 'mat', 'and', 'was', 'looking', 'at', 'a', 'tiny', 'mouse', 'running', 'by'] PREFIX -> encoder (9 tokens): ['the', 'cat', 'sat', 'on', 'the', 'mat', 'and', 'was', 'looking'] CONTINUATION -> decoder (6 tokens): ['at', 'a', 'tiny', 'mouse', 'running', 'by'] d_model=8, h=2, d_k=4, d_ff=32, V=500
2. Vocabulary & Token IDs
As requested, token IDs are random integers in [50, 500) — this mimics a real tokenizer, where the ID is an arbitrary index into a vocabulary, carrying no meaning by itself.
Three IDs are reserved below 50, as is standard:
<pad>= 0 — padding (not used here, sequences are not batched)<sos>= 1 — start of sequence, fed to the decoder at position 0<eos>= 2 — end of sequence, what the decoder must predict last
Prefix and continuation share one vocabulary of size $V = 500$ — they are halves of the same sentence. Note that the repeated word "the" receives the same ID at both of its positions — an embedding table is a lookup by ID, so identical words start from identical vectors. What makes them differ later is positional encoding and context.
Vocabulary: 17 entries used out of V=500 PREFIX token -> ID (encoder) the -> 83 cat -> 168 sat -> 458 on -> 328 the -> 83 mat -> 421 and -> 185 was -> 276 looking -> 274 CONTINUATION token -> ID (decoder) at -> 127 a -> 340 tiny -> 68 mouse -> 57 running -> 412 by -> 322 Note: 'the' appears at positions 0 and 4, both with ID 83.
3. Weight Initialization
The old notebook drew every weight from np.random.rand, i.e. $\mathcal{U}(0,1)$: all entries positive, mean 0.5.
For $y = xW$ with $W \in \mathbb{R}^{8\times 8}$ and all entries positive, each output is a sum of 8 positive products, so $\mathbb{E}[y] \approx 8 \cdot 0.5 \cdot \mathbb{E}[x] = 4\,\mathbb{E}[x]$. The signal is multiplied by ~4 at every matrix product, and it is never re-centred because there is no LayerNorm. Chaining ≈ 8 such products gives $4^8 \approx 65\,000$ — exactly the order of magnitude seen in the old output (≈ 80 000). Worse, since all rows grow in the same direction, they become nearly parallel and the row-to-row differences vanish: the tokens lose their identity.
Xavier/Glorot initialization fixes the scale:
$$W_{ij} \sim \mathcal{U}(-a, a), \qquad a = \sqrt{\frac{6}{fan_{in} + fan_{out}}}$$
which keeps $\mathrm{Var}(y) \approx \mathrm{Var}(x)$ and $\mathbb{E}[y] = 0$.
After 8 chained matmuls, no LayerNorm: U(0,1) init -> mean |value| = 33138.98 row-to-row std = 9033.2708 Xavier init -> mean |value| = 0.2531 row-to-row std = 0.1784 The U(0,1) rows have exploded AND collapsed onto each other.
4. Embeddings
The embedding table is $E \in \mathbb{R}^{V \times d_{model}} = \mathbb{R}^{500 \times 8}$. Looking up token ID $t$ means taking row $t$: this is exactly a one-hot vector times $E$,
$$\mathrm{emb}(t) = \mathbf{e}_t^\top E, \qquad \mathbf{e}_t \in \{0,1\}^{V}$$
but implemented as indexing, since the matrix product would be wasteful.
Scaling by $\sqrt{d_{model}}$. The paper multiplies the embeddings by $\sqrt{d_{model}}$ before adding the positional encoding (Section 3.4). The reason is scale matching: positional encodings are sines and cosines with values in $[-1, 1]$, while embeddings initialized with variance $1/d_{model}$ have entries around $\pm 0.35$. Without the scaling, the positional signal would dominate the semantic signal. Here $\sqrt{8} \approx 2.83$.
Embedding table E shape: (500, 8) src_emb = E[src_ids] * sqrt(8) = E[src_ids] * 2.8284 Prefix embeddings shape: (9, 8) the (id 83) -> [-0.1 -0.204 -0.878 -0.827 -0.226 0.367 0.914 -0.803] cat (id 168) -> [-0.685 0.791 -1.97 0.893 -1.211 0.731 0.014 -0.954] sat (id 458) -> [ 1.393 0.839 0.873 -0.966 -0.079 -0.3 -1.006 0.36 ] on (id 328) -> [ 1.221 -0.877 1.712 -1.748 0.435 0.475 -0.796 0.424] the (id 83) -> [-0.1 -0.204 -0.878 -0.827 -0.226 0.367 0.914 -0.803] mat (id 421) -> [ 0.234 1.035 -0.836 1.424 -0.092 -0.888 0.998 -0.287] and (id 185) -> [-0.536 -1.707 -1.117 1.236 -0.156 -0.548 0.16 0.502] was (id 276) -> [-0.864 1.279 1.042 0.584 -0.13 0.58 -0.707 0.856] looking (id 274) -> [ 0.234 -0.249 1.201 0.14 -1.967 -1.117 -0.186 0.31 ] Check that both occurrences of 'the' are identical before PE: position 0: [-0.1 -0.204 -0.878 -0.827 -0.226 0.367 0.914 -0.803] position 4: [-0.1 -0.204 -0.878 -0.827 -0.226 0.367 0.914 -0.803] identical? True
5. Positional Encoding
Self-attention is a set operation: it computes a weighted sum over positions, and the weights depend only on content. Nothing in $QK^\top$ knows that position 3 comes after position 2. Order must therefore be injected into the input itself.
The sinusoidal encoding from the paper, for position $pos$ and dimension index $i$:
$$PE_{(pos,\,2i)} = \sin\!\left(\frac{pos}{10000^{2i/d_{model}}}\right), \qquad PE_{(pos,\,2i+1)} = \cos\!\left(\frac{pos}{10000^{2i/d_{model}}}\right)$$
With $d_{model} = 8$ there are 4 frequency pairs, $i = 0, 1, 2, 3$:
| $i$ | columns | wavelength $2\pi \cdot 10000^{2i/8}$ |
|---|---|---|
| 0 | 0, 1 | $2\pi \approx 6.28$ |
| 1 | 2, 3 | $2\pi \cdot 10 \approx 62.8$ |
| 2 | 4, 5 | $2\pi \cdot 100 \approx 628$ |
| 3 | 6, 7 | $2\pi \cdot 1000 \approx 6283$ |
Fast columns encode fine-grained local position; slow columns encode coarse absolute position. The property that makes this choice useful is that $PE_{pos+k}$ is a fixed linear function of $PE_{pos}$ (a 2-D rotation by angle $k\omega_i$ in each frequency pair), so relative offsets are linearly recoverable by the attention weights.
PE_src shape: (9, 8) [[ 0. 1. 0. 1. 0. 1. 0. 1. ] [ 0.842 0.54 0.1 0.995 0.01 1. 0.001 1. ] [ 0.909 -0.416 0.199 0.98 0.02 1. 0.002 1. ] [ 0.141 -0.99 0.295 0.955 0.03 1. 0.003 1. ] [-0.757 -0.654 0.389 0.921 0.04 0.999 0.004 1. ] [-0.959 0.284 0.479 0.878 0.05 0.999 0.005 1. ] [-0.279 0.96 0.565 0.825 0.06 0.998 0.006 1. ] [ 0.657 0.754 0.644 0.765 0.07 0.998 0.007 1. ] [ 0.989 -0.145 0.717 0.697 0.08 0.997 0.008 1. ]] Hand check PE[3,2] = sin(3 / 10000^(2*1/8)) = sin(3 * 0.100000) = 0.295520 From array PE[3,2] = 0.295520 -> match: True
Encoder Input
$$X_{enc} = \mathrm{Embedding}(\text{src\_ids}) \cdot \sqrt{d_{model}} + PE$$
This is the step that was missing in the old notebook. Its effect is visible immediately: the two occurrences of "the" (positions 0 and 4 of the prefix), which had identical embeddings, now have different input vectors.
X_enc shape: (9, 8) 0 the [-0.1 0.796 -0.878 0.173 -0.226 1.367 0.914 0.197] 1 cat [ 0.157 1.331 -1.87 1.888 -1.201 1.731 0.015 0.046] 2 sat [ 2.302 0.423 1.071 0.014 -0.059 0.699 -1.004 1.36 ] 3 on [ 1.362 -1.867 2.008 -0.792 0.465 1.475 -0.793 1.424] 4 the [-0.857 -0.857 -0.489 0.094 -0.186 1.367 0.918 0.197] 5 mat [-0.725 1.319 -0.357 2.302 -0.042 0.111 1.003 0.713] 6 and [-0.815 -0.747 -0.552 2.061 -0.096 0.45 0.166 1.502] 7 was [-0.207 2.033 1.686 1.348 -0.06 1.577 -0.7 1.856] 8 looking [ 1.223 -0.394 1.918 0.837 -1.887 -0.12 -0.178 1.31 ] 'the' at position 0: [-0.1 0.796 -0.878 0.173 -0.226 1.367 0.914 0.197] 'the' at position 4: [-0.857 -0.857 -0.489 0.094 -0.186 1.367 0.918 0.197] identical now? False <- order information is now present
Embeddings + Positional Encoding
$X_{enc} = \text{scaled embeddings} + PE$ is element-wise: three matrices of shape $(9, 8)$, nothing mixes across rows or columns. Both operands are printed side by side below.
==============================================================================
X_enc = emb*sqrt(d) + PE all three are (9, 8) — element-wise, no mixing
==============================================================================
the |
emb*sqrt(d): -0.100 -0.204 -0.878 -0.827 -0.226 +0.367 +0.914 -0.803
PE: +0.000 +1.000 +0.000 +1.000 +0.000 +1.000 +0.000 +1.000
X_enc: -0.100 +0.796 -0.878 +0.173 -0.226 +1.367 +0.914 +0.197
cat |
emb*sqrt(d): -0.685 +0.791 -1.970 +0.893 -1.211 +0.731 +0.014 -0.954
PE: +0.841 +0.540 +0.100 +0.995 +0.010 +1.000 +0.001 +1.000
X_enc: +0.157 +1.331 -1.870 +1.888 -1.201 +1.731 +0.015 +0.046
sat |
emb*sqrt(d): +1.393 +0.839 +0.873 -0.966 -0.079 -0.300 -1.006 +0.360
PE: +0.909 -0.416 +0.199 +0.980 +0.020 +1.000 +0.002 +1.000
X_enc: +2.302 +0.423 +1.071 +0.014 -0.059 +0.699 -1.004 +1.360
on |
emb*sqrt(d): +1.221 -0.877 +1.712 -1.748 +0.435 +0.475 -0.796 +0.424
PE: +0.141 -0.990 +0.296 +0.955 +0.030 +1.000 +0.003 +1.000
X_enc: +1.362 -1.867 +2.008 -0.792 +0.465 +1.475 -0.793 +1.424
the |
emb*sqrt(d): -0.100 -0.204 -0.878 -0.827 -0.226 +0.367 +0.914 -0.803
PE: -0.757 -0.654 +0.389 +0.921 +0.040 +0.999 +0.004 +1.000
X_enc: -0.857 -0.857 -0.489 +0.094 -0.186 +1.367 +0.918 +0.197
mat |
emb*sqrt(d): +0.234 +1.035 -0.836 +1.424 -0.092 -0.888 +0.998 -0.287
PE: -0.959 +0.284 +0.479 +0.878 +0.050 +0.999 +0.005 +1.000
X_enc: -0.725 +1.319 -0.357 +2.302 -0.042 +0.111 +1.003 +0.713
and |
emb*sqrt(d): -0.536 -1.707 -1.117 +1.236 -0.156 -0.548 +0.160 +0.502
PE: -0.279 +0.960 +0.565 +0.825 +0.060 +0.998 +0.006 +1.000
X_enc: -0.815 -0.747 -0.552 +2.061 -0.096 +0.450 +0.166 +1.502
was |
emb*sqrt(d): -0.864 +1.279 +1.042 +0.584 -0.130 +0.580 -0.707 +0.856
PE: +0.657 +0.754 +0.644 +0.765 +0.070 +0.998 +0.007 +1.000
X_enc: -0.207 +2.033 +1.686 +1.348 -0.060 +1.577 -0.700 +1.856
looking |
emb*sqrt(d): +0.234 -0.249 +1.201 +0.140 -1.967 -1.117 -0.186 +0.310
PE: +0.989 -0.146 +0.717 +0.697 +0.080 +0.997 +0.008 +1.000
X_enc: +1.223 -0.394 +1.918 +0.837 -1.887 -0.120 -0.178 +1.310
check element [0,0]: -0.100154 + +0.000000 = -0.100154 MATCH: True
6. Layer Normalization
This is the component whose absence caused the 80 000-magnitude output in the old notebook. For each row (each token) independently:
$$\mu = \frac{1}{d}\sum_{j=1}^{d} x_j, \qquad \sigma^2 = \frac{1}{d}\sum_{j=1}^{d}(x_j - \mu)^2, \qquad \mathrm{LN}(x) = \gamma \odot \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta$$
Three points worth being precise about:
- Per row, not per column. Unlike BatchNorm, no statistic is shared across tokens. Token $i$ is normalized using only its own 8 numbers. This is what makes it valid at inference time with a single sequence.
- $\gamma, \beta \in \mathbb{R}^{d_{model}}$ are learned; initialized to $\mathbf{1}$ and $\mathbf{0}$, so LayerNorm starts as pure standardization.
- $\epsilon = 10^{-5}$ guards against division by zero.
The original paper uses post-norm: LayerNorm(x + Sublayer(x)). That is what is implemented below. (Modern implementations often use pre-norm, x + Sublayer(LayerNorm(x)), which trains more stably in deep stacks — but post-norm is the 2017 architecture and this notebook follows it.)
x = [-0.1 0.796 -0.878 0.173 -0.226 1.367 0.914 0.197]
mean = 0.2803
variance = 0.4513 (std = 0.6718)
LN(x) = [-0.566 0.768 -1.724 -0.16 -0.754 1.618 0.943 -0.124]
mean = -8.33e-17, std = 1.0000
7. Scaled Dot-Product Attention
$$\mathrm{Attention}(Q, K, V) = \mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}} + M\right) V$$
Reading it as a sequence of shapes, with $n_q$ queries and $n_{kv}$ keys/values:
| Step | Operation | Shape |
|---|---|---|
| 1 | $Q$, $K$, $V$ | $(n_q, d_k)$, $(n_{kv}, d_k)$, $(n_{kv}, d_v)$ |
| 2 | $S = QK^\top$ | $(n_q, n_{kv})$ — one score per (query, key) pair |
| 3 | $S / \sqrt{d_k}$ | same |
| 4 | $+ M$ (mask, optional) | same |
| 5 | $A = \mathrm{softmax}(\cdot)$ row-wise | same, each row sums to 1 |
| 6 | $AV$ | $(n_q, d_v)$ |
Why divide by $\sqrt{d_k}$? If the components of $q$ and $k$ are independent with mean 0 and variance 1, then $q \cdot k = \sum_{j=1}^{d_k} q_j k_j$ has mean 0 and variance $d_k$, i.e. standard deviation $\sqrt{d_k}$. Feeding scores of that magnitude into a softmax pushes it toward a one-hot distribution, where the gradient is nearly zero. Dividing by $\sqrt{d_k}$ restores unit variance and keeps the softmax in its responsive region. Here $\sqrt{d_k} = \sqrt{4} = 2$.
Softmax, applied to each row $s$ independently:
$$A_{ij} = \frac{\exp(s_{ij})}{\sum_{j'} \exp(s_{ij'})}$$
implemented as exp(s - max(s)) for numerical stability — subtracting the row max changes nothing mathematically (the constant cancels) but prevents exp from overflowing.
Numerical check on a tiny example: raw scores q.k = [[2. 0. 1.]] (2, 0, 1) scaled by sqrt(4) = [[1. 0. 0.5]] softmax weights = [[0.506 0.186 0.307]] sum = 1.0000 output = w @ V = [[5.065 1.863 3.072 0. ]] -> the query matches key 0 most, so V row 0 dominates the output.
8. Multi-Head Attention
The old notebook described "multi-head" but computed a single head and omitted the output projection $W^O$. Both are corrected here.
$$\mathrm{MultiHead}(Q,K,V) = \mathrm{Concat}(\mathrm{head}_1, \dots, \mathrm{head}_h)\,W^O, \qquad \mathrm{head}_i = \mathrm{Attention}(QW_i^Q,\, KW_i^K,\, VW_i^V)$$
Why several heads? A single softmax row is a probability distribution — it can only express one weighting of the context. Averaging cancels competing patterns. With $h$ heads, each head runs the same operation in its own $d_k$-dimensional subspace, so one head can track syntactic agreement while another tracks a co-reference, and $W^O$ recombines them.
Implementation. Rather than $h$ separate small matrices, one uses full $d_{model} \times d_{model}$ matrices $W^Q, W^K, W^V$ and reshapes the result — mathematically identical, since a block of columns of $W^Q$ is exactly $W_i^Q$:
X (n, 8) @ Wq (8, 8) -> Q (n, 8)
reshape -> (n, 2, 4)
transpose -> (2, n, 4) # (heads, tokens, d_k)
Total parameter count is unchanged: $h \cdot d_k = d_{model}$. Multi-head attention costs the same as single-head, but is strictly more expressive.
MHA class defined. Each instance owns Wq, Wk, Wv, Wo of shape (8, 8). Parameters per attention sub-layer: 4 * 8 * 8 = 256
9. Encoder: Self-Attention
"Self"-attention means $Q$, $K$ and $V$ all come from the same input $X_{enc}$. No causal mask is used here. This is the point that most often causes confusion in an autoregressive setup: the prefix is fully known input, not something being generated, so the encoder is free to be bidirectional — token 2 may look at token 7. The autoregressive constraint applies only to the decoder side, where tokens are produced one at a time.
Below the computation is unrolled, and one score is verified by hand against the definition.
[encoder self-attn] x_q (9, 8) , x_kv (9, 8) Q,K,V per head: (2, 9, 4), (2, 9, 4), (2, 9, 4) scores/weights: (2, 9, 9) concat: (9, 8) out: (9, 8) --- Head 0: Q, K, V (9 x 4 each) --- Q[head0]: [[-0.363 0.43 1.132 -1.267] [-1.14 -1.112 0.583 -0.527] [-0.04 -0.869 0.345 -0.223] [ 1.895 1.278 0.877 0.181] [ 0.557 1.335 1.018 -0.361] [-1.574 0.05 -0.869 -0.623] [-0.513 0.355 -0.385 0.953] [ 0.287 0.698 -0.771 -1.189] [-0.319 -0.98 -0.91 1.053]] K[head0]: [[-0.642 0.947 -0.654 -0.265] [-0.815 1.333 -0.331 0.315] [-0.212 0.556 1.075 0.39 ] [ 0.833 -0.997 1.658 -0.011] [ 0.126 -0.482 0.012 -0.353] [ 0.888 -0.151 0.471 -0.344] [ 1.368 -1.584 1.217 -0.221] [ 1.429 -0.377 0.183 0.313] [ 0.161 -1.964 2.485 1.113]] V[head0]: [[-1.375 -1.282 0.058 1.186] [-2.807 -1.31 -0.778 1.31 ] [ 0.743 0.139 -0.711 -1.884] [ 1.641 -0.434 0.404 -0.671] [-1.132 -1.533 0.746 2.018] [-3.082 -0.455 -1.491 0.839] [-2.619 -1.071 -0.747 1.683] [-2.389 0.491 -2.011 -0.774] [-1.053 -0.132 0.331 -1.287]] --- Head 0: raw scores Q K^T / sqrt(d_k), shape (9, 9) --- [[ 0.117 0.047 0.519 0.579 0.104 0.291 0.24 -0.436 0.25 ] [-0.282 -0.457 0.023 0.566 0.292 -0.194 0.514 -0.634 1.432] [-0.482 -0.656 -0.095 0.704 0.248 0.167 0.896 0.132 1.156] [-0.314 -0.037 0.661 0.878 -0.214 0.921 0.798 1.222 0.087] [ 0.168 0.437 0.788 0.412 -0.216 0.449 -0.016 0.183 -0.203] [ 0.895 0.72 -0.408 -1.398 -0.007 -0.8 -1.577 -1.311 -1.602] [ 0.332 0.66 0.133 -0.715 -0.288 -0.509 -0.972 -0.319 -0.338] [ 0.648 0.289 -0.483 -0.862 0.055 0.098 -0.695 -0.183 -2.283] [-0.204 -0.207 -0.522 -0.404 0.025 -0.463 -0.112 0.039 0.393]]
q_1 (head 0, token 'cat') = [-1.14 -1.112 0.583 -0.527]
k_8 (head 0, token 'looking') = [ 0.161 -1.964 2.485 1.113]
q . k = (-1.1400)(0.1606) + (-1.1117)(-1.9641) + (0.5834)(2.4849) + (-0.5266)(1.1127)
= 2.864096
/ sqrt(d_k) = 2.864096 / 2.0 = 1.432048
From the scores matrix: 1.432048 match: True
Softmax of row 1 (head 0), by hand:
[0.059 0.05 0.08 0.138 0.105 0.065 0.131 0.042 0.329]
sum = 1.000000
matches library row: True
Every row sums to 1: True Note the full square (no mask): the encoder is bidirectional over the prefix. The pattern is diffuse because the weights are random, not trained.
Self-Attention, Step by Step
This is the shape question spelled out. Follow the inner dimensions as they cancel:
| # | operation | shapes | result |
|---|---|---|---|
| 1 | $X_{enc} W^Q$ | $(9,8) @ (8,8)$ | $(9,8)$ |
| 2 | split into heads | reshape | $(2,9,4)$ |
| 3 | $Q_0 K_0^\top$ | $(9,4) @ (4,9)$ | $(9,\mathbf{9})$ |
| 4 | softmax | element-wise per row | $(9,9)$ |
| 5 | $A V_0$ | $(9,9) @ (9,4)$ | $(9,\mathbf{4})$ |
| 6 | concat heads | $(9,4) + (9,4)$ side by side | $(9,8)$ |
| 7 | $\cdot\, W^O$ | $(9,8) @ (8,8)$ | $(9,8)$ |
Why step 5 returns $(9,4)$ and not $(9,9)$: the $9 \times 9$ matrix is weights, not content. In $AV$ the shared dimension is the 9 keys, and it cancels — $(9,9)@(9,4) \to (9,4)$. Each output row is a weighted average of the 9 rows of $V$, so it has $V$'s width, which is $d_v = 4$. The count of tokens sets the number of rows, never the number of columns.
The block therefore returns $(9,8)$ — exactly the shape it received. That invariance is what allows 6 blocks to be stacked.
STEP 1 : project X_enc to Q (K and V are identical in form)
==============================================================================
==============================================================================
X_enc (9, 8) @ Wq (8, 8) = Q_full (9, 8)
shape rule: (9 x 8) @ (8 x 8) -> (9 x 8) inner dim 8 cancels
K uses Wk, V uses Wv — same shapes, same arithmetic
==============================================================================
X_enc =
the | -0.100 +0.796 -0.878 +0.173 -0.226 +1.367 +0.914 +0.197
cat | +0.157 +1.331 -1.870 +1.888 -1.201 +1.731 +0.015 +0.046
sat | +2.302 +0.423 +1.071 +0.014 -0.059 +0.699 -1.004 +1.360
on | +1.362 -1.867 +2.008 -0.792 +0.465 +1.475 -0.793 +1.424
the | -0.857 -0.857 -0.489 +0.094 -0.186 +1.367 +0.918 +0.197
mat | -0.725 +1.319 -0.357 +2.302 -0.042 +0.111 +1.003 +0.713
and | -0.815 -0.747 -0.552 +2.061 -0.096 +0.450 +0.166 +1.502
was | -0.207 +2.033 +1.686 +1.348 -0.060 +1.577 -0.700 +1.856
looking | +1.223 -0.394 +1.918 +0.837 -1.887 -0.120 -0.178 +1.310
Wq =
-0.356 -0.524 +0.249 -0.119 +0.514 +0.184 -0.502 +0.549
-0.296 -0.199 -0.120 -0.550 -0.004 -0.461 -0.121 -0.183
+0.334 +0.371 -0.424 -0.163 -0.155 -0.427 +0.366 +0.254
-0.379 -0.163 -0.483 +0.157 -0.464 +0.595 +0.344 -0.134
+0.096 +0.504 +0.066 -0.399 -0.149 +0.420 +0.177 +0.185
+0.449 +0.483 +0.514 -0.496 +0.054 +0.295 +0.315 +0.131
-0.408 +0.364 +0.270 -0.507 -0.044 -0.399 -0.199 -0.590
-0.122 +0.061 +0.146 +0.200 +0.412 -0.348 -0.465 +0.431
Q_full =
the | -0.363 +0.430 +1.132 -1.267 +0.149 -0.034 -0.191 -0.764
cat | -1.140 -1.112 +0.583 -0.527 -0.220 +1.320 +0.033 -0.870
sat | -0.040 -0.869 +0.345 -0.223 +1.661 -0.113 -1.033 +2.714
on | +1.895 +1.278 +0.877 +0.181 +1.397 +0.233 +0.047 +3.064
the | +0.557 +1.335 +1.018 -0.361 -0.263 +0.392 +0.511 -0.763
mat | -1.574 +0.050 -0.869 -0.623 -1.129 +0.147 +0.361 -1.317
and | -0.513 +0.355 -0.385 +0.953 -0.637 +1.160 +0.400 -0.138
was | +0.287 +0.698 -0.771 -1.189 -0.113 -0.821 +0.700 +1.169
looking | -0.319 -0.980 -0.910 +1.053 +0.767 -1.128 -0.522 +1.422
worked element Q_full[0,0] = row 0 of X_enc . column 0 of Wq
row 0 of X_enc = [-0.1 0.796 -0.878 0.173 -0.227 1.367 0.914 0.197]
col 0 of Wq = [-0.356 -0.296 0.334 -0.379 0.096 0.449 -0.408 -0.122]
(-0.1002)*(-0.3560) + (+0.7963)*(-0.2957) + (-0.8780)*(+0.3342) + (+0.1731)*(-0.3792) +
(-0.2265)*(+0.0962) + (+1.3674)*(+0.4493) + (+0.9136)*(-0.4081) + (+0.1968)*(-0.1221)
= -0.363201 | from array: -0.363201 | MATCH: True
K_full = X_enc @ Wk -> (9, 8)
V_full = X_enc @ Wv -> (9, 8)
STEP 2 : split the 8 columns into 2 heads of 4 ============================================================================== Q_full (9, 8) -> reshape(9, 2, 4) -> transpose -> (2, 9, 4) head 0 takes columns 0..3, head 1 takes columns 4..7 Q head 0 (9, 4): [[-0.363 0.43 1.132 -1.267] [-1.14 -1.112 0.583 -0.527] [-0.04 -0.869 0.345 -0.223] [ 1.895 1.278 0.877 0.181] [ 0.557 1.335 1.018 -0.361] [-1.574 0.05 -0.869 -0.623] [-0.513 0.355 -0.385 0.953] [ 0.287 0.698 -0.771 -1.189] [-0.319 -0.98 -0.91 1.053]] Q head 1 (9, 4): [[ 0.149 -0.034 -0.191 -0.764] [-0.22 1.32 0.033 -0.87 ] [ 1.661 -0.113 -1.033 2.714] [ 1.397 0.233 0.047 3.064] [-0.263 0.392 0.511 -0.763] [-1.129 0.147 0.361 -1.317] [-0.637 1.16 0.4 -0.138] [-0.113 -0.821 0.7 1.169] [ 0.767 -1.128 -0.522 1.422]] matches the column slices of Q_full: True
STEP 3 : scores = Q0 @ K0^T, THEN divide by sqrt(d_k)
==============================================================================
==============================================================================
Q_head0 (9, 4) @ K_head0^T (4, 9) = raw_scores (9, 9)
shape rule: (9 x 4) @ (4 x 9) -> (9 x 9) inner dim 4 cancels
K is TRANSPOSED: (9,4) becomes (4,9) so the inner dim 4 can cancel
==============================================================================
Q_head0 =
the | -0.363 +0.430 +1.132 -1.267
cat | -1.140 -1.112 +0.583 -0.527
sat | -0.040 -0.869 +0.345 -0.223
on | +1.895 +1.278 +0.877 +0.181
the | +0.557 +1.335 +1.018 -0.361
mat | -1.574 +0.050 -0.869 -0.623
and | -0.513 +0.355 -0.385 +0.953
was | +0.287 +0.698 -0.771 -1.189
looking | -0.319 -0.980 -0.910 +1.053
K_head0^T =
-0.642 -0.815 -0.212 +0.833 +0.126 +0.888 +1.368 +1.429 +0.161
+0.947 +1.333 +0.556 -0.997 -0.482 -0.151 -1.584 -0.377 -1.964
-0.654 -0.331 +1.075 +1.658 +0.012 +0.471 +1.217 +0.183 +2.485
-0.265 +0.315 +0.390 -0.011 -0.353 -0.344 -0.221 +0.313 +1.113
raw_scores =
the | +0.235 +0.094 +1.038 +1.159 +0.208 +0.582 +0.480 -0.871 +0.500
cat | -0.563 -0.913 +0.046 +1.132 +0.584 -0.388 +1.027 -1.269 +2.864
sat | -0.964 -1.311 -0.190 +1.408 +0.496 +0.335 +1.791 +0.264 +2.311
on | -0.628 -0.073 +1.321 +1.757 -0.429 +1.841 +1.597 +2.444 +0.174
the | +0.336 +0.874 +1.576 +0.824 -0.432 +0.897 -0.033 +0.366 -0.405
mat | +1.791 +1.441 -0.816 -2.796 -0.014 -1.601 -3.154 -2.623 -3.204
and | +0.665 +1.320 +0.265 -1.430 -0.577 -1.019 -1.944 -0.639 -0.675
was | +1.297 +0.578 -0.966 -1.724 +0.110 +0.195 -1.390 -0.367 -4.566
looking | -0.407 -0.414 -1.043 -0.807 +0.049 -0.926 -0.223 +0.078 +0.785
worked element raw_scores[0,0] = row 0 of Q_head0 . column 0 of K_head0^T
row 0 of Q_head0 = [-0.363 0.43 1.132 -1.267]
col 0 of K_head0^T = [-0.642 0.947 -0.654 -0.265]
(-0.3632)*(-0.6416) + (+0.4299)*(+0.9465) + (+1.1320)*(-0.6542) + (-1.2669)*(-0.2648)
= +0.234868 | from array: +0.234868 | MATCH: True
then divide every entry by sqrt(d_k) = sqrt(4) = 2:
[[ 0.117 0.047 0.519 0.579 0.104 0.291 0.24 -0.436 0.25 ]
[-0.282 -0.457 0.023 0.566 0.292 -0.194 0.514 -0.634 1.432]
[-0.482 -0.656 -0.095 0.704 0.248 0.167 0.896 0.132 1.156]
[-0.314 -0.037 0.661 0.878 -0.214 0.921 0.798 1.222 0.087]
[ 0.168 0.437 0.788 0.412 -0.216 0.449 -0.016 0.183 -0.203]
[ 0.895 0.72 -0.408 -1.398 -0.007 -0.8 -1.577 -1.311 -1.602]
[ 0.332 0.66 0.133 -0.715 -0.288 -0.509 -0.972 -0.319 -0.338]
[ 0.648 0.289 -0.483 -0.862 0.055 0.098 -0.695 -0.183 -2.283]
[-0.204 -0.207 -0.522 -0.404 0.025 -0.463 -0.112 0.039 0.393]]
matches enc_scores[0]: True
>>> this (9, 9) matrix is WEIGHTS, not output. 9 queries x 9 keys.
STEP 5 : the multiplication that answers the shape question
==============================================================================
==============================================================================
A (weights) (9, 9) @ V_head0 (9, 4) = head0_out (9, 4)
shape rule: (9 x 9) @ (9 x 4) -> (9 x 4) inner dim 9 cancels
the 9 KEYS are the shared dimension and cancel: (9,9)@(9,4) -> (9,4)
==============================================================================
A (weights) =
the | +0.100 +0.093 +0.149 +0.158 +0.098 +0.119 +0.113 +0.057 +0.114
cat | +0.059 +0.050 +0.080 +0.138 +0.105 +0.065 +0.131 +0.042 +0.329
sat | +0.046 +0.039 +0.068 +0.152 +0.096 +0.089 +0.184 +0.086 +0.239
on | +0.045 +0.060 +0.121 +0.150 +0.050 +0.156 +0.138 +0.211 +0.068
the | +0.100 +0.131 +0.186 +0.128 +0.068 +0.133 +0.083 +0.102 +0.069
mat | +0.325 +0.273 +0.088 +0.033 +0.132 +0.060 +0.027 +0.036 +0.027
and | +0.172 +0.238 +0.140 +0.060 +0.092 +0.074 +0.047 +0.089 +0.088
was | +0.243 +0.169 +0.078 +0.054 +0.134 +0.140 +0.063 +0.106 +0.013
looking | +0.103 +0.102 +0.075 +0.084 +0.129 +0.079 +0.112 +0.131 +0.186
V_head0 =
-1.375 -1.282 +0.058 +1.186
-2.807 -1.310 -0.778 +1.310
+0.743 +0.139 -0.711 -1.884
+1.641 -0.434 +0.404 -0.671
-1.132 -1.533 +0.746 +2.018
-3.082 -0.455 -1.491 +0.839
-2.619 -1.071 -0.747 +1.683
-2.389 +0.491 -2.011 -0.774
-1.053 -0.132 +0.331 -1.287
head0_out =
the | -1.056 -0.609 -0.373 +0.150
cat | -1.043 -0.545 -0.128 -0.077
sat | -1.195 -0.542 -0.307 +0.081
on | -1.373 -0.387 -0.771 +0.019
the | -1.177 -0.543 -0.567 +0.075
mat | -1.611 -1.020 -0.317 +0.854
and | -1.461 -0.731 -0.478 +0.355
was | -1.678 -0.833 -0.516 +0.723
looking | -1.426 -0.606 -0.399 +0.233
worked element head0_out[0,0] = row 0 of A (weights) . column 0 of V_head0
row 0 of A (weights) = [0.1 0.093 0.149 0.158 0.098 0.118 0.113 0.057 0.114]
col 0 of V_head0 = [-1.375 -2.807 0.743 1.641 -1.132 -3.082 -2.619 -2.389 -1.053]
(+0.0996)*(-1.3749) + (+0.0929)*(-2.8072) + (+0.1489)*(+0.7433) + (+0.1581)*(+1.6414) +
(+0.0983)*(-1.1323) + (+0.1185)*(-3.0817) + (+0.1126)*(-2.6194) + (+0.0573)*(-2.3888) +
(+0.1138)*(-1.0528)
= -1.055683 | from array: -1.055683 | MATCH: True
>>> (9, 9) @ (9, 4) = (9, 4). The 9x9 collapsed back to width 4.
>>> Row i of the result is a weighted average of the 9 rows of V,
>>> so it must have V's width (d_v = 4), not the token count.
STEPS 6 and 7 : concatenate the heads, then project with W_O
==============================================================================
head0_out (9, 4) + head1_out (9, 4) side by side -> concat (9, 8)
token | ---- head 0 (dims 0-3) ---- | ---- head 1 (dims 4-7) ----
--------------------------------------------------------------------------------------------
the | -1.056 -0.609 -0.373 +0.150 | +0.444 +0.965 +0.219 +0.641
cat | -1.043 -0.545 -0.128 -0.077 | +1.377 +0.827 +0.138 +0.455
sat | -1.195 -0.542 -0.307 +0.081 | -1.675 +1.218 +0.676 +0.960
on | -1.373 -0.387 -0.771 +0.019 | -1.276 +1.107 +0.730 +0.920
the | -1.177 -0.543 -0.567 +0.075 | +1.481 +0.723 +0.177 +0.401
mat | -1.611 -1.020 -0.317 +0.854 | +1.799 +0.757 +0.101 +0.299
and | -1.461 -0.731 -0.478 +0.355 | +1.617 +0.604 +0.223 +0.383
was | -1.678 -0.833 -0.516 +0.723 | +0.761 +0.511 +0.451 +0.556
looking | -1.426 -0.606 -0.399 +0.233 | -1.386 +1.132 +0.539 +0.899
matches concat_e: True
==============================================================================
concat (9, 8) @ W_O (8, 8) = attn_out (9, 8)
shape rule: (9 x 8) @ (8 x 8) -> (9 x 8) inner dim 8 cancels
W_O is what MIXES the two heads — without it they stay separate
==============================================================================
concat =
the | -1.056 -0.609 -0.373 +0.150 +0.444 +0.965 +0.219 +0.641
cat | -1.043 -0.545 -0.128 -0.077 +1.377 +0.827 +0.138 +0.455
sat | -1.195 -0.542 -0.307 +0.081 -1.675 +1.218 +0.676 +0.960
on | -1.373 -0.387 -0.771 +0.019 -1.276 +1.107 +0.730 +0.920
the | -1.177 -0.543 -0.567 +0.075 +1.481 +0.723 +0.177 +0.401
mat | -1.611 -1.020 -0.317 +0.854 +1.799 +0.757 +0.101 +0.299
and | -1.461 -0.731 -0.478 +0.355 +1.617 +0.604 +0.223 +0.383
was | -1.678 -0.833 -0.516 +0.723 +0.761 +0.511 +0.451 +0.556
looking | -1.426 -0.606 -0.399 +0.233 -1.386 +1.132 +0.539 +0.899
W_O =
-0.504 -0.106 +0.306 +0.237 +0.601 -0.272 -0.002 -0.209
-0.073 +0.437 +0.226 +0.468 -0.363 +0.457 -0.539 +0.476
-0.209 -0.227 +0.015 +0.540 -0.554 -0.181 +0.437 -0.127
-0.368 -0.092 -0.140 +0.509 +0.321 +0.006 +0.598 -0.043
-0.541 -0.291 -0.078 -0.420 +0.548 +0.327 -0.284 +0.111
+0.271 +0.540 +0.467 +0.019 +0.336 +0.076 +0.065 -0.472
-0.458 +0.036 -0.283 -0.587 +0.273 -0.573 +0.117 +0.048
+0.574 -0.019 -0.534 +0.063 +0.418 +0.609 -0.085 -0.220
attn_out =
the | +0.889 +0.304 -0.476 -0.917 +0.737 +0.560 +0.165 -0.565
cat | +0.298 -0.049 -0.437 -1.226 +0.878 +0.768 -0.166 -0.353
sat | +2.155 +1.103 -0.508 -0.271 -0.249 -0.125 +0.760 -0.912
on | +2.059 +1.127 -0.603 -0.726 +0.005 +0.145 +0.327 -0.632
the | +0.268 -0.033 -0.544 -1.488 +1.099 +0.857 -0.295 -0.200
mat | -0.004 -0.399 -0.824 -1.378 +1.245 +0.804 +0.451 -0.363
and | +0.166 -0.233 -0.781 -1.540 +1.076 +0.832 -0.026 -0.176
was | +0.589 -0.076 -1.056 -1.237 +0.755 +0.541 +0.482 -0.268
looking | +2.087 +0.973 -0.608 -0.374 -0.197 +0.056 +0.747 -0.810
worked element attn_out[0,0] = row 0 of concat . column 0 of W_O
row 0 of concat = [-1.056 -0.61 -0.373 0.15 0.444 0.965 0.219 0.641]
col 0 of W_O = [-0.504 -0.073 -0.209 -0.368 -0.541 0.271 -0.458 0.574]
(-1.0557)*(-0.5044) + (-0.6095)*(-0.0727) + (-0.3735)*(-0.2093) + (+0.1498)*(-0.3685) +
(+0.4438)*(-0.5410) + (+0.9654)*(+0.2712) + (+0.2191)*(-0.4579) + (+0.6409)*(+0.5742)
= +0.889213 | from array: +0.889213 | MATCH: True
Encoder Add & Norm (1)
$$X' = \mathrm{LayerNorm}\big(X_{enc} + \mathrm{MultiHead}(X_{enc})\big)$$
The residual + is what lets the original token identity survive an arbitrary number of layers; the LayerNorm is what keeps the magnitude from drifting. Watch the mean and standard deviation before and after.
attention output : mean=+0.0581 std=0.8073 after residual add : mean=+0.4630 std=1.2245 after LayerNorm : mean=+0.0000 std=1.0000 <- rescaled Encoder Add&Norm 1 output (shape (9, 8)): [[ 0.411 0.715 -1.678 -1.084 0.14 1.521 0.693 -0.718] [ 0.175 0.808 -1.938 0.333 -0.421 1.739 -0.288 -0.408] [ 2.439 0.46 -0.19 -0.744 -0.779 -0.183 -0.735 -0.268] [ 1.913 -0.932 0.534 -1.464 -0.105 0.682 -0.745 0.116] [-0.502 -0.768 -0.893 -1.211 0.82 1.975 0.565 0.014] [-1.366 0.494 -1.875 0.498 0.814 0.489 1.096 -0.149] [-0.831 -1.17 -1.532 0.369 0.84 1.15 -0.021 1.195] [-0.644 1.286 -0.341 -0.976 -0.26 1.483 -1.379 0.832] [ 1.965 0.004 0.529 -0.079 -1.908 -0.457 -0.002 -0.052]]
Add & Norm Operands
==============================================================================
residual = X_enc + attn_out all three are (9, 8) — element-wise, no mixing
==============================================================================
the |
X_enc: -0.100 +0.796 -0.878 +0.173 -0.226 +1.367 +0.914 +0.197
attn_out: +0.889 +0.304 -0.476 -0.917 +0.737 +0.560 +0.165 -0.565
residual: +0.789 +1.101 -1.354 -0.744 +0.511 +1.927 +1.078 -0.368
cat |
X_enc: +0.157 +1.331 -1.870 +1.888 -1.201 +1.731 +0.015 +0.046
attn_out: +0.298 -0.049 -0.437 -1.226 +0.878 +0.768 -0.166 -0.353
residual: +0.455 +1.282 -2.307 +0.662 -0.324 +2.499 -0.150 -0.307
sat |
X_enc: +2.302 +0.423 +1.071 +0.014 -0.059 +0.699 -1.004 +1.360
attn_out: +2.155 +1.103 -0.508 -0.271 -0.249 -0.125 +0.760 -0.912
residual: +4.457 +1.526 +0.563 -0.257 -0.308 +0.575 -0.243 +0.448
on |
X_enc: +1.362 -1.867 +2.008 -0.792 +0.465 +1.475 -0.793 +1.424
attn_out: +2.059 +1.127 -0.603 -0.726 +0.005 +0.145 +0.327 -0.632
residual: +3.421 -0.739 +1.404 -1.518 +0.470 +1.620 -0.466 +0.793
the |
X_enc: -0.857 -0.857 -0.489 +0.094 -0.186 +1.367 +0.918 +0.197
attn_out: +0.268 -0.033 -0.544 -1.488 +1.099 +0.857 -0.295 -0.200
residual: -0.589 -0.890 -1.033 -1.394 +0.912 +2.223 +0.623 -0.003
mat |
X_enc: -0.725 +1.319 -0.357 +2.302 -0.042 +0.111 +1.003 +0.713
attn_out: -0.004 -0.399 -0.824 -1.378 +1.245 +0.804 +0.451 -0.363
residual: -0.729 +0.920 -1.180 +0.924 +1.204 +0.915 +1.454 +0.350
and |
X_enc: -0.815 -0.747 -0.552 +2.061 -0.096 +0.450 +0.166 +1.502
attn_out: +0.166 -0.233 -0.781 -1.540 +1.076 +0.832 -0.026 -0.176
residual: -0.649 -0.980 -1.333 +0.521 +0.980 +1.282 +0.140 +1.326
was |
X_enc: -0.207 +2.033 +1.686 +1.348 -0.060 +1.577 -0.700 +1.856
attn_out: +0.589 -0.076 -1.056 -1.237 +0.755 +0.541 +0.482 -0.268
residual: +0.382 +1.958 +0.630 +0.111 +0.696 +2.119 -0.218 +1.587
looking |
X_enc: +1.223 -0.394 +1.918 +0.837 -1.887 -0.120 -0.178 +1.310
attn_out: +2.087 +0.973 -0.608 -0.374 -0.197 +0.056 +0.747 -0.810
residual: +3.311 +0.579 +1.310 +0.463 -2.084 -0.064 +0.569 +0.500
check element [0,0]: -0.100154 + +0.889213 = +0.789059 MATCH: True
Then LayerNorm, row by row (each token uses only its own 8 numbers):
token | mean std | before -> after (dim 0..7)
--------------------------------------------------------------------------------------------
the | +0.3675 1.0255 | +0.79 +1.10 -1.35 -0.74 +0.51 +1.93 +1.08 -0.37
| | +0.41 +0.71 -1.68 -1.08 +0.14 +1.52 +0.69 -0.72 <- normalized
cat | +0.2262 1.3070 | +0.46 +1.28 -2.31 +0.66 -0.32 +2.50 -0.15 -0.31
| | +0.18 +0.81 -1.94 +0.33 -0.42 +1.74 -0.29 -0.41 <- normalized
sat | +0.8451 1.4812 | +4.46 +1.53 +0.56 -0.26 -0.31 +0.57 -0.24 +0.45
| | +2.44 +0.46 -0.19 -0.74 -0.78 -0.18 -0.73 -0.27 <- normalized
on | +0.6231 1.4626 | +3.42 -0.74 +1.40 -1.52 +0.47 +1.62 -0.47 +0.79
| | +1.91 -0.93 +0.53 -1.46 -0.10 +0.68 -0.74 +0.12 <- normalized
the | -0.0187 1.1352 | -0.59 -0.89 -1.03 -1.39 +0.91 +2.22 +0.62 -0.00
| | -0.50 -0.77 -0.89 -1.21 +0.82 +1.98 +0.57 +0.01 <- normalized
mat | +0.4821 0.8865 | -0.73 +0.92 -1.18 +0.92 +1.20 +0.92 +1.45 +0.35
| | -1.37 +0.49 -1.88 +0.50 +0.81 +0.49 +1.10 -0.15 <- normalized
and | +0.1609 0.9748 | -0.65 -0.98 -1.33 +0.52 +0.98 +1.28 +0.14 +1.33
| | -0.83 -1.17 -1.53 +0.37 +0.84 +1.15 -0.02 +1.20 <- normalized
was | +0.9080 0.8164 | +0.38 +1.96 +0.63 +0.11 +0.70 +2.12 -0.22 +1.59
| | -0.64 +1.29 -0.34 -0.98 -0.26 +1.48 -1.38 +0.83 <- normalized
looking | +0.5729 1.3930 | +3.31 +0.58 +1.31 +0.46 -2.08 -0.06 +0.57 +0.50
| | +1.97 +0.00 +0.53 -0.08 -1.91 -0.46 -0.00 -0.05 <- normalized
10. Encoder: Feed-Forward Network
$$\mathrm{FFN}(x) = \max(0,\, xW_1 + b_1)\,W_2 + b_2, \qquad W_1 \in \mathbb{R}^{8 \times 32},\; W_2 \in \mathbb{R}^{32 \times 8}$$
"Position-wise" is the key word: the same $W_1, W_2$ are applied to each of the 15 token vectors independently. There is no mixing across positions here — all cross-token communication happened in the attention sub-layer. Equivalently, this is a 1-D convolution with kernel size 1.
Why widen to $d_{ff} = 4 d_{model}$ and come back? The expansion plus ReLU is what gives the block its non-linearity and its capacity; attention alone is (up to the softmax) a linear re-weighting of values. In a trained model the FFN holds the majority of the parameters.
ReLU: $\max(0, z)$ — it zeroes the negative half, and here roughly half the 32 hidden units will be inactive per token.
input (9, 8) @ W1 + b1 -> (9, 32) ReLU -> (9, 32) (47.9% of units zeroed) @ W2 + b2 -> (9, 8) FFN output: [[-0.356 -0.582 0.497 -0.98 -0.271 -0.415 0.712 -0.682] [-0.286 -0.416 0.454 -1.062 0.248 -1.02 0.6 -1.103] [ 0.929 -0.895 0.363 -0.304 0.117 -0.872 0.691 -0.045] [ 1.508 -0.685 -0.219 -0.446 -0.421 -0.528 0.863 0.345] [-0.052 -0.962 0.084 -0.346 0.016 -0.607 0.336 -0.609] [-0.388 -0.072 0.468 -1.447 -0.708 -0.897 0.807 -0.654] [ 0.321 -0.755 0.363 -0.291 -0.218 -1.309 0.955 -0.376] [ 0.28 -0.163 -0.124 0.022 0.478 -0.853 0.423 -0.951] [ 0.614 -0.12 -0.11 -0.333 0.078 -0.498 0.419 0.031]] Hand check, token 0 hidden unit 5: x . W1[:,5] + b1 = 0.586878 -> ReLU = 0.586878 from arrays: pre-ReLU 0.586878, post-ReLU 0.586878
FFN Matrices
The requested view: the $(9,8)$ input printed next to $W_1$ $(8,32)$, so the multiplication is visible rather than implied.
==============================================================================
FFN input (9,8) (9, 8) @ W1 (8,32) (8, 32) = pre-ReLU (9,32) (9, 32)
shape rule: (9 x 8) @ (8 x 32) -> (9 x 32) inner dim 8 cancels
each of the 9 tokens is expanded from 8 to 32 dims, independently
==============================================================================
FFN input (9,8) =
the | +0.411 +0.715 -1.678 -1.084 +0.140 +1.521 +0.693 -0.718
cat | +0.175 +0.808 -1.938 +0.333 -0.421 +1.739 -0.288 -0.408
sat | +2.439 +0.460 -0.190 -0.744 -0.779 -0.183 -0.735 -0.268
on | +1.913 -0.932 +0.534 -1.464 -0.105 +0.682 -0.745 +0.116
the | -0.502 -0.768 -0.893 -1.211 +0.820 +1.975 +0.565 +0.014
mat | -1.366 +0.494 -1.875 +0.498 +0.814 +0.489 +1.096 -0.149
and | -0.831 -1.170 -1.532 +0.369 +0.840 +1.150 -0.021 +1.195
was | -0.644 +1.286 -0.341 -0.976 -0.260 +1.483 -1.379 +0.832
looking | +1.965 +0.004 +0.529 -0.079 -1.908 -0.457 -0.002 -0.052
W1 (8,32) = [showing first 12 of 32 columns]
+0.315 -0.239 +0.089 -0.073 -0.355 -0.118 -0.182 +0.144 +0.348 -0.180 +0.241 -0.344
+0.349 -0.090 -0.059 +0.166 +0.080 -0.156 -0.069 -0.219 +0.361 -0.286 -0.099 -0.367
+0.372 -0.267 -0.212 -0.223 +0.227 -0.243 -0.336 +0.386 -0.368 +0.260 -0.104 +0.265
-0.129 +0.017 +0.066 -0.381 +0.202 +0.165 -0.104 -0.346 +0.197 +0.353 -0.201 -0.137
+0.186 -0.320 +0.131 -0.069 -0.023 -0.349 -0.239 +0.225 -0.036 -0.154 -0.090 +0.275
-0.308 +0.177 -0.233 -0.100 +0.001 +0.175 -0.300 -0.073 -0.012 +0.032 +0.247 -0.254
+0.309 +0.307 +0.144 +0.278 -0.319 +0.283 -0.139 -0.249 -0.202 -0.064 -0.193 -0.139
+0.043 +0.021 +0.120 -0.207 +0.282 -0.148 +0.150 -0.375 -0.099 -0.115 +0.034 -0.331
pre-ReLU (9,32) = [showing first 12 of 32 columns]
the | -0.367 +0.688 -0.043 +1.056 -1.113 +0.587 -0.142 -0.352 +0.713 -1.031 +0.626 -0.907
cat | -1.148 +0.753 -0.148 +0.287 -0.381 +0.808 +0.087 -1.012 +1.225 -0.464 +0.605 -1.298
sat | +0.627 -0.601 -0.015 +0.147 -0.848 -0.365 -0.032 +0.557 +1.142 -0.690 +0.870 -0.934
on | +0.210 -0.613 -0.250 -0.148 -0.657 -0.524 -0.370 +1.262 -0.022 -0.383 +1.118 -0.112
the | -0.885 +0.668 -0.160 +0.470 -0.523 +0.412 -0.295 +0.065 -0.530 -0.450 +0.597 +0.024
mat | -0.688 +0.950 +0.413 +0.641 -0.208 +0.755 +0.276 -1.271 +0.251 -0.369 -0.452 -0.279
and | -1.442 +0.672 +0.328 -0.358 +0.255 +0.438 +0.345 -0.920 -0.231 -0.012 +0.253 -0.194
was | -0.652 +0.053 -0.603 +0.023 +0.740 -0.365 +0.179 -0.302 +0.361 -0.606 +0.632 -0.739
looking | +0.612 -0.085 -0.092 -0.043 -0.566 +0.219 +0.059 +0.138 +0.553 +0.040 +0.492 -0.917
worked element pre-ReLU (9,32)[0,0] = row 0 of FFN input (9,8) . column 0 of W1 (8,32)
row 0 of FFN input (9,8) = [ 0.411 0.715 -1.678 -1.084 0.14 1.521 0.693 -0.718]
col 0 of W1 (8,32) = [ 0.315 0.349 0.372 -0.129 0.186 -0.308 0.309 0.043]
(+0.4110)*(+0.3150) + (+0.7148)*(+0.3490) + (-1.6782)*(+0.3724) + (-1.0841)*(-0.1285) +
(+0.1398)*(+0.1856) + (+1.5211)*(-0.3084) + (+0.6930)*(+0.3088) + (-0.7175)*(+0.0428)
= -0.366652 | from array: -0.366652 | MATCH: True
ReLU zeroes every negative entry — 138 of 288 entries (47.9%):
token | negatives killed | first 12 units after ReLU
----------------------------------------------------------------------------------------
the | 15 | 0.00 0.69 0.00 1.06 0.00 0.59 0.00 0.00 0.71 0.00 0.63 0.00
cat | 15 | 0.00 0.75 0.00 0.29 0.00 0.81 0.09 0.00 1.22 0.00 0.61 0.00
sat | 14 | 0.63 0.00 0.00 0.15 0.00 0.00 0.00 0.56 1.14 0.00 0.87 0.00
on | 17 | 0.21 0.00 0.00 0.00 0.00 0.00 0.00 1.26 0.00 0.00 1.12 0.00
the | 16 | 0.00 0.67 0.00 0.47 0.00 0.41 0.00 0.07 0.00 0.00 0.60 0.02
mat | 18 | 0.00 0.95 0.41 0.64 0.00 0.75 0.28 0.00 0.25 0.00 0.00 0.00
and | 13 | 0.00 0.67 0.33 0.00 0.26 0.44 0.35 0.00 0.00 0.00 0.25 0.00
was | 15 | 0.00 0.05 0.00 0.02 0.74 0.00 0.18 0.00 0.36 0.00 0.63 0.00
looking | 15 | 0.61 0.00 0.00 0.00 0.00 0.22 0.06 0.14 0.55 0.04 0.49 0.00
==============================================================================
post-ReLU (9,32) (9, 32) @ W2 (32,8) (32, 8) = FFN out (9,8) (9, 8)
shape rule: (9 x 32) @ (32 x 8) -> (9 x 8) inner dim 32 cancels
contraction back to d_model: each output element sums 32 products
==============================================================================
post-ReLU (9,32) = [showing first 12 of 32 columns]
the | +0.000 +0.688 +0.000 +1.056 +0.000 +0.587 +0.000 +0.000 +0.713 +0.000 +0.626 +0.000
cat | +0.000 +0.753 +0.000 +0.287 +0.000 +0.808 +0.087 +0.000 +1.225 +0.000 +0.605 +0.000
sat | +0.627 +0.000 +0.000 +0.147 +0.000 +0.000 +0.000 +0.557 +1.142 +0.000 +0.870 +0.000
on | +0.210 +0.000 +0.000 +0.000 +0.000 +0.000 +0.000 +1.262 +0.000 +0.000 +1.118 +0.000
the | +0.000 +0.668 +0.000 +0.470 +0.000 +0.412 +0.000 +0.065 +0.000 +0.000 +0.597 +0.024
mat | +0.000 +0.950 +0.413 +0.641 +0.000 +0.755 +0.276 +0.000 +0.251 +0.000 +0.000 +0.000
and | +0.000 +0.672 +0.328 +0.000 +0.255 +0.438 +0.345 +0.000 +0.000 +0.000 +0.253 +0.000
was | +0.000 +0.053 +0.000 +0.023 +0.740 +0.000 +0.179 +0.000 +0.361 +0.000 +0.632 +0.000
looking | +0.612 +0.000 +0.000 +0.000 +0.000 +0.219 +0.059 +0.138 +0.553 +0.040 +0.492 +0.000
W2 (32,8) =
-0.311 -0.208 +0.162 +0.251 +0.277 -0.066 -0.318 -0.048
-0.209 -0.162 +0.211 -0.329 +0.254 -0.056 +0.074 -0.318
-0.126 -0.318 +0.358 +0.238 -0.031 -0.214 +0.336 -0.120
-0.263 +0.049 +0.344 -0.384 -0.203 +0.340 -0.255 -0.300
+0.054 +0.280 -0.018 +0.383 +0.139 -0.330 +0.382 -0.071
-0.165 +0.045 -0.152 -0.320 -0.386 -0.039 +0.261 +0.056
+0.387 -0.165 -0.170 -0.264 -0.260 -0.166 +0.227 +0.314
+0.368 -0.207 +0.295 -0.133 -0.300 +0.287 +0.105 -0.260
-0.054 -0.165 +0.183 -0.036 +0.379 -0.345 -0.112 -0.014
+0.002 -0.143 +0.178 +0.290 +0.340 +0.236 +0.272 +0.039
+0.252 -0.152 -0.079 -0.142 -0.025 -0.166 +0.358 -0.145
-0.311 +0.021 -0.343 -0.302 -0.217 -0.028 -0.052 +0.232
+0.185 -0.057 -0.129 -0.142 +0.212 -0.347 -0.377 +0.229
-0.018 -0.286 -0.205 +0.335 +0.278 +0.162 -0.014 +0.016
+0.368 +0.361 -0.221 +0.241 -0.235 -0.222 +0.136 -0.074
+0.334 -0.318 +0.009 +0.324 +0.313 -0.266 -0.103 +0.092
-0.264 +0.363 +0.321 -0.258 -0.293 -0.321 +0.317 -0.237
-0.065 -0.040 +0.077 -0.332 -0.212 +0.334 +0.316 +0.218
-0.372 -0.134 +0.099 -0.174 +0.317 +0.104 -0.332 +0.317
+0.355 -0.327 +0.025 -0.152 -0.083 -0.248 +0.060 +0.355
+0.002 -0.323 +0.093 -0.124 +0.276 -0.051 -0.238 -0.239
+0.208 -0.182 +0.207 -0.228 -0.138 -0.027 +0.343 +0.334
-0.227 -0.110 -0.272 +0.008 -0.031 +0.310 +0.333 -0.090
-0.006 +0.307 +0.238 -0.384 -0.139 -0.061 -0.209 -0.233
-0.358 +0.111 -0.184 -0.347 -0.208 +0.234 +0.272 -0.013
-0.065 +0.370 -0.232 -0.220 -0.161 -0.253 -0.093 +0.318
+0.089 -0.068 +0.336 -0.350 -0.224 +0.051 +0.090 -0.357
-0.075 +0.116 +0.153 +0.107 -0.139 -0.315 -0.056 -0.288
+0.347 -0.019 -0.142 +0.206 -0.002 -0.147 +0.242 +0.284
+0.134 -0.351 +0.051 +0.204 -0.282 -0.210 +0.296 -0.372
+0.196 -0.020 -0.041 -0.382 +0.247 -0.109 +0.032 +0.137
+0.157 +0.353 +0.024 -0.078 +0.351 +0.110 -0.156 -0.289
FFN out (9,8) =
the | -0.356 -0.582 +0.497 -0.980 -0.271 -0.415 +0.712 -0.682
cat | -0.286 -0.416 +0.454 -1.062 +0.248 -1.020 +0.600 -1.103
sat | +0.929 -0.895 +0.363 -0.304 +0.117 -0.872 +0.691 -0.045
on | +1.508 -0.685 -0.219 -0.446 -0.421 -0.528 +0.863 +0.345
the | -0.052 -0.962 +0.084 -0.346 +0.016 -0.607 +0.336 -0.609
mat | -0.388 -0.072 +0.468 -1.447 -0.708 -0.897 +0.807 -0.654
and | +0.321 -0.755 +0.363 -0.291 -0.218 -1.309 +0.955 -0.376
was | +0.280 -0.163 -0.124 +0.022 +0.478 -0.853 +0.423 -0.951
looking | +0.614 -0.120 -0.110 -0.333 +0.078 -0.498 +0.419 +0.031
worked element FFN out (9,8)[0,0] = row 0 of post-ReLU (9,32) . column 0 of W2 (32,8)
row 0 of post-ReLU (9,32) = [0. 0.688 0. 1.056 0. 0.587 0. 0. 0.713 0. 0.626 0. 0.085 0. 0. 0.62 0.248 0. 0.013 0.14 0.106 0.729 1.216 0. 0. 0.385
0.111 0.805 0. 0.296 0. 0. ]
col 0 of W2 (32,8) = [-0.311 -0.209 -0.126 -0.263 0.054 -0.165 0.387 0.368 -0.054 0.002 0.252 -0.311 0.185 -0.018 0.368 0.334 -0.264 -0.065 -0.372 0.355 0.002 0.208
-0.227 -0.006 -0.358 -0.065 0.089 -0.075 0.347 0.134 0.196 0.157]
(+0.0000)*(-0.3110) + (+0.6883)*(-0.2086) + (+0.0000)*(-0.1260) + (+1.0558)*(-0.2626) +
(+0.0000)*(+0.0544) + (+0.5869)*(-0.1647) + (+0.0000)*(+0.3867) + (+0.0000)*(+0.3680) +
(+0.7128)*(-0.0544) + (+0.0000)*(+0.0024) + (+0.6262)*(+0.2524) + (+0.0000)*(-0.3114) +
(+0.0851)*(+0.1855) + (+0.0000)*(-0.0177) + (+0.0000)*(+0.3681) + (+0.6197)*(+0.3343) +
(+0.2479)*(-0.2638) + (+0.0000)*(-0.0648) + (+0.0125)*(-0.3723) + (+0.1395)*(+0.3546) +
(+0.1062)*(+0.0024) + (+0.7287)*(+0.2082) + (+1.2161)*(-0.2271) + (+0.0000)*(-0.0064) +
(+0.0000)*(-0.3580) + (+0.3846)*(-0.0650) + (+0.1105)*(+0.0893) + (+0.8052)*(-0.0751) +
(+0.0000)*(+0.3467) + (+0.2959)*(+0.1340) + (+0.0000)*(+0.1958) + (+0.0000)*(+0.1570)
= -0.355935 | from array: -0.355935 | MATCH: True
Final Add & Norm of the encoder:
==============================================================================
residual = LN1 out + FFN out all three are (9, 8) — element-wise, no mixing
==============================================================================
the |
LN1 out: +0.411 +0.715 -1.678 -1.084 +0.140 +1.521 +0.693 -0.718
FFN out: -0.356 -0.582 +0.497 -0.980 -0.271 -0.415 +0.712 -0.682
residual: +0.055 +0.133 -1.181 -2.064 -0.131 +1.106 +1.405 -1.399
cat |
LN1 out: +0.175 +0.808 -1.938 +0.333 -0.421 +1.739 -0.288 -0.408
FFN out: -0.286 -0.416 +0.454 -1.062 +0.248 -1.020 +0.600 -1.103
residual: -0.111 +0.392 -1.484 -0.729 -0.172 +0.719 +0.311 -1.511
sat |
LN1 out: +2.439 +0.460 -0.190 -0.744 -0.779 -0.183 -0.735 -0.268
FFN out: +0.929 -0.895 +0.363 -0.304 +0.117 -0.872 +0.691 -0.045
residual: +3.368 -0.435 +0.173 -1.048 -0.661 -1.055 -0.044 -0.313
on |
LN1 out: +1.913 -0.932 +0.534 -1.464 -0.105 +0.682 -0.745 +0.116
FFN out: +1.508 -0.685 -0.219 -0.446 -0.421 -0.528 +0.863 +0.345
residual: +3.421 -1.617 +0.315 -1.910 -0.525 +0.154 +0.119 +0.461
the |
LN1 out: -0.502 -0.768 -0.893 -1.211 +0.820 +1.975 +0.565 +0.014
FFN out: -0.052 -0.962 +0.084 -0.346 +0.016 -0.607 +0.336 -0.609
residual: -0.554 -1.730 -0.809 -1.557 +0.836 +1.368 +0.902 -0.595
mat |
LN1 out: -1.366 +0.494 -1.875 +0.498 +0.814 +0.489 +1.096 -0.149
FFN out: -0.388 -0.072 +0.468 -1.447 -0.708 -0.897 +0.807 -0.654
residual: -1.755 +0.422 -1.407 -0.949 +0.106 -0.408 +1.902 -0.803
and |
LN1 out: -0.831 -1.170 -1.532 +0.369 +0.840 +1.150 -0.021 +1.195
FFN out: +0.321 -0.755 +0.363 -0.291 -0.218 -1.309 +0.955 -0.376
residual: -0.510 -1.925 -1.169 +0.078 +0.622 -0.160 +0.934 +0.819
was |
LN1 out: -0.644 +1.286 -0.341 -0.976 -0.260 +1.483 -1.379 +0.832
FFN out: +0.280 -0.163 -0.124 +0.022 +0.478 -0.853 +0.423 -0.951
residual: -0.365 +1.123 -0.465 -0.955 +0.218 +0.630 -0.956 -0.119
looking |
LN1 out: +1.965 +0.004 +0.529 -0.079 -1.908 -0.457 -0.002 -0.052
FFN out: +0.614 -0.120 -0.110 -0.333 +0.078 -0.498 +0.419 +0.031
residual: +2.579 -0.116 +0.419 -0.412 -1.830 -0.955 +0.417 -0.021
check element [0,0]: +0.411049 + -0.355935 = +0.055114 MATCH: True
After LayerNorm -> encoder_output H (9, 8):
[[ 0.277 0.346 -0.812 -1.589 0.113 1.203 1.466 -1.004]
[ 0.268 0.905 -1.47 -0.513 0.191 1.319 0.803 -1.504]
[ 2.52 -0.324 0.131 -0.782 -0.493 -0.787 -0.031 -0.233]
[ 2.211 -1.096 0.173 -1.288 -0.379 0.067 0.044 0.268]
[-0.262 -1.339 -0.496 -1.181 1.01 1.498 1.07 -0.3 ]
[-1.276 0.717 -0.958 -0.538 0.428 -0.042 2.074 -0.405]
[-0.368 -1.872 -1.069 0.257 0.835 0.004 1.167 1.045]
[-0.368 1.791 -0.514 -1.225 0.478 1.076 -1.227 -0.011]
[ 2.148 -0.105 0.342 -0.353 -1.538 -0.807 0.34 -0.026]]
(the next cell assigns this to `encoder_output`)
Encoder Add & Norm (2)
$$\text{encoder\_output} = \mathrm{LayerNorm}\big(X' + \mathrm{FFN}(X')\big)$$
This $(9 \times 8)$ matrix, call it $H$, is the memory the decoder will read from through cross-attention. It is computed once and reused at every generation step. In a real model, $N = 6$ such blocks are stacked; here one block is enough to see the mechanics.
ENCODER OUTPUT H (shape (9, 8)): [[ 0.277 0.346 -0.812 -1.589 0.113 1.203 1.466 -1.004] [ 0.268 0.905 -1.47 -0.513 0.191 1.319 0.803 -1.504] [ 2.52 -0.324 0.131 -0.782 -0.493 -0.787 -0.031 -0.233] [ 2.211 -1.096 0.173 -1.288 -0.379 0.067 0.044 0.268] [-0.262 -1.339 -0.496 -1.181 1.01 1.498 1.07 -0.3 ] [-1.276 0.717 -0.958 -0.538 0.428 -0.042 2.074 -0.405] [-0.368 -1.872 -1.069 0.257 0.835 0.004 1.167 1.045] [-0.368 1.791 -0.514 -1.225 0.478 1.076 -1.227 -0.011] [ 2.148 -0.105 0.342 -0.353 -1.538 -0.807 0.34 -0.026]] mean = -0.0000, std = 1.0000 per-row std (should all be ~1 after LN): [1. 1. 1. 1. 1. 1. 1. 1. 1.] Compare with the old notebook, where the equivalent matrix had values around 80 000 and near-identical rows. Row-to-row spread here: 0.8410 (tokens remain distinguishable)
11. Decoder Input: Shift Right
This is the second bug from the old notebook (X_decoder_input = X). The decoder must not receive the prefix, and it must not receive the continuation unshifted.
The continuation is fed shifted right by one, with <sos> prepended:
| position $t$ | 0 | 1 | 2 | 3 | 4 | 5 | 6 |
|---|---|---|---|---|---|---|---|
| input (shifted) | <sos> |
at | a | tiny | mouse | running | by |
| label (to predict) | at | a | tiny | mouse | running | by | <eos> |
So at every position the decoder predicts the token one step to the right of its own input. This is exactly the autoregressive factorization
$$P(y_t \mid y_{<t},\, x)$$
written as a table.
Teacher forcing. During training the correct previous tokens are supplied rather than the model's own guesses, which lets all 7 positions be computed in a single parallel pass instead of 7 sequential ones. At inference time the true continuation does not exist, so the loop of section 19 feeds back the model's own predictions instead.
Combined with the causal mask (next section), the invariant is: the prediction at position $t$ depends only on the prefix $x$ and on continuation tokens $0 \dots t$ — never on $t+1$. Without the shift, or without the mask, the model could read the answer and the task would be trivial.
n_src = 9, n_tgt = 7 PREFIX (encoder): the cat sat on the mat and was looking pos decoder input id -> label (predict) id -------------------------------------------------------------- 0 <sos> 1 -> at 127 1 at 127 -> a 340 2 a 340 -> tiny 68 3 tiny 68 -> mouse 57 4 mouse 57 -> running 412 5 running 412 -> by 322 6 by 322 -> <eos> 2 X_dec = E[dec_input_ids] * sqrt(d_model) + PE shape (7, 8) [[ 1.382 1.649 -0.167 1.147 1.207 0.183 0.369 0.607] [ 2.015 0.91 -0.007 1.443 -1.561 -0.127 -1.193 1.143] [ 0.129 0.666 -1.151 0.565 0.054 1.111 -0.177 1.396] [ 0.449 -0.174 1.156 0.372 -0.137 1.282 -0.246 2.607] [-0.32 0.537 1.339 -0.564 -2.514 1.934 -1.363 0.775] [-1.805 -0.301 0.001 1.237 2.399 0.283 -0.412 0.937] [-0.123 1.984 0.74 -0.511 -0.352 1.13 -0.449 0.781]]
12. Decoder: Masked Self-Attention
Same operation as encoder self-attention, plus a look-ahead mask applied to the scores before the softmax:
$$M_{ij} = \begin{cases} 0 & j \le i \quad \text{(past and present: allowed)} \\[2pt] -\infty & j > i \quad \text{(future: blocked)} \end{cases}$$
$$A = \mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}} + M\right)$$
The mechanism: $\exp(-\infty) = 0$, so a blocked entry contributes 0 to both the numerator and the denominator of the softmax. The row still sums to 1 — the probability mass is simply redistributed over the allowed positions.
In code -1e9 is used rather than -np.inf. This matters: with -inf, a fully-masked row would produce 0/0 = NaN, and 0 * inf appears in some autograd paths. -1e9 is numerically indistinguishable after exp (it underflows to exactly 0 in float64) and is safe.
Consequences visible below:
- Row 0 (
<sos>) attends only to itself → weight exactly 1.0, all others 0. - Row $t$ has exactly $t+1$ non-zero entries.
- The weight matrix is lower-triangular.
Causal mask (7 x 7), True = blocked: [[0 1 1 1 1 1 1] [0 0 1 1 1 1 1] [0 0 0 1 1 1 1] [0 0 0 0 1 1 1] [0 0 0 0 0 1 1] [0 0 0 0 0 0 1] [0 0 0 0 0 0 0]] Allowed positions per row: [1 2 3 4 5 6 7]
[decoder masked self-attn] x_q (7, 8) , x_kv (7, 8)
Q,K,V per head: (2, 7, 4), (2, 7, 4), (2, 7, 4)
scores/weights: (2, 7, 7) concat: (7, 8) out: (7, 8)
--- Head 0: scores BEFORE masking (Q K^T / sqrt(d_k)) ---
[[ 0.5 0.56 0.12 1.41 0.44 0.59 0.41]
[ 0.97 2.16 -0.63 -0.04 0.29 -1.32 0.03]
[ 0.03 0.47 -0.49 0.88 0.46 0.09 -0.04]
[ 1.22 1.08 0.69 1.45 0.28 0.57 0.83]
[ 0.55 2.19 -0.59 -1. 0.48 -1.91 -0.17]
[-0.05 -0.91 0.52 2.1 0.36 1.87 0.51]
[ 0.26 0.75 -0.42 -0.87 -0.22 -1.02 -0.24]]
--- Head 0: scores AFTER masking ---
masked_scores = (blocked entries shown as -inf)
<sos> | +0.503 -inf -inf -inf -inf -inf -inf
at | +0.970 +2.161 -inf -inf -inf -inf -inf
a | +0.031 +0.475 -0.491 -inf -inf -inf -inf
tiny | +1.220 +1.080 +0.690 +1.452 -inf -inf -inf
mouse | +0.552 +2.191 -0.586 -1.000 +0.484 -inf -inf
running | -0.052 -0.905 +0.523 +2.104 +0.365 +1.867 -inf
by | +0.258 +0.751 -0.416 -0.868 -0.224 -1.020 -0.236
--- Head 0: attention weights after softmax ---
A = (blocked entries are exactly 0)
<sos> | 1.0000 0 0 0 0 0 0
at | 0.2330 0.7670 0 0 0 0 0
a | 0.3172 0.4946 0.1882 0 0 0 0
tiny | 0.2687 0.2338 0.1582 0.3392 0 0 0
mouse | 0.1313 0.6761 0.0421 0.0278 0.1227 0 0
running | 0.0496 0.0211 0.0881 0.4282 0.0752 0.3378 0
by | 0.2009 0.3289 0.1024 0.0652 0.1240 0.0560 0.1226
Row sums (must all be 1): [1. 1. 1. 1. 1. 1. 1.]
Row 0 = [1. 0. 0. 0. 0. 0. 0.] <- '<sos>' can only see itself
Strictly lower-triangular structure: True
Row 2, unmasked scores over positions 0..2: [ 0.031 0.475 -0.491] exp(s - max) : [0.641 1. 0.381] / sum = 2.0220 : [0.317 0.495 0.188] From the weights matrix (first 3 entries): [0.317 0.495 0.188] match: True remaining entries are exactly zero: True
The black upper triangle is the mask: no token attends to its future.
Masked Self-Attention Matrices
The same seven steps as the encoder, with both operands and the full result matrix printed at each one. Two things differ: every shape is 7 instead of 9, and step 4 applies the causal mask between the scores and the softmax.
STEP 1 : project X_dec to Q (K and V are identical in form)
==============================================================================
==============================================================================
X_dec (7, 8) @ Wq (8, 8) = Q_full (7, 8)
shape rule: (7 x 8) @ (8 x 8) -> (7 x 8) inner dim 8 cancels
K uses Wk, V uses Wv - same shapes, same arithmetic
==============================================================================
X_dec =
<sos> | +1.382 +1.649 -0.167 +1.147 +1.207 +0.183 +0.369 +0.607
at | +2.015 +0.910 -0.007 +1.443 -1.561 -0.127 -1.193 +1.143
a | +0.129 +0.666 -1.151 +0.565 +0.054 +1.111 -0.177 +1.396
tiny | +0.449 -0.174 +1.156 +0.372 -0.137 +1.282 -0.246 +2.607
mouse | -0.320 +0.537 +1.339 -0.564 -2.514 +1.934 -1.363 +0.775
running | -1.805 -0.301 +0.001 +1.237 +2.399 +0.283 -0.412 +0.937
by | -0.123 +1.984 +0.740 -0.511 -0.352 +1.130 -0.449 +0.781
Wq =
+0.420 -0.169 +0.197 -0.036 -0.369 +0.125 -0.538 -0.106
+0.057 -0.078 +0.305 +0.402 +0.332 -0.565 -0.375 +0.045
+0.533 +0.410 +0.420 -0.241 -0.086 -0.036 -0.420 -0.574
+0.547 -0.318 -0.318 -0.510 +0.525 +0.098 +0.205 +0.504
-0.189 +0.149 -0.565 -0.384 +0.406 +0.327 -0.183 -0.151
+0.041 -0.612 -0.317 -0.357 -0.306 +0.375 -0.270 -0.378
+0.005 +0.448 -0.317 -0.516 -0.176 +0.303 +0.062 +0.388
+0.563 +0.358 +0.127 -0.008 -0.563 +0.593 -0.143 -0.426
Q_full =
<sos> | +1.335 -0.345 -0.441 -0.656 +0.681 +0.293 -1.390 +0.235
at | +2.610 -1.152 +1.655 +0.810 -0.711 -0.363 -1.044 -0.108
a | +0.609 -0.977 -0.584 -0.086 -0.503 +0.945 -0.240 -0.129
tiny | +2.543 +0.313 +0.482 -0.855 -2.001 +2.057 -1.293 -2.201
mouse | +1.285 -1.150 +2.180 +1.163 -1.923 -0.497 -0.963 -2.205
running | -0.014 +0.271 -2.035 -1.504 +1.646 +1.387 +0.663 -0.226
by | +0.726 -0.334 +1.135 +0.840 -0.477 -0.576 -1.472 -1.461
worked element Q_full[0,0] = row 0 of X_dec . column 0 of Wq
row 0 of X_dec = [ 1.382 1.649 -0.167 1.147 1.206 0.183 0.369 0.607]
col 0 of Wq = [ 0.419 0.057 0.532 0.547 -0.189 0.041 0.005 0.563]
(+1.3822)*(+0.4195) + (+1.6487)*(+0.0571) + (-0.1671)*(+0.5325) + (+1.1467)*(+0.5473) +
(+1.2065)*(-0.1891) + (+0.1831)*(+0.0411) + (+0.3687)*(+0.0052) + (+0.6067)*(+0.5629)
= +1.335453 | from array: +1.335453 | MATCH: True
K_full = X_dec @ Wk -> (7, 8)
V_full = X_dec @ Wv -> (7, 8)
STEP 2 : split the 8 columns into 2 heads of 4 ============================================================================== Q_full (7, 8) -> reshape(7, 2, 4) -> transpose -> (2, 7, 4) head 0 takes columns 0..3, head 1 takes columns 4..7 Q head 0 (7, 4): [[ 1.335 -0.345 -0.441 -0.656] [ 2.61 -1.152 1.655 0.81 ] [ 0.609 -0.977 -0.584 -0.086] [ 2.543 0.313 0.482 -0.855] [ 1.285 -1.15 2.18 1.163] [-0.014 0.271 -2.035 -1.504] [ 0.726 -0.334 1.135 0.84 ]] Q head 1 (7, 4): [[ 0.681 0.293 -1.39 0.235] [-0.711 -0.363 -1.044 -0.108] [-0.503 0.945 -0.24 -0.129] [-2.001 2.057 -1.293 -2.201] [-1.923 -0.497 -0.963 -2.205] [ 1.646 1.387 0.663 -0.226] [-0.477 -0.576 -1.472 -1.461]] matches the column slices of Q_full: True
STEP 3 : scores = Q0 @ K0^T, THEN divide by sqrt(d_k)
==============================================================================
==============================================================================
Q_head0 (7, 4) @ K_head0^T (4, 7) = raw_scores (7, 7)
shape rule: (7 x 4) @ (4 x 7) -> (7 x 7) inner dim 4 cancels
(7,4) @ (4,7) -> (7,7): 7 queries x 7 keys
==============================================================================
Q_head0 =
<sos> | +1.335 -0.345 -0.441 -0.656
at | +2.610 -1.152 +1.655 +0.810
a | +0.609 -0.977 -0.584 -0.086
tiny | +2.543 +0.313 +0.482 -0.855
mouse | +1.285 -1.150 +2.180 +1.163
running | -0.014 +0.271 -2.035 -1.504
by | +0.726 -0.334 +1.135 +0.840
K_head0^T =
+0.669 +0.249 -0.194 +0.637 -0.471 +0.225 +0.134
+0.121 -1.721 +0.590 -0.852 -1.635 +0.651 +0.002
+0.471 +1.726 +0.724 -0.595 +0.982 -1.013 +0.458
-0.553 -1.444 -1.567 -2.151 -2.104 -0.996 -1.296
raw_scores =
<sos> | +1.007 +1.113 +0.246 +2.818 +0.882 +1.175 +0.826
at | +1.939 +4.322 -1.256 -0.083 +0.577 -2.646 +0.055
a | +0.061 +0.950 -0.982 +1.752 +0.917 +0.178 -0.077
tiny | +2.439 +2.161 +1.380 +2.905 +0.562 +1.139 +1.669
mouse | +1.104 +4.382 -1.172 -2.001 +0.969 -3.825 -0.340
running | -0.104 -1.811 +1.046 +4.208 +0.729 +3.734 +1.017
by | +0.516 +1.502 -0.832 -1.736 -0.449 -2.041 -0.473
worked element raw_scores[0,0] = row 0 of Q_head0 . column 0 of K_head0^T
row 0 of Q_head0 = [ 1.335 -0.345 -0.441 -0.656]
col 0 of K_head0^T = [ 0.669 0.121 0.471 -0.552]
(+1.3355)*(+0.6694) + (-0.3451)*(+0.1215) + (-0.4407)*(+0.4709) + (-0.6556)*(-0.5525)
= +1.006698 | from array: +1.006698 | MATCH: True
then divide every entry by sqrt(d_k) = sqrt(4) = 2:
<sos> | +0.503 +0.557 +0.123 +1.409 +0.441 +0.588 +0.413
at | +0.970 +2.161 -0.628 -0.042 +0.288 -1.323 +0.027
a | +0.031 +0.475 -0.491 +0.876 +0.458 +0.089 -0.039
tiny | +1.220 +1.080 +0.690 +1.452 +0.281 +0.569 +0.835
mouse | +0.552 +2.191 -0.586 -1.000 +0.484 -1.913 -0.170
running | -0.052 -0.905 +0.523 +2.104 +0.365 +1.867 +0.508
by | +0.258 +0.751 -0.416 -0.868 -0.224 -1.020 -0.236
>>> still a full (7, 7) matrix - nothing is blocked yet.
STEP 4 : apply the causal mask, BEFORE the softmax
==============================================================================
Scaled scores + M. Entries above the diagonal are set to -inf,
so exp(-inf) = 0 and a future token can carry no weight.
masked_scores = (blocked entries shown as -inf)
<sos> | +0.503 -inf -inf -inf -inf -inf -inf
at | +0.970 +2.161 -inf -inf -inf -inf -inf
a | +0.031 +0.475 -0.491 -inf -inf -inf -inf
tiny | +1.220 +1.080 +0.690 +1.452 -inf -inf -inf
mouse | +0.552 +2.191 -0.586 -1.000 +0.484 -inf -inf
running | -0.052 -0.905 +0.523 +2.104 +0.365 +1.867 -inf
by | +0.258 +0.751 -0.416 -0.868 -0.224 -1.020 -0.236
(-1e9 is used in code as a stand-in for -inf: exp(-1e9) underflows to
exactly 0.0 in float64, while avoiding the 0 * inf = NaN edge case.)
After the softmax those blocked entries become hard zeros:
A = softmax(masked_scores) = (blocked entries are exactly 0)
<sos> | 1.0000 0 0 0 0 0 0
at | 0.2330 0.7670 0 0 0 0 0
a | 0.3172 0.4946 0.1882 0 0 0 0
tiny | 0.2687 0.2338 0.1582 0.3392 0 0 0
mouse | 0.1313 0.6761 0.0421 0.0278 0.1227 0 0
running | 0.0496 0.0211 0.0881 0.4282 0.0752 0.3378 0
by | 0.2009 0.3289 0.1024 0.0652 0.1240 0.0560 0.1226
upper triangle is exactly 0 : True
every row still sums to 1 : [1. 1. 1. 1. 1. 1. 1.]
non-zero entries per row : [1 2 3 4 5 6 7] <- row t has t+1
STEP 5 : weights @ V, the masked weighted sum
==============================================================================
==============================================================================
A (masked weights) (7, 7) @ V_head0 (7, 4) = head0_out (7, 4)
shape rule: (7 x 7) @ (7 x 4) -> (7 x 4) inner dim 7 cancels
(7,7)@(7,4) -> (7,4). Zero weights mean those V rows contribute nothing
==============================================================================
A (masked weights) =
<sos> | +1.000 +0.000 +0.000 +0.000 +0.000 +0.000 +0.000
at | +0.233 +0.767 +0.000 +0.000 +0.000 +0.000 +0.000
a | +0.317 +0.495 +0.188 +0.000 +0.000 +0.000 +0.000
tiny | +0.269 +0.234 +0.158 +0.339 +0.000 +0.000 +0.000
mouse | +0.131 +0.676 +0.042 +0.028 +0.123 +0.000 +0.000
running | +0.050 +0.021 +0.088 +0.428 +0.075 +0.338 +0.000
by | +0.201 +0.329 +0.102 +0.065 +0.124 +0.056 +0.123
V_head0 =
+0.597 -0.274 +1.040 -1.039
+1.296 +0.654 +1.575 -2.918
+0.529 -0.749 +1.454 +0.569
+1.053 -0.130 +2.242 +0.536
+1.675 +1.939 +0.404 -1.420
-0.656 -1.042 -0.113 +2.285
+1.490 +0.772 -0.180 -0.791
head0_out =
<sos> | +0.597 -0.274 +1.040 -1.039
at | +1.133 +0.438 +1.451 -2.480
a | +0.930 +0.096 +1.383 -1.666
tiny | +0.904 -0.083 +1.639 -0.690
mouse | +1.212 +0.609 +1.375 -2.245
running | +0.459 -0.328 +1.165 +0.832
by | +1.023 +0.352 +1.044 -1.221
worked element head0_out[0,0] = row 0 of A (masked weights) . column 0 of V_head0
row 0 of A (masked weights) = [1. 0. 0. 0. 0. 0. 0.]
col 0 of V_head0 = [ 0.597 1.296 0.529 1.053 1.675 -0.656 1.49 ]
(+1.0000)*(+0.5966) + (+0.0000)*(+1.2962) + (+0.0000)*(+0.5293) + (+0.0000)*(+1.0534) +
(+0.0000)*(+1.6746) + (+0.0000)*(-0.6561) + (+0.0000)*(+1.4896)
= +0.596564 | from array: +0.596564 | MATCH: True
>>> row 0 has a single non-zero weight (1.0 on itself), so its output
>>> equals V row 0 exactly: True
STEPS 6 and 7 : concatenate the heads, then project with W_O
==============================================================================
dec_head0 (7, 4) + dec_head1 (7, 4) side by side -> concat (7, 8)
token | ---- head 0 (dims 0-3) ---- | ---- head 1 (dims 4-7) ----
--------------------------------------------------------------------------------------------
<sos> | +0.597 -0.274 +1.040 -1.039 | -1.693 +0.565 +0.092 -1.168
at | +1.133 +0.438 +1.451 -2.480 | -1.629 +0.983 -0.221 -0.983
a | +0.930 +0.096 +1.383 -1.666 | -1.266 +0.861 -0.252 -0.567
tiny | +0.904 -0.083 +1.639 -0.690 | -1.002 +0.785 +0.393 -1.023
mouse | +1.212 +0.609 +1.375 -2.245 | -1.065 +0.435 +0.340 -0.832
running | +0.459 -0.328 +1.165 +0.832 | -1.040 +1.642 -0.355 -0.613
by | +1.023 +0.352 +1.044 -1.221 | -0.754 +0.231 +0.862 -0.728
==============================================================================
concat (7, 8) @ W_O (8, 8) = dec_attn_out (7, 8)
shape rule: (7 x 8) @ (8 x 8) -> (7 x 8) inner dim 8 cancels
W_O is what MIXES the two heads - this is the sub-layer's final output
==============================================================================
concat =
<sos> | +0.597 -0.274 +1.040 -1.039 -1.693 +0.565 +0.092 -1.168
at | +1.133 +0.438 +1.451 -2.480 -1.629 +0.983 -0.221 -0.983
a | +0.930 +0.096 +1.383 -1.666 -1.266 +0.861 -0.252 -0.567
tiny | +0.904 -0.083 +1.639 -0.690 -1.002 +0.785 +0.393 -1.023
mouse | +1.212 +0.609 +1.375 -2.245 -1.065 +0.435 +0.340 -0.832
running | +0.459 -0.328 +1.165 +0.832 -1.040 +1.642 -0.355 -0.613
by | +1.023 +0.352 +1.044 -1.221 -0.754 +0.231 +0.862 -0.728
W_O =
-0.227 +0.099 +0.282 -0.109 +0.158 +0.244 +0.259 -0.549
+0.441 -0.051 -0.291 -0.316 -0.021 +0.474 -0.149 +0.306
-0.494 -0.029 +0.602 +0.331 -0.355 +0.002 -0.326 -0.111
+0.343 +0.332 -0.333 +0.517 +0.113 -0.147 +0.002 +0.129
+0.073 -0.017 +0.006 -0.396 +0.383 +0.051 -0.452 +0.326
+0.530 +0.548 -0.160 -0.446 +0.023 -0.474 +0.532 -0.388
-0.303 -0.114 -0.143 +0.361 -0.054 -0.403 -0.285 +0.011
-0.432 +0.137 +0.081 +0.263 +0.536 +0.505 +0.546 -0.313
dec_attn_out =
<sos> | -0.474 -0.135 +1.012 -0.026 -1.654 -0.811 +0.254 -1.065
at | -0.739 -0.320 +1.676 -1.196 -1.742 -0.106 +0.533 -1.575
a | -0.739 -0.063 +1.466 -0.658 -1.290 -0.138 +0.563 -1.420
tiny | -0.623 +0.079 +1.225 +0.033 -1.450 -0.813 -0.091 -1.099
mouse | -1.048 -0.594 +1.548 -0.898 -1.425 +0.099 -0.070 -1.172
running | +0.627 +1.178 +0.381 +0.259 -0.911 -1.163 +0.898 -1.162
by | -0.891 -0.412 +0.997 -0.193 -1.074 -0.265 -0.311 -0.824
worked element dec_attn_out[0,0] = row 0 of concat . column 0 of W_O
row 0 of concat = [ 0.597 -0.274 1.04 -1.038 -1.694 0.565 0.092 -1.168]
col 0 of W_O = [-0.227 0.441 -0.494 0.343 0.073 0.53 -0.303 -0.432]
(+0.5966)*(-0.2273) + (-0.2739)*(+0.4406) + (+1.0401)*(-0.4938) + (-1.0385)*(+0.3430) +
(-1.6935)*(+0.0734) + (+0.5651)*(+0.5296) + (+0.0917)*(-0.3028) + (-1.1684)*(-0.4318)
= -0.474334 | from array: -0.474334 | MATCH: True
Decoder Add & Norm (1)
$$Y_1 = \mathrm{LayerNorm}\big(X_{dec} + \mathrm{MaskedMultiHead}(X_{dec})\big)$$
==============================================================================
residual = X_dec + attn_out all three are (7, 8) — element-wise, no mixing
==============================================================================
<sos> |
X_dec: +1.382 +1.649 -0.167 +1.147 +1.207 +0.183 +0.369 +0.607
attn_out: -0.474 -0.135 +1.012 -0.026 -1.654 -0.811 +0.254 -1.065
residual: +0.908 +1.514 +0.845 +1.120 -0.447 -0.628 +0.623 -0.458
at |
X_dec: +2.015 +0.910 -0.007 +1.443 -1.561 -0.127 -1.193 +1.143
attn_out: -0.739 -0.320 +1.676 -1.196 -1.742 -0.106 +0.533 -1.575
residual: +1.276 +0.590 +1.669 +0.247 -3.303 -0.232 -0.660 -0.432
a |
X_dec: +0.129 +0.666 -1.151 +0.565 +0.054 +1.111 -0.177 +1.396
attn_out: -0.739 -0.063 +1.466 -0.658 -1.290 -0.138 +0.563 -1.420
residual: -0.609 +0.603 +0.315 -0.093 -1.236 +0.973 +0.386 -0.024
tiny |
X_dec: +0.449 -0.174 +1.156 +0.372 -0.137 +1.282 -0.246 +2.607
attn_out: -0.623 +0.079 +1.225 +0.033 -1.450 -0.813 -0.091 -1.099
residual: -0.175 -0.096 +2.381 +0.406 -1.587 +0.469 -0.337 +1.509
mouse |
X_dec: -0.320 +0.537 +1.339 -0.564 -2.514 +1.934 -1.363 +0.775
attn_out: -1.048 -0.594 +1.548 -0.898 -1.425 +0.099 -0.070 -1.172
residual: -1.368 -0.057 +2.887 -1.462 -3.939 +2.033 -1.433 -0.396
running |
X_dec: -1.805 -0.301 +0.001 +1.237 +2.399 +0.283 -0.412 +0.937
attn_out: +0.627 +1.178 +0.381 +0.259 -0.911 -1.163 +0.898 -1.162
residual: -1.178 +0.877 +0.382 +1.497 +1.489 -0.880 +0.487 -0.225
by |
X_dec: -0.123 +1.984 +0.740 -0.511 -0.352 +1.130 -0.449 +0.781
attn_out: -0.891 -0.412 +0.997 -0.193 -1.074 -0.265 -0.311 -0.824
residual: -1.014 +1.572 +1.737 -0.704 -1.426 +0.865 -0.760 -0.043
check element [0,0]: +1.382159 + -0.474334 = +0.907825 MATCH: True
then LayerNorm, row by row -> Y1 (shape (7, 8)):
<sos> | +0.613 +1.398 +0.531 +0.889 -1.142 -1.377 +0.244 -1.156
at | +0.967 +0.487 +1.242 +0.247 -2.239 -0.089 -0.388 -0.229
a | -0.987 +0.858 +0.420 -0.201 -1.941 +1.421 +0.527 -0.097
tiny | -0.440 -0.370 +1.826 +0.075 -1.692 +0.131 -0.584 +1.053
mouse | -0.447 +0.204 +1.665 -0.494 -1.724 +1.241 -0.480 +0.035
running | -1.579 +0.607 +0.081 +1.266 +1.258 -1.261 +0.192 -0.565
by | -0.916 +1.356 +1.501 -0.644 -1.277 +0.735 -0.693 -0.063
mean=+0.0000, per-row std=[1. 1. 1. 1. 1. 1. 1.]
13. Decoder: Cross-Attention
This is the one place where the encoder and the decoder communicate, and the shapes stop being square:
$$Q = Y_1 W^Q \quad \text{from the \textbf{decoder}} \qquad K = H W^K, \;\; V = H W^V \quad \text{from the \textbf{encoder output } } H$$
| comes from | shape | |
|---|---|---|
| $Q$ | decoder, $Y_1$ | $(7, 8) \rightarrow$ per head $(2, 7, 4)$ |
| $K$, $V$ | encoder output $H$ | $(9, 8) \rightarrow$ per head $(2, 9, 4)$ |
| scores $QK^\top$ | $(2, \mathbf{7}, \mathbf{9})$ | |
| output | $(7, 8)$ |
Read the $7 \times 9$ score matrix as: for each of the 7 decoder positions, how much attention is paid to each of the 9 prefix words. Each row sums to 1 across the 9 prefix tokens. This is the conditioning on $x$ in $P(y_t \mid y_{<t}, x)$ — without this sub-layer the decoder would be a plain language model that had never seen the prefix.
Two things to be careful about:
- No causal mask here. The decoder is allowed to look at the entire prefix at every step. The prefix is given, not generated, so there is no future to hide. The causal constraint lives only in the decoder's self-attention.
- $K$ and $V$ are computed once and reused at every generation step — only $Q$ changes, because only the decoder side grows. This is why encoder–decoder models cache the encoder output.
[cross-attention] x_q (7, 8) , x_kv (9, 8) Q,K,V per head: (2, 7, 4), (2, 9, 4), (2, 9, 4) scores/weights: (2, 7, 9) concat: (7, 8) out: (7, 8) Q comes from the DECODER : (7, 8) -> per head (2, 7, 4) K,V come from the ENCODER: (9, 8) -> per head (2, 9, 4) Score matrix is rectangular: (2, 7, 9) (heads, n_tgt, n_src) --- Head 0: cross-attention weights (7 decoder x 9 prefix) --- [[0.144 0.088 0.091 0.091 0.068 0.018 0.002 0.448 0.05 ] [0.005 0.004 0.429 0.338 0.004 0. 0.002 0.005 0.212] [0.034 0.039 0.222 0.204 0.044 0.029 0.184 0.019 0.225] [0.004 0.006 0.387 0.279 0.007 0.001 0.041 0.004 0.271] [0.004 0.004 0.362 0.289 0.005 0.001 0.036 0.001 0.298] [0.069 0.074 0.001 0.001 0.07 0.43 0.009 0.345 0.001] [0.035 0.029 0.269 0.255 0.037 0.015 0.077 0.011 0.272]] Each row sums to 1 over the 9 prefix tokens: True
Decoder position 3 ('tiny') attends most to prefix tokens:
'sat' (pos 2): 0.3868
'on' (pos 3): 0.2788
'looking' (pos 8): 0.2714
(With untrained weights this is arbitrary — a trained model would concentrate on the prefix words that actually constrain the next token.)
Cross-Attention Matrices
Here the two stacks meet, and the operands come from different places. This is the clearest place to see why shapes work out:
| operation | shapes | result |
|---|---|---|
| $Q$ from decoder $Y_1$ | $(7,8) @ (8,8)$ | $(7,8) \to$ head $(7,4)$ |
| $K$ from encoder $H$ | $(9,8) @ (8,8)$ | $(9,8) \to$ head $(9,4)$ |
| $Q_0 K_0^\top$ | $(7,4) @ (4,9)$ | $(\mathbf{7},\mathbf{9})$ |
| $A V_0$ | $(7,9) @ (9,4)$ | $(\mathbf{7},4)$ |
The output has 7 rows (one per decoder position) and 4 columns per head — the 9 prefix positions appear only as the summed-over dimension. The decoder gets back exactly its own shape, having absorbed information from a sequence of a different length.
Q source = dec_norm1 (7, 8) (decoder)
K,V source = encoder_output (9, 8) (encoder)
==============================================================================
Q_dec_head0 (7, 4) @ K_enc_head0^T (4, 9) = raw_scores (7, 9)
shape rule: (7 x 4) @ (4 x 9) -> (7 x 9) inner dim 4 cancels
rows indexed by DECODER positions, columns by PREFIX positions
==============================================================================
Q_dec_head0 =
<sos> | -1.283 +0.863 -0.584 +2.871
at | +0.054 -0.599 +1.694 +2.945
a | +0.547 -0.652 +0.976 -0.296
tiny | +1.658 -0.365 +1.646 +1.304
mouse | +1.061 -1.398 +1.859 +0.975
running | -0.135 +1.784 -2.773 -0.212
by | +0.329 -1.214 +0.509 +0.816
K_enc_head0^T =
-1.745 -1.226 +0.050 -0.148 -0.856 -0.898 +1.177 -1.513 +0.188
-0.360 -0.008 -1.309 -1.395 -0.341 -0.161 -1.431 +1.839 -1.672
-0.470 -0.308 +2.101 +1.958 -0.500 -1.599 +0.232 -0.248 +1.714
-0.051 -0.234 +1.244 +1.151 -0.186 -1.401 -1.305 +0.230 +0.913
raw_scores =
<sos> | +2.059 +1.076 +1.152 +1.148 +0.563 -2.075 -6.626 +4.335 -0.063
at | -0.823 -1.272 +8.009 +7.532 -1.236 -6.786 -2.530 -0.925 +6.602
a | -1.164 -0.897 +2.564 +2.399 -0.679 -1.533 +2.188 -2.337 +2.596
tiny | -3.602 -2.843 +5.643 +4.988 -2.360 -5.890 +1.153 -3.289 +4.935
mouse | -2.271 -2.091 +7.002 +6.555 -1.541 -5.066 +2.408 -4.413 +6.614
running | +0.907 +1.056 -8.431 -8.142 +0.932 +4.565 -3.078 +4.125 -7.955
by | -0.418 -0.741 +3.691 +3.581 -0.273 -2.057 +1.178 -2.669 +3.709
worked element raw_scores[0,0] = row 0 of Q_dec_head0 . column 0 of K_enc_head0^T
row 0 of Q_dec_head0 = [-1.283 0.863 -0.584 2.871]
col 0 of K_enc_head0^T = [-1.745 -0.36 -0.47 -0.051]
(-1.2834)*(-1.7453) + (+0.8625)*(-0.3597) + (-0.5840)*(-0.4697) + (+2.8708)*(-0.0505)
= +2.058955 | from array: +2.058955 | MATCH: True
column labels are the prefix tokens: ['the', 'cat', 'sat', 'on', 'the', 'mat', 'and', 'was', 'looking']
==============================================================================
A (7x9) (7, 9) @ V_enc_head0 (9x4) (9, 4) = head0_out (7x4) (7, 4)
shape rule: (7 x 9) @ (9 x 4) -> (7 x 4) inner dim 9 cancels
the 9 PREFIX positions cancel — output keeps the decoder's 7 rows
==============================================================================
A (7x9) =
<sos> | +0.144 +0.088 +0.091 +0.091 +0.068 +0.018 +0.002 +0.448 +0.050
at | +0.005 +0.004 +0.429 +0.338 +0.004 +0.000 +0.002 +0.005 +0.212
a | +0.034 +0.039 +0.222 +0.204 +0.044 +0.029 +0.184 +0.019 +0.225
tiny | +0.004 +0.006 +0.387 +0.279 +0.007 +0.001 +0.041 +0.004 +0.271
mouse | +0.004 +0.004 +0.362 +0.289 +0.005 +0.001 +0.036 +0.001 +0.298
running | +0.069 +0.074 +0.001 +0.001 +0.070 +0.430 +0.009 +0.345 +0.001
by | +0.035 +0.029 +0.269 +0.255 +0.037 +0.015 +0.077 +0.011 +0.272
V_enc_head0 (9x4) =
+0.972 -0.975 +0.443 +0.039
+0.085 -1.478 +1.483 -0.670
-0.539 -0.849 -0.472 +0.348
+0.392 -0.988 -0.585 +0.547
+2.163 -1.204 +0.166 +0.024
+0.894 +0.418 -0.328 -0.174
+1.643 -0.586 -0.894 -0.251
+0.186 +1.391 +1.781 -0.570
-1.243 -0.802 -1.008 +0.702
head0_out (7x4) =
<sos> | +0.322 +0.071 +0.850 -0.194
at | -0.343 -0.879 -0.598 +0.477
a | +0.143 -0.784 -0.510 +0.261
tiny | -0.348 -0.859 -0.637 +0.461
mouse | -0.376 -0.867 -0.663 +0.481
running | +0.688 +0.391 +0.616 -0.318
by | -0.125 -0.843 -0.539 +0.378
worked element head0_out (7x4)[0,0] = row 0 of A (7x9) . column 0 of V_enc_head0 (9x4)
row 0 of A (7x9) = [0.144 0.088 0.091 0.091 0.068 0.018 0.002 0.448 0.05 ]
col 0 of V_enc_head0 (9x4) = [ 0.972 0.085 -0.539 0.392 2.163 0.894 1.643 0.186 -1.243]
(+0.1437)*(+0.9724) + (+0.0879)*(+0.0852) + (+0.0913)*(-0.5387) + (+0.0911)*(+0.3924) +
(+0.0680)*(+2.1627) + (+0.0182)*(+0.8942) + (+0.0019)*(+1.6430) + (+0.4483)*(+0.1859) +
(+0.0497)*(-1.2427)
= +0.321661 | from array: +0.321661 | MATCH: True
>>> (7,9) @ (9,4) = (7,4). Encoder length 9 vanishes into the sum.
>>> This is how a decoder of length 7 reads a source of length 9.
Decoder Add & Norm (2)
$$Y_2 = \mathrm{LayerNorm}\big(Y_1 + \mathrm{CrossAttention}(Y_1, H)\big)$$
Note the residual adds $Y_1$ — the decoder stream — not the encoder output. The encoder information enters only through the attention output. This keeps the decoder's own representation as the backbone.
Decoder Add&Norm 2 (shape (7, 8)): [[ 0.764 1.392 0.269 0.846 -1.436 -0.636 0.269 -1.467] [ 1.128 0.09 1.533 0.505 -1.823 -0.364 -0.328 -0.741] [-1.147 0.963 0.417 -0.026 -1.832 1.408 0.438 -0.221] [-0.447 -0.56 2.124 0.323 -1.345 -0.258 -0.62 0.783] [-0.493 -0.006 2.077 -0.293 -1.413 0.988 -0.514 -0.345] [-1.54 1.181 -0.292 1.414 0.74 -0.902 0.25 -0.851] [-0.974 1.342 1.715 -0.487 -0.987 0.564 -0.8 -0.373]]
14. Decoder: Feed-Forward Network
Identical in form to the encoder FFN, with its own weights:
$$\mathrm{FFN}(x) = \max(0, xW_1 + b_1)W_2 + b_2, \qquad \text{decoder\_output} = \mathrm{LayerNorm}\big(Y_2 + \mathrm{FFN}(Y_2)\big)$$
(7, 8) -> W1 -> (7, 32) -> ReLU -> W2 -> (7, 8) (49.6% of hidden units zeroed by ReLU) DECODER OUTPUT (shape (7, 8)): [[ 0.427 1.034 0.123 1.27 -0.264 -0.185 -0.171 -2.233] [ 0.851 -0.119 1.272 1.394 -0.955 -0.262 -0.678 -1.503] [-1.206 0.158 0.225 0.322 -0.799 2.268 -0.245 -0.722] [-0.107 -1.762 1.586 0.781 -0.874 0.396 -0.657 0.636] [-0.541 -0.844 1.685 0.095 -0.302 1.61 -0.936 -0.766] [-1.514 0.464 -0.833 1.901 0.461 0.283 0.168 -0.93 ] [-0.707 0.527 1.426 -0.643 0.119 1.46 -1.444 -0.738]] mean = +0.0000, std = 1.0000 Magnitudes stay O(1) through the whole network — contrast with ~80 000 before.
15. Output Projection
The decoder produced 7 vectors of dimension 8. To predict tokens they must be projected onto the vocabulary:
$$\text{logits} = \text{decoder\_output} \cdot W_{out} + b_{out}, \qquad W_{out} \in \mathbb{R}^{d_{model} \times V} = \mathbb{R}^{8 \times 500}$$
$$(7 \times 8) \cdot (8 \times 500) = (7 \times 500)$$
Row $t$, column $v$ of the result is the unnormalized score ("logit") that token $v$ is the correct continuation at position $t$.
Weight tying. The paper shares this matrix with the embedding table: $W_{out} = E^\top$. Two reasons —
- Parameter economy: $V \times d_{model} = 4000$ parameters saved here, hundreds of millions in a real model.
- Interpretation: the logit becomes a dot product between the decoder's output vector and the embedding of each candidate token, $\text{logit}_{t,v} = \langle \text{decoder\_output}_t,\; E_v \rangle$. So the model is literally asking "which token's embedding does my prediction vector point at most?" — a nearest-neighbour search in embedding space, measured by inner product.
This is why keeping the random IDs in [50, 500) is fine: the projection covers all 500 vocabulary slots, and the handful that correspond to real words simply have meaningful embeddings while the rest are unused noise.
decoder_output (7, 8) @ W_out (8, 500) = logits (7, 500)
Logits for the last position (t = 6), first 12 vocabulary entries:
[-0.721 -0.798 -0.249 -0.334 0.354 3.275 -0.734 0.531 -0.554 0.007 0.205 -0.104]
range over the whole matrix: [-4.153, 3.611]
Hand check logit[t=6, v=2 ('<eos>')] = <decoder_output[6], E[2]>
= (-0.707)(0.010) + (0.527)(0.452) + (1.426)(0.068) + (-0.643)(0.016) + (0.119)(-0.481) + (1.460)(0.264) + (-1.444)(0.228) + (-0.738)(0.765)
= -0.249453
from logits matrix: -0.249453 match: True
Output Projection Matrix
==============================================================================
decoder_output (7,8) (7, 8) @ W_out = E^T (8,500) (8, 500) = logits (7,500) (7, 500)
shape rule: (7 x 8) @ (8 x 500) -> (7 x 500) inner dim 8 cancels
8 collapses, 500 opens up: one score per vocabulary entry per position
==============================================================================
decoder_output (7,8) =
<sos> | +0.427 +1.034 +0.123 +1.270 -0.264 -0.185 -0.171 -2.233
at | +0.851 -0.119 +1.272 +1.394 -0.955 -0.262 -0.678 -1.503
a | -1.206 +0.158 +0.225 +0.322 -0.799 +2.268 -0.245 -0.722
tiny | -0.107 -1.762 +1.586 +0.781 -0.874 +0.396 -0.657 +0.636
mouse | -0.541 -0.844 +1.685 +0.095 -0.302 +1.610 -0.936 -0.766
running | -1.514 +0.464 -0.833 +1.901 +0.461 +0.283 +0.168 -0.930
by | -0.707 +0.527 +1.426 -0.643 +0.119 +1.460 -1.444 -0.738
W_out = E^T (8,500) = [showing first 10 of 500 columns]
+0.206 +0.489 +0.010 -0.109 -0.158 -0.198 +0.156 +0.230 -0.085 +0.201
-0.127 +0.229 +0.452 +0.077 +0.069 +0.133 -0.178 -0.035 +0.124 +0.158
+0.209 -0.059 +0.068 +0.088 +0.380 +0.553 -0.361 +0.653 -0.442 +0.227
+0.392 +0.052 +0.016 +0.558 -0.363 -0.023 +0.250 -0.378 +0.510 +0.470
+0.290 +0.427 -0.481 -0.034 +0.047 -0.196 +0.086 -0.539 -0.029 +0.069
+0.179 -0.289 +0.264 +0.099 -0.248 +0.665 -0.199 -0.245 +0.395 +0.251
+0.377 +0.130 +0.228 +0.215 +0.423 -0.512 -0.453 -0.016 +0.121 -0.032
+0.413 -0.139 +0.765 +0.066 -0.539 -0.777 +0.308 +0.086 +0.161 +0.509
logits (7,500) = [showing first 10 of 500 columns]
<sos> | -0.617 +0.734 -1.168 +0.559 +0.753 +1.844 -0.441 -0.340 +0.239 -0.322
at | -0.198 +0.175 -0.851 +0.549 +0.378 +2.016 -0.144 +0.964 -0.340 +0.221
a | -0.312 -1.477 +0.454 +0.493 -0.144 +2.729 -0.849 -0.440 +0.960 +0.137
tiny | +0.672 -1.171 +0.183 +0.420 -0.545 +0.923 +0.258 +1.217 -0.308 +0.810
mouse | -0.084 -1.162 -0.501 +0.112 +0.236 +3.130 -0.678 +0.687 -0.348 +0.207
running | +0.063 -0.220 -0.650 +1.175 -0.212 +0.591 +0.079 -2.028 +1.494 +0.097
by | -0.721 -0.798 -0.249 -0.334 +0.354 +3.275 -0.734 +0.531 -0.554 +0.007
worked element logits (7,500)[0,0] = row 0 of decoder_output (7,8) . column 0 of W_out = E^T (8,500)
row 0 of decoder_output (7,8) = [ 0.427 1.034 0.123 1.27 -0.264 -0.185 -0.171 -2.233]
col 0 of W_out = E^T (8,500) = [ 0.206 -0.127 0.209 0.392 0.29 0.179 0.377 0.413]
(+0.4266)*(+0.2064) + (+1.0344)*(-0.1270) + (+0.1228)*(+0.2088) + (+1.2699)*(+0.3920) +
(-0.2637)*(+0.2901) + (-0.1854)*(+0.1793) + (-0.1713)*(+0.3771) + (-2.2332)*(+0.4134)
= -0.617476 | from array: -0.617476 | MATCH: True
>>> the only step where the width EXPANDS: 8 -> 500.
>>> 7 * 8 * 500 = 28,000 multiply-adds, more than the rest of the decoder combined.
>>> then softmax turns each of the 7 rows into a distribution over 500 tokens.
16. Softmax: Next-Token Distribution
$$P(y_t = v \mid y_{<t},\, x) = \frac{\exp(\text{logit}_{t,v})}{\sum_{v'=1}^{V} \exp(\text{logit}_{t,v'})}$$
Each of the 7 rows becomes a probability distribution over all 500 tokens, summing to 1. This is the autoregressive factorization made concrete: row $t$ is one factor of the product $\prod_t P(y_t \mid y_{<t}, x)$.
The prediction of the next token is the last row. Position $t = 6$ has seen the prefix (through cross-attention) plus <sos> at a tiny mouse running by — the complete continuation so far — so its distribution answers "what comes after by?" (the correct answer being <eos>, since the sentence ends there).
Rows $0 \dots 5$ are predictions too — they are what the loss is computed on during training — but only the last one is used when generating.
probs shape (7, 500), every row sums to 1: True Prefix (encoder): 'the cat sat on the mat and was looking' Context (decoder): '<sos> at a tiny mouse running by' Question: what token comes next? rank id token logit probability -------------------------------------------------- 1 5 <unused_5> 3.2750 0.033622 2 342 <unused_342> 2.9872 0.025214 3 190 <unused_190> 2.7904 0.020711 4 487 <unused_487> 2.5049 0.015566 5 208 <unused_208> 2.4846 0.015253 PREDICTED NEXT TOKEN: id 5 -> '<unused_5>' True next token : id 2 -> '<eos>' Probability assigned to the true token: 0.000991 Uniform baseline (1/V) : 0.002000 With random weights the prediction is essentially chance — as expected. The arithmetic is correct; the model is simply untrained.
max probability = 0.03362 at id 5 entropy = 5.6914 nats (maximum possible = log(500) = 6.2146)
17. Decoding Strategies
Turning the distribution into a token is a separate decision, not part of the model.
Greedy — take the mode: $$\hat{v} = \arg\max_v P(v)$$ Deterministic, but tends to produce repetitive text.
Temperature $T$ — rescale the logits before the softmax: $$P_T(v) = \frac{\exp(\text{logit}_v / T)}{\sum_{v'} \exp(\text{logit}_{v'} / T)}$$ * $T \to 0$: distribution collapses onto the argmax (greedy). * $T = 1$: unchanged. * $T > 1$: flatter, more random.
Temperature is a monotone transformation of the logits, so it never changes the ranking — only how sharply the mass concentrates. Entropy is the quantity to watch.
Top-$k$ — keep the $k$ highest logits, set the rest to $-\infty$, renormalize, sample. This truncates the long tail of 495 implausible tokens that collectively can still hold non-trivial probability mass.
T entropy (nats) max prob argmax id ------------------------------------------------ 0.1 0.2585 0.93886 5 0.5 4.0472 0.19117 5 1.0 5.6914 0.03362 5 2.0 6.0906 0.00931 5 5.0 6.1952 0.00381 5 The argmax is identical for every T — temperature reshapes, never reorders. Top-5 renormalized distribution: id 5 '<unused>' : 0.033622 -> 0.3046 id 342 '<unused>' : 0.025214 -> 0.2285 id 190 '<unused>' : 0.020711 -> 0.1877 id 487 '<unused>' : 0.015566 -> 0.1410 id 208 '<unused>' : 0.015253 -> 0.1382 (the other 495 tokens: probability exactly 0) 8 samples from the top-5 distribution: [208 487 342 5 5 487 5 342]
18. Cross-Entropy Loss
Prediction and training are two sides of the same forward pass. The loss at position $t$ is the negative log-probability the model assigned to the correct token:
$$\mathcal{L}_t = -\log P_t(\text{label}_t), \qquad \mathcal{L} = \frac{1}{n_{tgt}} \sum_{t=0}^{n_{tgt}-1} \mathcal{L}_t$$
Two reference values are useful for sanity-checking any implementation:
- An untrained model should score $\mathcal{L} \approx \log V = \log 500 \approx 6.21$, i.e. chance level. In practice it lands slightly above $\log V$: the random distribution is not uniform, so it is confidently wrong on some positions, and $-\log p$ punishes that asymmetrically. Anything in the range $6$–$8$ here is normal; a value near $0$ would mean the labels are leaking into the input.
- Perplexity $= e^{\mathcal{L}}$ is the effective number of tokens the model is choosing between. At initialization it should be $\approx V$.
If a fresh implementation starts far from $\log V$, something is wrong (usually a mask leak or a label misalignment). All 7 positions contribute — this is why teacher forcing is efficient: one forward pass yields 7 training signals, one per factor of the autoregressive product.
t input label P(label) -log P ---------------------------------------------------- 0 <sos> at 0.002379 6.0412 1 at a 0.000329 8.0187 2 a tiny 0.001004 6.9041 3 tiny mouse 0.001992 6.2188 4 mouse running 0.000786 7.1490 5 running by 0.000530 7.5420 6 by <eos> 0.000991 6.9170 ---------------------------------------------------- Mean cross-entropy loss : 6.9701 log(V) = log(500) : 6.2146 <- untrained reference Perplexity exp(loss) : 1064.3 (vs V = 500) The loss is at chance level (slightly above log V, as expected for a random non-uniform distribution). Training would drive it toward 0; backpropagation is out of scope for this notebook.
19. Autoregressive Generation
Everything above computed all 7 positions at once, using the true continuation as decoder input (teacher forcing). At inference time the true continuation does not exist, so the decoder must be run repeatedly, feeding back its own predictions. This is the loop that the word "autoregressive" actually names:
H = Encoder(prefix) # computed ONCE
y = [<sos>]
repeat:
logits = Decoder(y, H) # full re-run over the prefix of y
next = argmax(softmax(logits[-1]))
y.append(next)
until next == <eos> or len(y) > max_len
Three points that are easy to miss:
- The encoder runs once. $H$, and therefore cross-attention's $K$ and $V$, never change. Only the decoder is re-executed. Every generated token is conditioned on the same 9-token prefix.
- Step $t$ costs a full decoder pass over $t{+}1$ positions. Positions $0 \dots t-1$ recompute identical values, because the causal mask guarantees they cannot depend on the newly appended token. Real implementations exploit this with a KV cache, storing $K$ and $V$ per position; the naive loop below recomputes them so the mechanism stays visible.
- The causal mask is still required even though nothing to the right exists yet — it is what makes point 2 true, and it keeps training and inference behaviour identical.
The two functions below wrap the exact same weights and operations used above.
encoder_forward matches encoder_output : True decoder_forward matches logits : True
PREFIX GIVEN TO ENCODER: the cat sat on the mat and was looking step context length predicted id P entropy ------------------------------------------------------------------ 0 1 <unused_479> 479 0.02287 5.7252 1 2 <unused_479> 479 0.03229 5.6743 2 3 <unused_314> 314 0.03234 5.6773 3 4 <unused_314> 314 0.03286 5.6893 4 5 <unused_314> 314 0.03137 5.6932 5 6 <unused_314> 314 0.03170 5.6736 6 7 <unused_314> 314 0.03167 5.6559 7 8 <unused_314> 314 0.03266 5.6519 GENERATED: <sos> <unused_479> <unused_479> <unused_314> <unused_314> <unused_314> <unused_314> <unused_314> <unused_314> TRUE CONTINUATION: <sos> at a tiny mouse running by <eos> The output is gibberish because the weights are random. What the loop demonstrates is the mechanism: encoder once, decoder per step, each step conditioned on all tokens generated so far.
T=0.7 -> <sos> <42> <130> <375> <233> <21> <179> <194> <71> T=1.0 -> <sos> <43> <102> <388> <253> <45> <197> <206> <72> T=1.5 -> <sos> <43> <110> <394> looking <46> <205> <243> <84> Greedy is deterministic; sampling varies with T and the seed. The forward pass is identical in all cases.
20. Summary
| Stage | Operation | Output shape |
|---|---|---|
| Prefix tokens | split | 9 |
| Prefix IDs | vocabulary lookup | $(9,)$ |
| Prefix embeddings | $E[\text{ids}] \cdot \sqrt{d_{model}}$ | $(9, 8)$ |
| $+$ positional encoding | $X_{enc}$ | $(9, 8)$ |
| Encoder self-attn (no mask) | $Q,K,V$ per head | $(2, 9, 4)$ |
| scores / weights | $(2, 9, 9)$ | |
| concat $\cdot W^O$ | $(9, 8)$ | |
| Add & Norm | $(9, 8)$ | |
| Encoder FFN | $8 \to 32 \to 8$ | $(9, 8)$ |
| Add & Norm | encoder output $H$ | $(9, 8)$ |
| Decoder input | <sos> + continuation, shifted |
$(7,)$ |
| $+$ positional encoding | $X_{dec}$ | $(7, 8)$ |
| Masked self-attn (causal) | scores / weights | $(2, 7, 7)$ |
| Add & Norm | $(7, 8)$ | |
| Cross-attn | $Q$ from decoder, $K,V$ from $H$ | scores $(2, 7, 9)$ |
| Add & Norm | $(7, 8)$ | |
| Decoder FFN | $8 \to 32 \to 8$ | $(7, 8)$ |
| Add & Norm | decoder output | $(7, 8)$ |
| Output projection | $\cdot\, W_{out} = E^\top$ | $(7, 500)$ |
| Softmax | per row over vocabulary | $(7, 500)$ |
| Next token | argmax of the last row | scalar |
Where Autoregression Happens
| Requirement | Enforced by | Section |
|---|---|---|
| $y_t$ predicted from $y_{<t}$ only | shift right (<sos> prepended) |
11 |
| no attention to the future | causal mask on decoder self-attention | 12 |
| conditioning on the prefix $x$ | cross-attention, $Q$ decoder / $K,V$ encoder | 13 |
| one factor $P(y_t \mid y_{<t}, x)$ per row | softmax over the vocabulary | 16 |
| generation = repeated conditioning | the loop, decoder re-run per step | 19 |
Note that the encoder is deliberately not autoregressive: it is bidirectional over the prefix, because the prefix is given input rather than generated output.
Parameter Count
| Component | Count |
|---|---|
| Embedding table $E$ (tied with $W_{out}$) | $500 \times 8 = 4000$ |
| Encoder self-attention ($W^Q, W^K, W^V, W^O$) | $4 \times 64 = 256$ |
| Encoder FFN | $8{\cdot}32 + 32 + 32{\cdot}8 + 8 = 552$ |
| Decoder self-attention | 256 |
| Cross-attention | 256 |
| Decoder FFN | 552 |
| LayerNorm ($\gamma, \beta$) $\times 5$ | $5 \times 16 = 80$ |
| Total | ≈ 5952 |
What's Left Out
- Backpropagation. Every weight here is random and stays random; the loss is computed but never minimized.
- Stacking. The paper uses $N = 6$ encoder and 6 decoder blocks. The block above is one of them; stacking means feeding
encoder_outputinto an identical block with fresh weights. - Dropout, learning-rate warmup, label smoothing — training details, not forward-pass mechanics.
- Padding masks and batching. With a single unpadded sequence, no
<pad>mask is needed; a real batch requires masking padded key positions in both self- and cross-attention. - KV caching during generation, as noted in section 19.
---
Appendix: Every Number, By Hand
Everything above showed matrices. This appendix takes one token and follows it through every single arithmetic operation of the network, writing out each multiply-and-add explicitly, and checking each result against the arrays computed earlier.
Two traces:
- Trace A — encoder token 1,
"cat", from embedding lookup toencoder_output[1]. - Trace B — decoder position 6,
"by", from embedding lookup to the final probability of the next token.
Every line printed below is verified with np.isclose against the corresponding entry of the matrices computed in the main body. If a line says MATCH: True, the hand calculation and the vectorized code agree exactly.
Helpers defined. Every step below prints its terms, its result, and a match check.
Trace A: Encoder Token 'cat'
A.1 Embedding Lookup and Scaling
$$\text{emb} = E[\,\text{id}\,] \cdot \sqrt{d_{model}} = E[\,\text{id}\,] \cdot 2.828427$$
Component by component.
Token 'cat' at position 1, vocabulary id = 168 dim E[168] x sqrt(8) = emb ------------------------------------------------ 0 -0.242053 2.828427 -0.684630 1 0.279528 2.828427 0.790626 2 -0.696537 2.828427 -1.970104 3 0.315581 2.828427 0.892597 4 -0.428214 2.828427 -1.211172 5 0.258364 2.828427 0.730764 6 0.005046 2.828427 0.014273 7 -0.337269 2.828427 -0.953939 matches src_emb[1]: True
A.2 Positional Encoding
$$X_{enc}[1, d] = \text{emb}[d] + PE[1, d]$$
with, for even $d = 2i$: $PE[1, 2i] = \sin(1 \cdot \omega_i)$ and for odd $d = 2i+1$: $PE[1, 2i+1] = \cos(1 \cdot \omega_i)$, where $\omega_i = 10000^{-2i/8}$.
dim i fn omega_i PE[1,d] emb X_enc[1,d] -------------------------------------------------------------------- 0 0 sin 1.000000 0.841471 -0.684630 0.156841 1 0 cos 1.000000 0.540302 0.790626 1.330928 2 1 sin 0.100000 0.099833 -1.970104 -1.870270 3 1 cos 0.100000 0.995004 0.892597 1.887602 4 2 sin 0.010000 0.010000 -1.211172 -1.201172 5 2 cos 0.010000 0.999950 0.730764 1.730714 6 3 sin 0.001000 0.001000 0.014273 0.015273 7 3 cos 0.001000 1.000000 -0.953939 0.046060 matches X_enc[1]: True x = X_enc[1] = [ 0.157 1.331 -1.87 1.888 -1.201 1.731 0.015 0.046]
A.3 Projection to Q, K, V
$$Q[1, c] = \sum_{d=0}^{7} X_{enc}[1, d] \cdot W^Q[d, c]$$
One output component is a sum of 8 products. Below, component $c = 0$ is expanded in full; the remaining 7 are tabulated.
Recall the head split: columns 0–3 of $Q$ become head 0, columns 4–7 become head 1.
Q[1, 0] = x . Wq[:, 0]
(+0.1568)*(-0.3560) + (+1.3309)*(-0.2957) + (-1.8703)*(+0.3342) + (+1.8876)*(-0.3792) +
(-1.2012)*(+0.0962) + (+1.7307)*(+0.4493) + (+0.0153)*(-0.4081) + (+0.0461)*(-0.1221)
= -1.140044 | from array: -1.140044 | MATCH: True
All 8 components of Q[1] (full 8-dim, before head split):
c Q[1,c] head
------------------------------
0 -1.140044 head 0
1 -1.111687 head 0
2 0.583413 head 0
3 -0.526634 head 0
4 -0.219950 head 1
5 1.319576 head 1
6 0.033149 head 1
7 -0.870463 head 1
head 0 slice = columns 0..3 = [-1.14 -1.112 0.583 -0.527]
matches Qe[0, 1] = [-1.14 -1.112 0.583 -0.527] MATCH: True
head 1 slice = columns 4..7 = [-0.22 1.32 0.033 -0.87 ]
matches Qe[1, 1] = [-0.22 1.32 0.033 -0.87 ] MATCH: True
A.4 Attention Scores
$$s_j = \frac{q_1 \cdot k_j}{\sqrt{d_k}} = \frac{q_1 \cdot k_j}{2}, \qquad j = 0 \dots 8$$
Nine dot products of 4-dimensional vectors (head 0).
q_1 (head 0) = [-1.14 -1.112 0.583 -0.527]
j token q.k_j /2 = score array match
--------------------------------------------------------------------
0 the -0.563007 -0.281504 -0.281504 True
1 cat -0.913154 -0.456577 -0.456577 True
2 sat 0.045592 0.022796 0.022796 True
3 on 1.131739 0.565869 0.565869 True
4 the 0.584410 0.292205 0.292205 True
5 mat -0.388204 -0.194102 -0.194102 True
6 and 1.027013 0.513506 0.513506 True
7 was -1.268738 -0.634369 -0.634369 True
8 looking 2.864096 1.432048 1.432048 True
Expansion of j = 8 ('looking'):
(-1.1400)*(+0.1606) + (-1.1117)*(-1.9641) + (+0.5834)*(+2.4849) + (-0.5266)*(+1.1127)
= +2.864096, divided by sqrt(4) = 2 -> +1.432048
A.5 Softmax
$$A_{1j} = \frac{e^{s_j - m}}{\sum_{j'} e^{s_{j'} - m}}, \qquad m = \max_j s_j$$
The shift by $m$ is the numerical-stability trick: it cancels in the ratio, since $e^{s_j - m} / \sum e^{s_{j'} - m} = e^{s_j} / \sum e^{s_{j'}}$.
row max m = 1.432048 (at position 8, 'looking') j token s_j s_j - m exp /Z = A_1j array ---------------------------------------------------------------------------- 0 the -0.281504 -1.713551 0.180225 0.059310 0.059310 1 cat -0.456577 -1.888625 0.151280 0.049785 0.049785 2 sat 0.022796 -1.409252 0.244326 0.080405 0.080405 3 on 0.565869 -0.866179 0.420556 0.138401 0.138401 4 the 0.292205 -1.139843 0.319869 0.105266 0.105266 5 mat -0.194102 -1.626150 0.196685 0.064727 0.064727 6 and 0.513506 -0.918541 0.399101 0.131340 0.131340 7 was -0.634369 -2.066417 0.126639 0.041676 0.041676 8 looking 1.432048 0.000000 1.000000 0.329090 0.329090 ---------------------------------------------------------------------------- denominator Z = 0.1802 + 0.1513 + 0.2443 + 0.4206 + 0.3199 + 0.1967 + 0.3991 + 0.1266 + 1.0000 = 3.038680 sum of weights = 1.0000000000 matches array row: True
A.6 Attention Output
$$\text{head}_0[1] = \sum_{j=0}^{8} A_{1j} \, v_j$$
This is the operation attention actually performs: a convex combination of the 9 value vectors, with the softmax weights as coefficients. Component 0 is expanded term by term.
Component 0: sum_j A_1j * V[j, 0]
(+0.0593)*(-1.3749) + (+0.0498)*(-2.8072) + (+0.0804)*(+0.7433) + (+0.1384)*(+1.6414) +
(+0.1053)*(-1.1323) + (+0.0647)*(-3.0817) + (+0.1313)*(-2.6194) + (+0.0417)*(-2.3888) +
(+0.3291)*(-1.0528)
= -1.043069 | from array: -1.043069 | MATCH: True
All 4 components of head 0 output for token 1:
dim 0: -1.043069
dim 1: -0.544660
dim 2: -0.127629
dim 3: -0.076856
head 0 = [-1.043 -0.545 -0.128 -0.077]
head 1 = [1.377 0.827 0.138 0.455]
concat = [-1.043 -0.545 -0.128 -0.077 1.377 0.827 0.138 0.455]
matches concat_e[1]: True
Sanity: weights are non-negative and sum to 1, so every component of the output
lies between min and max of the corresponding V column:
V[:,0] range = [-3.0817, +1.6414], output[0] = -1.0431
A.7 Output Projection
$$\text{attn\_out}[1, c] = \sum_{d=0}^{7} \text{concat}[1, d] \cdot W^O[d, c]$$
This is the step the old notebook omitted entirely. Without it the two heads are simply glued together and never mixed.
attn_out[1, 0] = concat . Wo[:, 0]
(-1.0431)*(-0.5044) + (-0.5447)*(-0.0727) + (-0.1276)*(-0.2093) + (-0.0769)*(-0.3685) +
(+1.3767)*(-0.5410) + (+0.8272)*(+0.2712) + (+0.1382)*(-0.4579) + (+0.4548)*(+0.5742)
= +0.298177 | from array: +0.298177 | MATCH: True
full row: [ 0.298 -0.049 -0.437 -1.226 0.878 0.768 -0.166 -0.353]
array : [ 0.298 -0.049 -0.437 -1.226 0.878 0.768 -0.166 -0.353]
MATCH : True
A.8 Residual Addition
$$r = x + \text{attn\_out}$$
dim x + attn_out = r ---------------------------------------------- 0 0.156841 0.298177 0.455018 1 1.330928 -0.049324 1.281604 2 -1.870270 -0.436947 -2.307218 3 1.887602 -1.225554 0.662048 4 -1.201172 0.877659 -0.323513 5 1.730714 0.767867 2.498581 6 0.015273 -0.165709 -0.150436 7 0.046060 -0.352609 -0.306549 matches enc_res1[1]: True
A.9 Layer Normalization
$$\mu = \frac{1}{8}\sum_{d} r_d, \qquad \sigma^2 = \frac{1}{8}\sum_{d} (r_d - \mu)^2, \qquad \mathrm{LN}(r)_d = \gamma_d \frac{r_d - \mu}{\sqrt{\sigma^2 + 10^{-5}}} + \beta_d$$
With $\gamma = \mathbf{1}$, $\beta = \mathbf{0}$ this is plain standardization of the 8 numbers of this one token.
mu = (+0.4550 + +1.2816 + -2.3072 + +0.6620 + -0.3235 + +2.4986 + -0.1504 + -0.3065) / 8
= +1.809535 / 8 = +0.226192
var = ((+0.2288)^2 + (+1.0554)^2 + (-2.5334)^2 + (+0.4359)^2 + (-0.5497)^2 + (+2.2724)^2 + (-0.3766)^2 + (-0.5327)^2) / 8
= +13.665984 / 8 = 1.708248
denominator sqrt(var + 1e-5) = 1.307003
dim r_d r_d - mu / den array
------------------------------------------------------------
0 0.455018 0.228826 0.175077 0.175077
1 1.281604 1.055412 0.807505 0.807505
2 -2.307218 -2.533410 -1.938334 -1.938334
3 0.662048 0.435856 0.333477 0.333477
4 -0.323513 -0.549705 -0.420584 -0.420584
5 2.498581 2.272390 1.738626 1.738626
6 -0.150436 -0.376628 -0.288162 -0.288162
7 -0.306549 -0.532741 -0.407605 -0.407605
MATCH: True
check: mean of result = +2.78e-17, std = 0.999997
A.10 Feed-Forward Network
$$z_u = \sum_{d=0}^{7} \text{LN}[1,d]\, W_1[d,u] + b_{1,u}, \qquad a_u = \max(0, z_u), \qquad \text{ffn}_c = \sum_{u=0}^{31} a_u W_2[u,c] + b_{2,c}$$
The first layer produces 32 hidden units from 8 inputs; two of them are expanded below (one that survives ReLU, one that is killed). The second layer contracts 32 back to 8 — that sum has 32 terms.
hidden unit u = 8 (survives)
(+0.1751)*(+0.3476) + (+0.8075)*(+0.3610) + (-1.9383)*(-0.3681) + (+0.3335)*(+0.1975) +
(-0.4206)*(-0.0357) + (+1.7386)*(-0.0116) + (-0.2882)*(-0.2019) + (-0.4076)*(-0.0986)
+ b1[8] = +0.0000
z_8 = +1.224936 -> ReLU -> a_8 = +1.224936
array: pre-ReLU +1.224936, post +1.224936 MATCH: True
hidden unit u = 24 (killed by ReLU)
(+0.1751)*(+0.0218) + (+0.8075)*(-0.3468) + (-1.9383)*(+0.3620) + (+0.3335)*(-0.3270) +
(-0.4206)*(-0.2873) + (+1.7386)*(-0.3211) + (-0.2882)*(-0.1196) + (-0.4076)*(-0.0938)
+ b1[24] = +0.0000
z_24 = -1.451721 -> ReLU -> a_24 = +0.000000
array: pre-ReLU -1.451721, post +0.000000 MATCH: True
ReLU killed 15 of 32 units for this token.
Second layer, output component 0 = sum over 32 hidden units:
(only the 17 non-zero units contribute; the rest add 0)
(+0.7532)*(-0.2086) + (+0.2872)*(-0.2626) + (+0.8079)*(-0.1647) + (+0.0870)*(+0.3867) +
(+1.2249)*(-0.0544) + (+0.6054)*(+0.2524) + (+0.0026)*(-0.0177) + (+0.7401)*(+0.3343) +
(+0.4529)*(-0.2638) + (+0.4463)*(+0.0024) + (+1.1784)*(-0.2271) + (+0.4412)*(-0.0650) +
(+0.6174)*(+0.0893) + (+1.2025)*(-0.0751) + (+0.1314)*(+0.1340) + (+0.4955)*(+0.1958) +
(+0.3005)*(+0.1570)
+ b2[0] = +0.0000
= -0.286386 | from array: -0.286386 | MATCH: True
A.11 Final Residual + Norm
$$H[1] = \mathrm{LayerNorm}\big(\text{LN}[1] + \text{ffn}[1]\big)$$
End of Trace A: this row is what cross-attention will later read as a key and a value.
dim LN1 + ffn = r2 -> LN2 array -------------------------------------------------------------------------- 0 0.175077 -0.286386 -0.111309 0.268271 0.268271 1 0.807505 -0.415618 0.391887 0.905322 0.905322 2 -1.938334 0.453973 -1.484362 -1.470026 -1.470026 3 0.333477 -1.062228 -0.728750 -0.513415 -0.513415 4 -0.420584 0.248086 -0.172498 0.190806 0.190806 5 1.738626 -1.019688 0.718937 1.319370 1.319370 6 -0.288162 0.599531 0.311369 0.803385 0.803385 7 -0.407605 -1.103366 -1.510971 -1.503713 -1.503713 mu = -0.323212, var = 0.623907 MATCH encoder_output[1]: True === TRACE A COMPLETE: 'cat' -> [ 0.268 0.905 -1.47 -0.513 0.191 1.319 0.803 -1.504] ===
Trace B: Decoder Position 6 ('by')
This is the position whose output becomes the actual prediction. It carries the whole autoregressive condition: <sos> at a tiny mouse running by on the decoder side, and the 9-token prefix through cross-attention.
B.1 Masked Self-Attention
Position 6 is the last row, so all 7 positions are visible — the mask blocks nothing here. (Compare with position 0, shown afterwards, where 6 of the 7 are blocked.)
Decoder position 6 = 'by'
q (head 0) = [ 0.726 -0.334 1.135 0.84 ]
j token masked? score exp weight array
----------------------------------------------------------------------------
0 <sos> no 0.258070 0.610869 0.200936 0.200936
1 at no 0.750942 1.000000 0.328935 0.328935
2 a no -0.415989 0.311321 0.102404 0.102404
3 tiny no -0.867940 0.198120 0.065169 0.065169
4 mouse no -0.224371 0.377074 0.124033 0.124033
5 running no -1.020254 0.170129 0.055961 0.055961
6 by no -0.236296 0.372604 0.122562 0.122562
----------------------------------------------------------------------------
Z = 3.040119, weights sum = 1.0000000000, MATCH: True
Contrast, position 0 ('<sos>'): [1. 0. 0. 0. 0. 0. 0.]
-> 6 of 7 scores were set to -1e9, exp(-1e9 - max) underflows to exactly 0,
-> so the single visible position takes the whole mass: weight = 1.0
B.2 Output Projection $W^O$
$$\text{attn\_out}[6, c] = \sum_{d=0}^{7} \text{concat}[6, d] \cdot W^O[d, c]$$
The two heads are concatenated into 8 numbers, then mixed by $W^O$. Without this projection the heads would remain two independent 4-dimensional streams.
head 0 output at position 6 = [ 1.023 0.352 1.044 -1.221]
head 1 output at position 6 = [-0.754 0.231 0.862 -0.728]
concat (8 numbers) = [ 1.023 0.352 1.044 -1.221 -0.754 0.231 0.862 -0.728]
matches concat_d[6]: True
dec_attn_out[6, 0] = concat . W_O[:, 0]
(+1.0226)*(-0.2273) + (+0.3517)*(+0.4406) + (+1.0439)*(-0.4938) + (-1.2206)*(+0.3430) +
(-0.7539)*(+0.0734) + (+0.2307)*(+0.5296) + (+0.8617)*(-0.3028) + (-0.7277)*(-0.4318)
= -0.891437 | from array: -0.891437 | MATCH: True
full row: [-0.891 -0.412 0.997 -0.193 -1.074 -0.265 -0.311 -0.824]
array : [-0.891 -0.412 0.997 -0.193 -1.074 -0.265 -0.311 -0.824]
MATCH : True
B.3 Residual + LayerNorm (Add & Norm 1)
$$r = X_{dec}[6] + \text{attn\_out}[6], \qquad Y_1[6] = \mathrm{LayerNorm}(r)$$
Same arithmetic as A.8 and A.9, now on the decoder side.
dim X_dec + attn_out = r
----------------------------------------------
0 -0.122722 -0.891437 -1.014158
1 1.983701 -0.411742 1.571959
2 0.739929 0.997242 1.737171
3 -0.511390 -0.192861 -0.704251
4 -0.351859 -1.073738 -1.425598
5 1.130129 -0.265326 0.864802
6 -0.448548 -0.311432 -0.759979
7 0.781329 -0.824303 -0.042975
mu = (-1.0142 + +1.5720 + +1.7372 + -0.7043 + -1.4256 + +0.8648 + -0.7600 + -0.0430) / 8
= +0.226970 / 8 = +0.028371
var = 1.295811, sqrt(var + 1e-5) = 1.138342
dim r_d - mu / den array
------------------------------------------------
0 -1.042530 -0.915832 -0.915832
1 1.543587 1.355997 1.355997
2 1.708799 1.501131 1.501131
3 -0.732622 -0.643588 -0.643588
4 -1.453969 -1.277270 -1.277270
5 0.836431 0.734780 0.734780
6 -0.788351 -0.692543 -0.692543
7 -0.071346 -0.062675 -0.062675
MATCH dec_norm1[6]: True
B.4 Cross-Attention
$$s_j = \frac{q^{cross}_6 \cdot k^{cross}_j}{2}, \qquad j = 0 \dots 8 \;\; (\textbf{prefix positions})$$
Note what changes here: $q$ comes from the decoder, but $j$ now indexes the 9 prefix tokens, because $K$ and $V$ were built from the encoder output $H$. This is the only arithmetic in the whole network that mixes the two stacks.
q_cross (head 0, decoder position 6) = [ 0.329 -1.214 0.509 0.816]
K_cross rows come from H = encoder_output (9 prefix tokens)
j prefix token score exp weight array
------------------------------------------------------------------
0 the -0.208753 0.127016 0.034525 0.034525
1 cat -0.370644 0.108031 0.029364 0.029364
2 sat 1.845304 0.990657 0.269274 0.269274
3 on 1.790264 0.937605 0.254854 0.254854
4 the -0.136596 0.136520 0.037108 0.037108
5 mat -1.028494 0.055956 0.015210 0.015210
6 and 0.588864 0.282006 0.076653 0.076653
7 was -1.334534 0.041204 0.011200 0.011200
8 looking 1.854691 1.000000 0.271813 0.271813
------------------------------------------------------------------
MATCH: True, sum = 1.0000000000
Expansion of the largest score, j = 8 ('looking'):
(+0.3289)*(+0.1877) + (-1.2142)*(-1.6720) + (+0.5089)*(+1.7143) + (+0.8162)*(+0.9130)
= +3.709381 / 2 = +1.854691
cross-attention head-0 output = sum_j weight_j * V_j = [-0.125 -0.843 -0.539 0.378]
component 0: (+0.0345)*(+0.9724) + (+0.0294)*(+0.0852) + (+0.2693)*(-0.5387) + (+0.2549)*(+0.3924) +
(+0.0371)*(+2.1627) + (+0.0152)*(+0.8942) + (+0.0767)*(+1.6430) + (+0.0112)*(+0.1859) +
(+0.2718)*(-1.2427)
= -0.124896 | from array: -0.124896 | MATCH: True
B.5 Cross-Attention $W^O$ + Add & Norm 2
$$Y_2[6] = \mathrm{LayerNorm}\big(Y_1[6] + \mathrm{CrossAttention}(Y_1, H)[6]\big)$$
The cross-attention sub-layer also has its own output projection; both heads are mixed before the residual is added.
cross-attention concat at position 6 = [-0.125 -0.843 -0.539 0.378 -0.512 0.096 0.056 0.095]
cross_out[6, 0] = concat . W_O[:, 0]
(-0.1249)*(-0.0181) + (-0.8432)*(-0.1343) + (-0.5387)*(-0.2618) + (+0.3780)*(-0.2699) +
(-0.5121)*(+0.0703) + (+0.0958)*(+0.4548) + (+0.0564)*(+0.5068) + (+0.0951)*(+0.3894)
= +0.227632 | from array: +0.227632 | MATCH: True
full row: [ 0.228 0.083 0.281 0.403 0.577 -0.01 0.165 -0.073]
array : [ 0.228 0.083 0.281 0.403 0.577 -0.01 0.165 -0.073]
MATCH : True
Add & Norm 2:
dim Y1 + cross = r -> LN array
--------------------------------------------------------------------------
0 -0.915832 0.227632 -0.688200 -0.974226 -0.974226
1 1.355997 0.082857 1.438854 1.341625 1.341625
2 1.501131 0.280904 1.782034 1.715267 1.715267
3 -0.643588 0.402607 -0.240981 -0.487312 -0.487312
4 -1.277270 0.577426 -0.699844 -0.986904 -0.986904
5 0.734780 -0.009865 0.724915 0.564317 0.564317
6 -0.692543 0.164553 -0.527990 -0.799796 -0.799796
7 -0.062675 -0.073286 -0.135962 -0.372971 -0.372971
mu = +0.206603, var = 0.843587
MATCH dec_norm2[6]: True
B.6 Feed-Forward Network
$$z_u = \sum_{d=0}^{7} Y_2[6,d]\, W_1[d,u] + b_{1,u}, \qquad a_u = \max(0, z_u), \qquad \text{ffn}_c = \sum_{u=0}^{31} a_u W_2[u,c] + b_{2,c}$$
Identical in form to A.10, with the decoder's own weights.
hidden unit u = 8 (survives)
(-0.9742)*(-0.0613) + (+1.3416)*(+0.2223) + (+1.7153)*(+0.3452) + (-0.4873)*(-0.3863) +
(-0.9869)*(-0.1973) + (+0.5643)*(+0.0113) + (-0.7998)*(-0.3303) + (-0.3730)*(-0.3614)
+ b1[8] = +0.0000
z_8 = +1.738386 -> ReLU -> a_8 = +1.738386
array: pre-ReLU +1.738386, post +1.738386 MATCH: True
hidden unit u = 27 (killed by ReLU)
(-0.9742)*(+0.3458) + (+1.3416)*(-0.1154) + (+1.7153)*(-0.3581) + (-0.4873)*(-0.2504) +
(-0.9869)*(+0.3746) + (+0.5643)*(-0.1724) + (-0.7998)*(+0.1468) + (-0.3730)*(+0.0328)
+ b1[27] = +0.0000
z_27 = -1.580555 -> ReLU -> a_27 = +0.000000
array: pre-ReLU -1.580555, post +0.000000 MATCH: True
ReLU killed 19 of 32 units at this position.
Second layer, output component 0 = sum over 32 hidden units:
(only the 13 non-zero units contribute; the rest add 0)
(+0.2118)*(-0.3195) + (+1.0305)*(+0.1828) + (+0.4628)*(-0.1684) + (+1.7384)*(-0.2064) +
(+0.0673)*(-0.2320) + (+0.0340)*(-0.0034) + (+0.2111)*(-0.3802) + (+0.5934)*(+0.2627) +
(+0.8853)*(+0.1140) + (+0.3648)*(-0.2992) + (+0.9836)*(+0.2737) + (+0.6189)*(+0.3108) +
(+0.5985)*(-0.1744)
+ b2[0] = +0.0000
= +0.092907 | from array: +0.092907 | MATCH: True
B.7 Final Residual + Norm $\rightarrow$ decoder_output[6]
$$\text{decoder\_output}[6] = \mathrm{LayerNorm}\big(Y_2[6] + \text{ffn}[6]\big)$$
This 8-vector is the one the output projection turns into 500 logits in B.8.
dim Y2 + ffn = r3 -> LN array -------------------------------------------------------------------------- 0 -0.974226 0.092907 -0.881320 -0.706660 -0.706660 1 1.341625 -1.137906 0.203719 0.527177 0.527177 2 1.715267 -0.721423 0.993844 1.425657 1.425657 3 -0.487312 -0.338064 -0.825376 -0.643044 -0.643044 4 -0.986904 0.831840 -0.155063 0.119193 0.119193 5 0.564317 0.459811 1.024128 1.460093 1.460093 6 -0.799796 -0.730027 -1.529823 -1.444096 -1.444096 7 -0.372971 -0.536190 -0.909161 -0.738319 -0.738319 mu = -0.259882, var = 0.773338 MATCH decoder_output[6]: True >>> the full decoder block is now traced by hand for position 6.
B.8 From Output to Logit
$$\text{logit}_v = \langle \text{decoder\_output}_6,\; E_v \rangle$$
Because $W_{out} = E^\top$ (weight tying), a logit is literally the inner product between the decoder's output vector and the embedding of candidate token $v$. Below: the true token <eos>, the argmax token, and the geometry that separates them.
decoder_output[6] = [-0.707 0.527 1.426 -0.643 0.119 1.46 -1.444 -0.738] (norm = 2.8284)
--- v = 2 '<eos>' [TRUE next token] ---
(-0.7067)*(+0.0102) + (+0.5272)*(+0.4520) + (+1.4257)*(+0.0676) + (-0.6430)*(+0.0164) +
(+0.1192)*(-0.4808) + (+1.4601)*(+0.2638) + (-1.4441)*(+0.2282) + (-0.7383)*(+0.7648)
= -0.249453 | from array: -0.249453 | MATCH: True
cosine with decoder output = -0.0823, |E_v| = 1.0710
--- v = 5 '<unused_5>' [ARGMAX (model's pick)] ---
(-0.7067)*(-0.1976) + (+0.5272)*(+0.1334) + (+1.4257)*(+0.5535) + (-0.6430)*(-0.0232) +
(+0.1192)*(-0.1963) + (+1.4601)*(+0.6651) + (-1.4441)*(-0.5120) + (-0.7383)*(-0.7774)
= +3.274961 | from array: +3.274961 | MATCH: True
cosine with decoder output = +0.8852, |E_v| = 1.3081
The logit is large when the embedding is both well-aligned in direction
and long in magnitude — an inner product, not a pure angle.
B.9 Softmax Over 500 Tokens
$$P(v) = \frac{e^{\ell_v - m}}{\sum_{v'=1}^{500} e^{\ell_{v'} - m}}$$
The denominator is a sum over the entire vocabulary, which is why a large vocabulary makes the output layer the expensive part of a language model. Below, the sum is verified and the top of the ranking is expanded.
max logit m = 3.274961 (token 5) denominator Z = sum of 500 exponentials = 29.742343 rank v token logit l-m exp /Z = P array ------------------------------------------------------------------------------------ 1 5 <unused_5> 3.2750 0.0000 1.000000 0.033622 0.033622 2 342 <unused_342> 2.9872 -0.2878 0.749925 0.025214 0.025214 3 190 <unused_190> 2.7904 -0.4845 0.615986 0.020711 0.020711 4 487 <unused_487> 2.5049 -0.7701 0.462966 0.015566 0.015566 5 208 <unused_208> 2.4846 -0.7904 0.453665 0.015253 0.015253 6 424 <unused_424> 2.3796 -0.8954 0.408449 0.013733 0.013733 ... 289 2 <eos> -0.2495 -3.5244 0.029469 0.000991 0.000991 <- TRUE ------------------------------------------------------------------------------------ sum of all 500 probabilities = 1.0000000000 MATCH with probs[6]: True True token '<eos>' ranks 289 out of 500 — chance-level, as expected from random weights.
B.10 Loss at This Position
$$\mathcal{L}_6 = -\log P(y_6 = \texttt{<eos>}) = -\big(\ell_{\texttt{<eos>}} - m - \log Z\big)$$
The second form is the log-sum-exp identity — it is how the loss is computed in practice, without ever forming the 500 probabilities explicitly.
direct : -log(P) = -log(0.00099081) = 6.916985 log-sum-exp: -(l_v - m - log Z) = -(-0.249453 - 3.274961 - 3.392572) = 6.916985 identical: True reference values: log(V) = log(500) = 6.214608 (uniform guess) this position = 6.916985 perfect prediction (P = 1) = 0.000000 === TRACE B COMPLETE ===
Every Operation, Counted
For the record, the full forward pass of this block on these two sequences:
| Operation | Where | Multiply-adds |
|---|---|---|
| Encoder $Q,K,V$ projections | $9 \times 8 \times 8 \times 3$ | 1 728 |
| Encoder attention scores | $2 \times 9 \times 9 \times 4$ | 648 |
| Encoder attention $\times V$ | $2 \times 9 \times 9 \times 4$ | 648 |
| Encoder $W^O$ | $9 \times 8 \times 8$ | 576 |
| Encoder FFN | $9 \times (8{\cdot}32 + 32{\cdot}8)$ | 4 608 |
| Decoder self $Q,K,V$ + $W^O$ | $7 \times 8 \times 8 \times 4$ | 1 792 |
| Decoder self attention | $2 \times 7 \times 7 \times 4 \times 2$ | 784 |
| Cross $Q$ (decoder) | $7 \times 8 \times 8$ | 448 |
| Cross $K, V$ (encoder, 9 rows) | $9 \times 8 \times 8 \times 2$ | 1 152 |
| Cross scores + $\times V$ | $2 \times 7 \times 9 \times 4 \times 2$ | 1 008 |
| Cross $W^O$ | $7 \times 8 \times 8$ | 448 |
| Decoder FFN | $7 \times (8{\cdot}32 + 32{\cdot}8)$ | 3 584 |
| Output projection | $7 \times 8 \times 500$ | 28 000 |
Two observations worth carrying forward:
- The output projection dominates — 28 000 of roughly 45 000 multiply-adds, purely because $V \gg d_{model}$. In real models with $V \approx 50\,000$ this layer and its softmax are a major cost, which is why techniques like sampled softmax and tied embeddings exist.
- Attention is not the expensive part here. At $n = 9$ the quadratic term is negligible; the FFN and the output layer dominate. The $O(n^2)$ cost of attention only takes over at long sequence lengths — which is exactly where the various efficient-attention methods aim.