Causal mask Implemented Direction ForwardEquation Causal mask Show code: Causal mask #Link to equation: Causal mask Copy LaTeX S i j ← { S i j j ≤ i − ∞ j > i S_{ij} \leftarrow \begin{cases} S_{ij} & j \le i \\ -\infty & j > i \end{cases} S ij ← { S ij − ∞ j ≤ i j > i File and function transformer.pyattention_forwardTensor shapes The mask has one layer of shape (T, T). NumPy applies it to every head S, scaled (H, T, T)(2, 16, 16)np.triu(..., 1) (T, T)(16, 16)S, masked (H, T, T)(2, 16, 16)Links Explanation Numerical example transformer.py:70-86 Merging heads Implemented Direction ForwardEquation Heads side by side Show code: Heads side by side #Link to equation: Heads side by side Copy LaTeX O = merge ( A V h ) = [ A ( 0 ) V ( 0 ) ⋯ A ( H − 1 ) V ( H − 1 ) ] O = \operatorname{merge}(A V_h) = \begin{bmatrix} A^{(0)} V^{(0)} & \cdots & A^{(H-1)} V^{(H-1)} \end{bmatrix} O = merge ( A V h ) = [ A ( 0 ) V ( 0 ) ⋯ A ( H − 1 ) V ( H − 1 ) ] File and function transformer.pymerge_headsTensor shapes The transpose, then the reshape, in merge_heads A @ Vh (H, T, dh)(2, 16, 16)transpose (T, H, dh)(16, 2, 16)O (T, d)(16, 32)Links Explanation Numerical example transformer.py:65-67 Output projection Implemented Direction ForwardEquation Output projection Show code: Output projection #Link to equation: Output projection Copy LaTeX o u t = O W o + 1 b o ⊤ \mathrm{out} = O W_o + \mathbf{1} b_o^\top out = O W o + 1 b o ⊤ File and function transformer.pyattention_forwardTensor shapes The output has the shape of the input Xn O (T, d)(16, 32)wo (d, d)(32, 32)out (T, d)(16, 32)Links Explanation Numerical example transformer.py:70-86 Q, K, V projections Implemented Direction ForwardEquation Query projection Show code: Query projection #Link to equation: Query projection Copy LaTeX Q = X n W q + 1 b q ⊤ Q = X_n W_q + \mathbf{1} b_q^\top Q = X n W q + 1 b q ⊤ Key projection Show code: Key projection #Link to equation: Key projection Copy LaTeX K = X n W k + 1 b k ⊤ K = X_n W_k + \mathbf{1} b_k^\top K = X n W k + 1 b k ⊤ Value projection Show code: Value projection #Link to equation: Value projection Copy LaTeX V = X n W v + 1 b v ⊤ V = X_n W_v + \mathbf{1} b_v^\top V = X n W v + 1 b v ⊤ File and function transformer.pyattention_forwardTensor shapes One projection. Q, K and V all have this shape Xn (T, d)(16, 32)wq (d, d)(32, 32)Q (T, d)(16, 32)Links Explanation Numerical example transformer.py:70-86 Scaling by sqrt(d_h) Implemented Direction ForwardEquation Scaled scores Show code: Scaled scores #Link to equation: Scaled scores Copy LaTeX S = Q h K h ⊤ d h S = \frac{Q_h K_h^\top}{\sqrt{d_h}} S = d h Q h K h ⊤ File and function transformer.pyattention_forwardTensor shapes Scaling divides by one number, so the shape does not change Qh @ Kh^T (H, T, T)(2, 16, 16)S (H, T, T)(2, 16, 16)Links Explanation Numerical example transformer.py:70-86 Scores Implemented Direction ForwardEquation Scores, one matrix per head Show code: Scores, one matrix per head #Link to equation: Scores, one matrix per head Copy LaTeX Q h K h ⊤ , ( Q h K h ⊤ ) i j = ∑ c ( Q h ) i c ( K h ) j c Q_h K_h^\top, \qquad (Q_h K_h^\top)_{ij} = \sum_{c} (Q_h)_{ic} (K_h)_{jc} Q h K h ⊤ , ( Q h K h ⊤ ) ij = c ∑ ( Q h ) i c ( K h ) j c File and function transformer.pyattention_forwardTensor shapes Scores, before scaling. Both this step and the next are on one line of code Qh (H, T, dh)(2, 16, 16)Kh.transpose(0, 2, 1) (H, dh, T)(2, 16, 16)Qh @ Kh^T (H, T, T)(2, 16, 16)Links Explanation Numerical example transformer.py:70-86 Softmax per row Implemented Direction ForwardEquation Weights, one row at a time Show code: Weights, one row at a time #Link to equation: Weights, one row at a time Copy LaTeX A = softmax ( S ) , ∑ j A i j = 1 A = \operatorname{softmax}(S), \qquad \sum_{j} A_{ij} = 1 A = softmax ( S ) , j ∑ A ij = 1 A masked cell has a weight of exactly 0 Show code: A masked cell has a weight of exactly 0 #Link to equation: A masked cell has a weight of exactly 0 Copy LaTeX exp ( − ∞ ) = 0 ⇒ A i j = 0 for j > i \exp(-\infty) = 0 \quad\Rightarrow\quad A_{ij} = 0 \text{ for } j > i exp ( − ∞ ) = 0 ⇒ A ij = 0 for j > i File and function transformer.pyattention_forwardTensor shapes Softmax keeps the shape. It normalises along the last axis S (H, T, T)(2, 16, 16)A (H, T, T)(2, 16, 16)Links Explanation Numerical example transformer.py:70-86 Splitting into heads Implemented Direction ForwardEquation Split each matrix into heads Show code: Split each matrix into heads #Link to equation: Split each matrix into heads Copy LaTeX Q h , K h , V h = split ( Q , H ) , split ( K , H ) , split ( V , H ) Q_h,\ K_h,\ V_h = \operatorname{split}(Q, H),\ \operatorname{split}(K, H),\ \operatorname{split}(V, H) Q h , K h , V h = split ( Q , H ) , split ( K , H ) , split ( V , H ) File and function transformer.pysplit_headsTensor shapes The reshape, then the transpose, in split_heads Q (T, d)(16, 32)reshape (T, H, dh)(16, 2, 16)transpose (H, T, dh)(2, 16, 16)Links Explanation Numerical example transformer.py:60-62 Weighted values Implemented Direction ForwardEquation Weighted values, one matrix per head Show code: Weighted values, one matrix per head #Link to equation: Weighted values, one matrix per head Copy LaTeX The same product, for one row Derived from implementation Show code: The same product, for one row #Link to equation: The same product, for one row Copy LaTeX ( A V h ) i = ∑ j ≤ i A i j ( V h ) j (A V_h)_{i} = \sum_{j \le i} A_{ij} (V_h)_{j} ( A V h ) i = j ≤ i ∑ A ij ( V h ) j File and function transformer.pyattention_forwardTensor shapes Weighted values have the shape of the values A (H, T, T)(2, 16, 16)Vh (H, T, dh)(2, 16, 16)A @ Vh (H, T, dh)(2, 16, 16)Links Explanation Numerical example transformer.py:70-86 Shuffling and batching Derived from implementation Direction Neither (not differentiated)Equation Batches in one epoch Show code: Batches in one epoch #Link to equation: Batches in one epoch Copy LaTeX batches per epoch = ⌈ W B ⌉ \text{batches per epoch} = \left\lceil \frac{W}{B} \right\rceil batches per epoch = ⌈ B W ⌉ File and function data.pyiterate_batchesTensor shapes From all windows to one sequence tr_x (W, T)(W, 16)xb (B, T)(16, 16)x (T,)(16,)Links Explanation Numerical example data.py:36-43 Sliding windows Implemented Direction Neither (not differentiated)Equation Where windows start Derived from implementation Show code: Where windows start #Link to equation: Where windows start Copy LaTeX s ∈ { 0 , k , 2 k , … } , s < N − T s \in \{0,\, k,\, 2k,\, \dots\}, \qquad s < N - T s ∈ { 0 , k , 2 k , … } , s < N − T Inputs Show code: Inputs #Link to equation: Inputs Copy LaTeX x t ( s ) = ids s + t , t = 0 , … , T − 1 x^{(s)}_{t} = \text{ids}_{s+t}, \qquad t = 0, \dots, T-1 x t ( s ) = ids s + t , t = 0 , … , T − 1 Targets Show code: Targets #Link to equation: Targets Copy LaTeX y t ( s ) = ids s + t + 1 , t = 0 , … , T − 1 y^{(s)}_{t} = \text{ids}_{s+t+1}, \qquad t = 0, \dots, T-1 y t ( s ) = ids s + t + 1 , t = 0 , … , T − 1 Targets are inputs shifted by one Show code: Targets are inputs shifted by one #Link to equation: Targets are inputs shifted by one Copy LaTeX y t ( s ) = x t + 1 ( s ) , t = 0 , … , T − 2 y^{(s)}_{t} = x^{(s)}_{t+1}, \qquad t = 0, \dots, T-2 y t ( s ) = x t + 1 ( s ) , t = 0 , … , T − 2 File and function data.pymake_windowsTensor shapes Inputs of the training windows train_ids (N,)tr_x (W, T)(W, 16)Targets of the training windows train_ids (N,)tr_y (W, T)(W, 16)Links Explanation Numerical example data.py:24-33 The split Implemented Direction Neither (not differentiated)Equation Where the stream is cut Show code: Where the stream is cut #Link to equation: Where the stream is cut Copy LaTeX cut = ⌊ ( 1 − ρ ) N ⌋ \text{cut} = \lfloor (1 - \rho)\, N \rfloor cut = ⌊( 1 − ρ ) N ⌋ File and function data.pysplit_idsTensor shapes None declared Links Explanation Numerical example data.py:19-21 Window count Derived from implementation Direction Neither (not differentiated)Equation Number of windows Show code: Number of windows #Link to equation: Number of windows Copy LaTeX W = ⌈ N − T k ⌉ , N > T W = \left\lceil \frac{N - T}{k} \right\rceil, \qquad N > T W = ⌈ k N − T ⌉ , N > T File and function data.pymake_windowsTensor shapes None declared Links Explanation Numerical example data.py:24-33 Positional embedding Implemented Direction ForwardEquation The first T rows of the position table Show code: The first T rows of the position table #Link to equation: The first T rows of the position table Copy LaTeX ( P [ : T ] ) t = P t , t = 0 , … , T − 1 \left(P[:T]\right)_t = P_t, \quad t = 0, \ldots, T - 1 ( P [ : T ] ) t = P t , t = 0 , … , T − 1 File and function transformer.pyembed_forwardTensor shapes The table is cut to the length of the sequence pos_emb (Tmax, d)(32, 32)pos_emb[:T] (T, d)(16, 32)Links Explanation Numerical example transformer.py:52-53 Token embedding Implemented Direction ForwardEquation Row t of the lookup is the table row of the t-th id Show code: Row t of the lookup is the table row of the t-th id #Link to equation: Row t of the lookup is the table row of the t-th id Copy LaTeX ( E [ ids ] ) t = E i t \left(E[\text{ids}]\right)_t = E_{i_t} ( E [ ids ] ) t = E i t File and function transformer.pyembed_forwardTensor shapes From ids to vectors ids (T,)(16,)tok_emb[ids] (T, d)(16, 32)The table itself. V depends on the corpus rows (V,)tok_emb (V, d)(V, 32)Links Explanation Numerical example transformer.py:52-53 Autoregressive loop Implemented Direction Neither (not differentiated)Equation The last window ids Show code: The last window ids #Link to equation: The last window ids Copy LaTeX x = i d s [ − w : ] x = \mathrm{ids}[-w:] x = ids [ − w : ] Logits of the last position only Show code: Logits of the last position only #Link to equation: Logits of the last position only Copy LaTeX Z = f o r w a r d ( x ) , z = Z T − 1 Z = \mathrm{forward}(x), \qquad z = Z_{T-1} Z = forward ( x ) , z = Z T − 1 Sample and append Show code: Sample and append #Link to equation: Sample and append Copy LaTeX i d s ← i d s + [ s a m p l e _ n e x t ( z , τ , k , ρ ) ] \mathrm{ids} \leftarrow \mathrm{ids} + [\,\mathrm{sample\_next}(z, \tau, k, \rho)\,] ids ← ids + [ sample_next ( z , τ , k , ρ ) ] File and function generate.pygenerateTensor shapes One step of the loop x (T,)(16,)logits (T, V)(16, V)logits[-1] (V,)Links Explanation No numerical example generate.py:28-35 Checking every parameter Implemented Direction Neither (not differentiated)Equation Worst relative error for one parameter name Show code: Worst relative error for one parameter name #Link to equation: Worst relative error for one parameter name Copy LaTeX worst [ n ] = max s = 1 , … , S rel ( ∂ L ∂ θ n [ i s ] , L ( θ n [ i s ] + h ) − L ( θ n [ i s ] − h ) 2 h ) \operatorname{worst}[n] = \max_{s = 1,\ldots,S}\ \operatorname{rel}\!\left(\frac{\partial L}{\partial \theta_n}[i_s],\ \frac{L(\theta_n[i_s] + h) - L(\theta_n[i_s] - h)}{2h}\right) worst [ n ] = s = 1 , … , S max rel ( ∂ θ n ∂ L [ i s ] , 2 h L ( θ n [ i s ] + h ) − L ( θ n [ i s ] − h ) ) File and function gradcheck.pycheck_paramsTensor shapes One name, at the sizes in config.py. The test itself uses a smaller ModelConfig W = pos_emb (Tmax, d)(32, 32)grads["pos_emb"] (Tmax, d)(32, 32)Links Explanation Numerical example gradcheck.py:19-33 Central difference Implemented Direction Neither (not differentiated)Equation Central difference for one entry Show code: Central difference for one entry #Link to equation: Central difference for one entry Copy LaTeX g i = f ( x + h e i ) − f ( x − h e i ) 2 h g_i = \frac{f(x + h\, e_i) - f(x - h\, e_i)}{2h} g i = 2 h f ( x + h e i ) − f ( x − h e i ) File and function gradcheck.pynumerical_gradTensor shapes None declared Links Explanation Numerical example gradcheck.py:3-14 Relative error Implemented Direction Neither (not differentiated)Equation Relative error with the epsilon guard Show code: Relative error with the epsilon guard #Link to equation: Relative error with the epsilon guard Copy LaTeX rel ( a , b ) = ∣ a − b ∣ max ( ∣ a ∣ , ∣ b ∣ , ε ) \operatorname{rel}(a, b) = \frac{|a - b|}{\max\left(|a|,\, |b|,\, \varepsilon\right)} rel ( a , b ) = max ( ∣ a ∣ , ∣ b ∣ , ε ) ∣ a − b ∣ File and function gradcheck.pyrelative_errorTensor shapes None declared Links Explanation Numerical example gradcheck.py:16-17 Heads Implemented Direction ForwardEquation Width of one head Show code: Width of one head #Link to equation: Width of one head Copy LaTeX d h = d H , d m o d H = 0 d_h = \frac{d}{H}, \qquad d \bmod H = 0 d h = H d , d mod H = 0 Head h owns one block of columns Show code: Head h owns one block of columns #Link to equation: Head h owns one block of columns Copy LaTeX M t , i ( h ) = M t , h d h + i M^{(h)}_{t,i} = M_{t,\; h d_h + i} M t , i ( h ) = M t , h d h + i File and function transformer.pysplit_headstransformer.pymerge_headsTensor shapes Split, for the query matrix Q (T, d)(16, 32)Qh (H, T, dh)(2, 16, 16)Merge, after the weighted values A @ Vh (H, T, dh)(2, 16, 16)O (T, d)(16, 32)Links Explanation Numerical example transformer.py:60-62 transformer.py:65-67 Matrix multiplication Derived from implementation Direction ForwardEquation One entry of the product Show code: One entry of the product #Link to equation: One entry of the product Copy LaTeX ( X W ) t j = ∑ k X t k W k j (XW)_{tj} = \sum_{k} X_{tk} W_{kj} ( X W ) t j = k ∑ X t k W k j The shape rule Show code: The shape rule #Link to equation: The shape rule Copy LaTeX X ∈ R T × a , W ∈ R a × c ⇒ X W ∈ R T × c X \in \mathbb{R}^{T \times a},\; W \in \mathbb{R}^{a \times c} \;\Rightarrow\; XW \in \mathbb{R}^{T \times c} X ∈ R T × a , W ∈ R a × c ⇒ X W ∈ R T × c A linear layer, with b added to every row Show code: A linear layer, with b added to every row #Link to equation: A linear layer, with b added to every row Copy LaTeX Y t j = ( X W ) t j + b j Y_{tj} = (XW)_{tj} + b_j Y t j = ( X W ) t j + b j File and function ops.pylinear_forwardTensor shapes Query projection, width stays the same Xn (T, d)(16, 32)Q (T, d)(16, 32)First feed-forward layer, width grows Xn (T, d)(16, 32)Z (T, F)(16, 128)Linear head, width becomes the vocabulary size Xf (T, d)(16, 32)logits (T, V)(16, V)Links Explanation Numerical example ops.py:7-8 Cross-entropy, backward Implemented Direction BackwardEquation The derivative of the loss at one position Show code: The derivative of the loss at one position #Link to equation: The derivative of the loss at one position Copy LaTeX ∂ ℓ t ∂ logits t v = P t v − 1 [ v = y t ] \frac{\partial \ell_t}{\partial \text{logits}_{tv}} = P_{tv} - \mathbf{1}[v = y_t] ∂ logits t v ∂ ℓ t = P t v − 1 [ v = y t ] The combined backward, where E_y is one-hot at each correct token Show code: The combined backward, where E_y is one-hot at each correct token #Link to equation: The combined backward, where E_y is one-hot at each correct token Copy LaTeX d Z = P − E y T dZ = \frac{P - E_y}{T} d Z = T P − E y File and function ops.pycross_entropy_backwardTensor shapes The start of the backward pass probs (T, V)(16, V)dZ (T, V)(16, V)Links Explanation Numerical example ops.py:36-40 Cross-entropy Implemented Direction ForwardEquation The same shift as softmax, written inline Show code: The same shift as softmax, written inline #Link to equation: The same shift as softmax, written inline Copy LaTeX z t v = logits t v − max k logits t k z_{tv} = \text{logits}_{tv} - \max_{k} \text{logits}_{tk} z t v = logits t v − k max logits t k Log-softmax Show code: Log-softmax #Link to equation: Log-softmax Copy LaTeX log p t v = z t v − log ∑ k e z t k \log p_{tv} = z_{tv} - \log \sum_{k} e^{z_{tk}} log p t v = z t v − log k ∑ e z t k Mean over positions Show code: Mean over positions #Link to equation: Mean over positions Copy LaTeX L = − 1 T ∑ t log p t , y t L = -\frac{1}{T} \sum_{t} \log p_{t,\, y_t} L = − T 1 t ∑ log p t , y t The probabilities it returns Show code: The probabilities it returns #Link to equation: The probabilities it returns Copy LaTeX P t v = e log p t v P_{tv} = e^{\log p_{tv}} P t v = e l o g p t v File and function ops.pycross_entropy_forwardTensor shapes From logits to probabilities. The loss is a single number logits (T, V)(16, V)logp (T, V)(16, V)probs (T, V)(16, V)Links Explanation Numerical example ops.py:29-34 LayerNorm, backward Implemented Direction BackwardEquation Gradient on the normalised row Show code: Gradient on the normalised row #Link to equation: Gradient on the normalised row Copy LaTeX d X ^ t i = d Y t i γ i d\hat{X}_{ti} = dY_{ti}\, \gamma_i d X ^ t i = d Y t i γ i Gradient on the input Show code: Gradient on the input #Link to equation: Gradient on the input Copy LaTeX d X t j = s t ( d X ^ t j − mean i ( d X ^ t i ) − X ^ t j mean i ( d X ^ t i X ^ t i ) ) dX_{tj} = s_t \left( d\hat{X}_{tj} - \operatorname{mean}_i\!\left(d\hat{X}_{ti}\right) - \hat{X}_{tj}\, \operatorname{mean}_i\!\left(d\hat{X}_{ti} \hat{X}_{ti}\right) \right) d X t j = s t ( d X ^ t j − mean i ( d X ^ t i ) − X ^ t j mean i ( d X ^ t i X ^ t i ) ) Gradient on the scale, summed over rows Show code: Gradient on the scale, summed over rows #Link to equation: Gradient on the scale, summed over rows Copy LaTeX d γ i = ∑ t d Y t i X ^ t i d\gamma_i = \sum_{t} dY_{ti}\, \hat{X}_{ti} d γ i = t ∑ d Y t i X ^ t i Gradient on the shift, summed over rows Show code: Gradient on the shift, summed over rows #Link to equation: Gradient on the shift, summed over rows Copy LaTeX d β i = ∑ t d Y t i d\beta_i = \sum_{t} dY_{ti} d β i = t ∑ d Y t i File and function ops.pylayer_norm_backwardTensor shapes The final LayerNorm dXf (T, d)(16, 32)dX (T, d)(16, 32)dg (d,)(32,)db (d,)(32,)LayerNorm before attention, inside a block dXn (T, d)(16, 32)dX_branch (T, d)(16, 32)dg (d,)(32,)db (d,)(32,)Links Explanation Numerical example ops.py:49-57 Deriving the backward formula Derived from implementation Direction BackwardEquation Step 1: from the output to the normalised row Show code: Step 1: from the output to the normalised row #Link to equation: Step 1: from the output to the normalised row Copy LaTeX ∂ L ∂ X ^ t i = ∂ L ∂ Y t i ∂ Y t i ∂ X ^ t i = d Y t i γ i = d X ^ t i \frac{\partial L}{\partial \hat{X}_{ti}} = \frac{\partial L}{\partial Y_{ti}}\, \frac{\partial Y_{ti}}{\partial \hat{X}_{ti}} = dY_{ti}\, \gamma_i = d\hat{X}_{ti} ∂ X ^ t i ∂ L = ∂ Y t i ∂ L ∂ X ^ t i ∂ Y t i = d Y t i γ i = d X ^ t i Step 2: the derivative of the mean Show code: Step 2: the derivative of the mean #Link to equation: Step 2: the derivative of the mean Copy LaTeX ∂ μ t ∂ X t j = 1 d \frac{\partial \mu_t}{\partial X_{tj}} = \frac{1}{d} ∂ X t j ∂ μ t = d 1 Step 3: the derivative of the variance, where the mu term cancels Show code: Step 3: the derivative of the variance, where the mu term cancels #Link to equation: Step 3: the derivative of the variance, where the mu term cancels Copy LaTeX ∂ σ t 2 ∂ X t j = 2 d ∑ i ( X t i − μ t ) ( δ i j − 1 d ) = 2 d ( X t j − μ t ) \frac{\partial \sigma_t^2}{\partial X_{tj}} = \frac{2}{d} \sum_{i} \left(X_{ti} - \mu_t\right)\left(\delta_{ij} - \frac{1}{d}\right) = \frac{2}{d}\left(X_{tj} - \mu_t\right) ∂ X t j ∂ σ t 2 = d 2 i ∑ ( X t i − μ t ) ( δ ij − d 1 ) = d 2 ( X t j − μ t ) Step 4: the derivative of s, with epsilon absorbed into s Show code: Step 4: the derivative of s, with epsilon absorbed into s #Link to equation: Step 4: the derivative of s, with epsilon absorbed into s Copy LaTeX ∂ s t ∂ X t j = − 1 2 ( σ t 2 + ϵ ) − 3 / 2 2 d ( X t j − μ t ) = − 1 d s t 3 ( X t j − μ t ) \frac{\partial s_t}{\partial X_{tj}} = -\frac{1}{2}\left(\sigma_t^2 + \epsilon\right)^{-3/2} \frac{2}{d}\left(X_{tj} - \mu_t\right) = -\frac{1}{d}\, s_t^3 \left(X_{tj} - \mu_t\right) ∂ X t j ∂ s t = − 2 1 ( σ t 2 + ϵ ) − 3/2 d 2 ( X t j − μ t ) = − d 1 s t 3 ( X t j − μ t ) Step 5: the product rule on the normalised row Show code: Step 5: the product rule on the normalised row #Link to equation: Step 5: the product rule on the normalised row Copy LaTeX ∂ X ^ t i ∂ X t j = s t ( δ i j − 1 d ) + ( X t i − μ t ) ∂ s t ∂ X t j = s t ( δ i j − 1 d − 1 d X ^ t i X ^ t j ) \frac{\partial \hat{X}_{ti}}{\partial X_{tj}} = s_t\left(\delta_{ij} - \frac{1}{d}\right) + \left(X_{ti} - \mu_t\right)\frac{\partial s_t}{\partial X_{tj}} = s_t\left(\delta_{ij} - \frac{1}{d} - \frac{1}{d}\,\hat{X}_{ti}\hat{X}_{tj}\right) ∂ X t j ∂ X ^ t i = s t ( δ ij − d 1 ) + ( X t i − μ t ) ∂ X t j ∂ s t = s t ( δ ij − d 1 − d 1 X ^ t i X ^ t j ) Step 6: the chain rule over all features of the row Show code: Step 6: the chain rule over all features of the row #Link to equation: Step 6: the chain rule over all features of the row Copy LaTeX d X t j = ∑ i d X ^ t i ∂ X ^ t i ∂ X t j = s t ( d X ^ t j − 1 d ∑ i d X ^ t i − X ^ t j 1 d ∑ i d X ^ t i X ^ t i ) dX_{tj} = \sum_{i} d\hat{X}_{ti}\, \frac{\partial \hat{X}_{ti}}{\partial X_{tj}} = s_t \left( d\hat{X}_{tj} - \frac{1}{d}\sum_{i} d\hat{X}_{ti} - \hat{X}_{tj}\, \frac{1}{d}\sum_{i} d\hat{X}_{ti}\hat{X}_{ti} \right) d X t j = i ∑ d X ^ t i ∂ X t j ∂ X ^ t i = s t ( d X ^ t j − d 1 i ∑ d X ^ t i − X ^ t j d 1 i ∑ d X ^ t i X ^ t i ) The parameters, each of which is shared by every row Show code: The parameters, each of which is shared by every row #Link to equation: The parameters, each of which is shared by every row Copy LaTeX ∂ Y t i ∂ γ i = X ^ t i , ∂ Y t i ∂ β i = 1 \frac{\partial Y_{ti}}{\partial \gamma_i} = \hat{X}_{ti},\quad \frac{\partial Y_{ti}}{\partial \beta_i} = 1 ∂ γ i ∂ Y t i = X ^ t i , ∂ β i ∂ Y t i = 1 File and function ops.pylayer_norm_forwardops.pylayer_norm_backwardTensor shapes Each row has its own matrix of derivatives. The code never builds it dXhat row (d,)(32,)dX/dXhat matrix (d, d)(32, 32)dX row (d,)(32,)Links Explanation Numerical example ops.py:42-47 ops.py:49-57 LayerNorm Implemented Direction ForwardEquation Mean of one row, over the last axis Show code: Mean of one row, over the last axis #Link to equation: Mean of one row, over the last axis Copy LaTeX μ t = 1 d ∑ i X t i \mu_t = \frac{1}{d} \sum_{i} X_{ti} μ t = d 1 i ∑ X t i Biased variance of one row. It divides by d Show code: Biased variance of one row. It divides by d #Link to equation: Biased variance of one row. It divides by d Copy LaTeX σ t 2 = 1 d ∑ i ( X t i − μ t ) 2 \sigma_t^2 = \frac{1}{d} \sum_{i} \left(X_{ti} - \mu_t\right)^2 σ t 2 = d 1 i ∑ ( X t i − μ t ) 2 Inverse standard deviation Show code: Inverse standard deviation #Link to equation: Inverse standard deviation Copy LaTeX s t = 1 σ t 2 + ϵ s_t = \frac{1}{\sqrt{\sigma_t^2 + \epsilon}} s t = σ t 2 + ϵ 1 The normalised row Show code: The normalised row #Link to equation: The normalised row Copy LaTeX X ^ t i = ( X t i − μ t ) s t \hat{X}_{ti} = \left(X_{ti} - \mu_t\right) s_t X ^ t i = ( X t i − μ t ) s t Scale and shift, one pair per feature Show code: Scale and shift, one pair per feature #Link to equation: Scale and shift, one pair per feature Copy LaTeX Y t i = γ i X ^ t i + β i Y_{ti} = \gamma_i \hat{X}_{ti} + \beta_i Y t i = γ i X ^ t i + β i File and function ops.pylayer_norm_forwardTensor shapes Before attention X (T, d)(16, 32)Xn (T, d)(16, 32)Before the feed-forward layer X1 (T, d)(16, 32)Xn2 (T, d)(16, 32)At the end of the model, before the head X (T, d)(16, 32)Xf (T, d)(16, 32)Links Explanation Numerical example ops.py:42-47 Linear layer, backward Implemented Direction BackwardEquation The derivatives of one output entry Derived from implementation Show code: The derivatives of one output entry #Link to equation: The derivatives of one output entry Copy LaTeX ∂ Y t j ∂ W k j = X t k , ∂ Y t j ∂ b j = 1 , ∂ Y t j ∂ X t k = W k j \frac{\partial Y_{tj}}{\partial W_{kj}} = X_{tk},\quad \frac{\partial Y_{tj}}{\partial b_{j}} = 1,\quad \frac{\partial Y_{tj}}{\partial X_{tk}} = W_{kj} ∂ W k j ∂ Y t j = X t k , ∂ b j ∂ Y t j = 1 , ∂ X t k ∂ Y t j = W k j Gradient on the weights Show code: Gradient on the weights #Link to equation: Gradient on the weights Copy LaTeX d W = X ⊤ d Y dW = X^{\top} dY d W = X ⊤ d Y Gradient on the bias, the sum of dY over rows Show code: Gradient on the bias, the sum of dY over rows #Link to equation: Gradient on the bias, the sum of dY over rows Copy LaTeX d b = ∑ t d Y t db = \sum_{t} dY_{t} d b = t ∑ d Y t Gradient on the input Show code: Gradient on the input #Link to equation: Gradient on the input Copy LaTeX d X = d Y W ⊤ dX = dY\, W^{\top} d X = d Y W ⊤ File and function ops.pylinear_backwardTensor shapes Linear head dZ (T, V)(16, V)dW (d, V)(32, V)db (V,)dXf (T, d)(16, 32)Second feed-forward layer dOut (T, d)(16, 32)dW2 (F, d)(128, 32)db2 (d,)(32,)dA (T, F)(16, 128)First feed-forward layer dZ (T, F)(16, 128)dW1 (d, F)(32, 128)db1 (F,)(128,)dXn (T, d)(16, 32)Links Explanation Numerical example ops.py:10-14 Linear layer Implemented Direction ForwardEquation The linear layer, with b added to every row of XW Show code: The linear layer, with b added to every row of XW #Link to equation: The linear layer, with b added to every row of XW Copy LaTeX File and function ops.pylinear_forwardTensor shapes Attention output projection, width stays the same O (T, d)(16, 32)out (T, d)(16, 32)First feed-forward layer, width grows Xn (T, d)(16, 32)Z (T, F)(16, 128)Second feed-forward layer, width shrinks back A (T, F)(16, 128)output (T, d)(16, 32)Linear head, width becomes the vocabulary size Xf (T, d)(16, 32)logits (T, V)(16, V)Links Explanation Numerical example ops.py:7-8 Perplexity Implemented Direction Neither (not differentiated)Equation Perplexity Show code: Perplexity #Link to equation: Perplexity Copy LaTeX perplexity = e L \text{perplexity} = e^{L} perplexity = e L File and function ops.pyperplexityTensor shapes None declared Links Explanation Numerical example ops.py:59-60 ReLU, backward Implemented Direction BackwardEquation The derivative as the code uses it, with 0 at exactly Z = 0 Derived from implementation Show code: The derivative as the code uses it, with 0 at exactly Z = 0 #Link to equation: The derivative as the code uses it, with 0 at exactly Z = 0 Copy LaTeX ∂ A t i ∂ Z t i = { 1 Z t i > 0 0 Z t i ≤ 0 \frac{\partial A_{ti}}{\partial Z_{ti}} = \begin{cases} 1 & Z_{ti} > 0 \\ 0 & Z_{ti} \le 0 \end{cases} ∂ Z t i ∂ A t i = { 1 0 Z t i > 0 Z t i ≤ 0 The mask-based backward Show code: The mask-based backward #Link to equation: The mask-based backward Copy LaTeX d Z = d A ⊙ 1 [ Z > 0 ] dZ = dA \odot \mathbf{1}[Z > 0] d Z = d A ⊙ 1 [ Z > 0 ] File and function ops.pyrelu_backwardTensor shapes Inside the feed-forward backward dA (T, F)(16, 128)dZ (T, F)(16, 128)Links Explanation Numerical example ops.py:19-20 ReLU Implemented Direction ForwardEquation ReLU, on each entry Show code: ReLU, on each entry #Link to equation: ReLU, on each entry Copy LaTeX A t i = max ( 0 , Z t i ) A_{ti} = \max(0,\, Z_{ti}) A t i = max ( 0 , Z t i ) File and function ops.pyrelu_forwardTensor shapes Inside the feed-forward layer Z (T, F)(16, 128)A (T, F)(16, 128)Links Explanation Numerical example ops.py:16-17 Softmax, backward Implemented Direction BackwardEquation The derivative of softmax. The code never builds this matrix Show code: The derivative of softmax. The code never builds this matrix #Link to equation: The derivative of softmax. The code never builds this matrix Copy LaTeX ∂ y i ∂ x j = y i ( δ i j − y j ) \frac{\partial y_i}{\partial x_j} = y_i\left(\delta_{ij} - y_j\right) ∂ x j ∂ y i = y i ( δ ij − y j ) The backward line, which multiplies dY by that derivative Show code: The backward line, which multiplies dY by that derivative #Link to equation: The backward line, which multiplies dY by that derivative Copy LaTeX d x j = y j ( d Y j − ∑ i y i d Y i ) dx_j = y_j\left(dY_j - \sum_i y_i\, dY_i\right) d x j = y j ( d Y j − i ∑ y i d Y i ) File and function ops.pysoftmax_backwardTensor shapes Attention dA (H, T, T)(2, 16, 16)dS (H, T, T)(2, 16, 16)Links Explanation Numerical example ops.py:26-27 Softmax Implemented Direction ForwardEquation Shift by the maximum, then exponentiate Show code: Shift by the maximum, then exponentiate #Link to equation: Shift by the maximum, then exponentiate Copy LaTeX e i = exp ( x i − max k x k ) e_i = \exp\left(x_i - \max_k x_k\right) e i = exp ( x i − k max x k ) Normalize Show code: Normalize #Link to equation: Normalize Copy LaTeX p i = e i ∑ k e k p_i = \frac{e_i}{\sum_k e_k} p i = ∑ k e k e i Why the shift is safe. The factor from any constant c cancels Show code: Why the shift is safe. The factor from any constant c cancels #Link to equation: Why the shift is safe. The factor from any constant c cancels Copy LaTeX e x i − c ∑ k e x k − c = e − c e x i e − c ∑ k e x k = e x i ∑ k e x k \frac{e^{x_i - c}}{\sum_k e^{x_k - c}} = \frac{e^{-c}\, e^{x_i}}{e^{-c} \sum_k e^{x_k}} = \frac{e^{x_i}}{\sum_k e^{x_k}} ∑ k e x k − c e x i − c = e − c ∑ k e x k e − c e x i = ∑ k e x k e x i File and function ops.pysoftmax_forwardTensor shapes Attention weights S (H, T, T)(2, 16, 16)A (H, T, T)(2, 16, 16)Sampling logits / temperature (V,)p (V,)Links Explanation Numerical example ops.py:22-24 Xavier uniform initialisation Implemented Direction Neither (not differentiated)Equation The limit Show code: The limit #Link to equation: The limit Copy LaTeX a = 6 fan_in + fan_out a = \sqrt{\frac{6}{\text{fan\_in} + \text{fan\_out}}} a = fan_in + fan_out 6 Each entry is drawn evenly between -a and a Show code: Each entry is drawn evenly between -a and a #Link to equation: Each entry is drawn evenly between -a and a Copy LaTeX W i j ∼ U [ − a , a ] W_{ij} \sim U[-a,\, a] W ij ∼ U [ − a , a ] The variance this gives. The code sets a and never computes this Show code: The variance this gives. The code sets a and never computes this #Link to equation: The variance this gives. The code sets a and never computes this Copy LaTeX Var ( W i j ) = a 2 3 = 2 fan_in + fan_out \operatorname{Var}(W_{ij}) = \frac{a^2}{3} = \frac{2}{\text{fan\_in} + \text{fan\_out}} Var ( W ij ) = 3 a 2 = fan_in + fan_out 2 File and function ops.pyxavier_uniformTensor shapes A query, key, value or output matrix fan_in, fan_out (d, d)(32, 32)w{q,k,v,o} (d, d)(32, 32)The first feed-forward matrix fan_in, fan_out (d, F)(32, 128)ff1_w (d, F)(32, 128)Links Explanation Numerical example ops.py:3-5 AdamW update Implemented Direction Neither (not differentiated)Equation Moving average of the gradient Show code: Moving average of the gradient #Link to equation: Moving average of the gradient Copy LaTeX m t = β 1 m t − 1 + ( 1 − β 1 ) g t m_t = \beta_1 m_{t-1} + (1 - \beta_1)\, g_t m t = β 1 m t − 1 + ( 1 − β 1 ) g t Moving average of the squared gradient Show code: Moving average of the squared gradient #Link to equation: Moving average of the squared gradient Copy LaTeX v t = β 2 v t − 1 + ( 1 − β 2 ) g t 2 v_t = \beta_2 v_{t-1} + (1 - \beta_2)\, g_t^2 v t = β 2 v t − 1 + ( 1 − β 2 ) g t 2 Bias correction of m Show code: Bias correction of m #Link to equation: Bias correction of m Copy LaTeX m ^ t = m t 1 − β 1 t \hat m_t = \frac{m_t}{1 - \beta_1^{\,t}} m ^ t = 1 − β 1 t m t Bias correction of v Show code: Bias correction of v #Link to equation: Bias correction of v Copy LaTeX v ^ t = v t 1 − β 2 t \hat v_t = \frac{v_t}{1 - \beta_2^{\,t}} v ^ t = 1 − β 2 t v t The update before decay Show code: The update before decay #Link to equation: The update before decay Copy LaTeX u t = m ^ t v ^ t + ϵ u_t = \frac{\hat m_t}{\sqrt{\hat v_t} + \epsilon} u t = v ^ t + ϵ m ^ t Decoupled weight decay Show code: Decoupled weight decay #Link to equation: Decoupled weight decay Copy LaTeX u t ← u t + λ w t − 1 only where w has two or more dimensions u_t \leftarrow u_t + \lambda\, w_{t-1} \qquad \text{only where } w \text{ has two or more dimensions} u t ← u t + λ w t − 1 only where w has two or more dimensions The parameter update, in place Show code: The parameter update, in place #Link to equation: The parameter update, in place Copy LaTeX w t = w t − 1 − η t u t w_t = w_{t-1} - \eta_t\, u_t w t = w t − 1 − η t u t File and function optim.pystepTensor shapes One key with two dimensions, head_w, whose update includes decay grads[k] (d, V)(32, V)self.m[k] (d, V)(32, V)self.v[k] (d, V)(32, V)update (d, V)(32, V)params[k] (d, V)(32, V)One key with one dimension, head_b, whose update has no decay grads[k] (V,)self.m[k] (V,)self.v[k] (V,)update (V,)params[k] (V,)Links Explanation Numerical example optim.py:13-30 Gradient clipping Implemented Direction Neither (not differentiated)Equation Global norm, over all tensors Show code: Global norm, over all tensors #Link to equation: Global norm, over all tensors Copy LaTeX ∥ g ∥ = ∑ k ∑ i g k , i 2 \lVert g \rVert = \sqrt{\sum_{k} \sum_{i} g_{k,i}^2} ∥ g ∥ = k ∑ i ∑ g k , i 2 Scaling by one factor Show code: Scaling by one factor #Link to equation: Scaling by one factor Copy LaTeX g k ← g k ⋅ c ∥ g ∥ for every key k , only when ∥ g ∥ > c g_k \leftarrow g_k \cdot \frac{c}{\lVert g \rVert} \qquad \text{for every key } k, \text{ only when } \lVert g \rVert > c g k ← g k ⋅ ∥ g ∥ c for every key k , only when ∥ g ∥ > c File and function optim.pyclip_grad_normTensor shapes Clipping changes values, not shapes. head_w as one example grads[k] before (d, V)(32, V)grads[k] after (d, V)(32, V)Links Explanation Numerical example optim.py:40-50 Learning-rate schedule Implemented Direction Neither (not differentiated)Equation Linear warmup Show code: Linear warmup #Link to equation: Linear warmup Copy LaTeX η ( t ) = η max ⋅ t t warm t ≤ t warm \eta(t) = \eta_{\max} \cdot \frac{t}{t_{\text{warm}}} \qquad t \le t_{\text{warm}} η ( t ) = η m a x ⋅ t warm t t ≤ t warm How far through the decay Show code: How far through the decay #Link to equation: How far through the decay Copy LaTeX f = t − t warm t total − t warm t > t warm f = \frac{t - t_{\text{warm}}}{t_{\text{total}} - t_{\text{warm}}} \qquad t > t_{\text{warm}} f = t total − t warm t − t warm t > t warm Cosine decay Show code: Cosine decay #Link to equation: Cosine decay Copy LaTeX η ( t ) = η min + 1 2 ( η max − η min ) ( 1 + cos ( π f ) ) \eta(t) = \eta_{\min} + \tfrac{1}{2} (\eta_{\max} - \eta_{\min}) \big(1 + \cos(\pi f)\big) η ( t ) = η m i n + 2 1 ( η m a x − η m i n ) ( 1 + cos ( π f ) ) File and function optim.pylr_scheduleTensor shapes None declared Links Explanation Numerical example optim.py:32-37 Random choice Implemented Direction Neither (not differentiated)Equation One draw from the final probabilities Show code: One draw from the final probabilities #Link to equation: One draw from the final probabilities Copy LaTeX P ( n e x t = i ) = p i P(\mathrm{next} = i) = p_i P ( next = i ) = p i File and function generate.pysample_nextTensor shapes The array before the draw. The draw itself returns one int, with no axes p after top-k (V,)p after top-p (V,)Links Explanation No numerical example generate.py:7-26 Greedy guard Implemented Direction Neither (not differentiated)Equation Greedy guard Show code: Greedy guard #Link to equation: Greedy guard Copy LaTeX τ < 10 − 6 ⟹ n e x t = arg max i z i \tau < 10^{-6} \implies \mathrm{next} = \arg\max_i z_i τ < 1 0 − 6 ⟹ next = arg i max z i File and function generate.pysample_nextTensor shapes From the model's logits to the row sample_next reads. The result is one int, with no axes logits (T, V)(16, V)logits[-1] (V,)Links Explanation No numerical example generate.py:7-26 Temperature Implemented Direction Neither (not differentiated)Equation Softmax of the logits divided by temperature Show code: Softmax of the logits divided by temperature #Link to equation: Softmax of the logits divided by temperature Copy LaTeX p = s o f t m a x ( z τ ) p = \mathrm{softmax}\left(\frac{z}{\tau}\right) p = softmax ( τ z ) File and function generate.pysample_nextTensor shapes Temperature keeps the shape logits (V,)p (V,)Links Explanation No numerical example generate.py:7-26 Top-k Derived from implementation Direction Neither (not differentiated)Equation The cut-off Show code: The cut-off #Link to equation: The cut-off Copy LaTeX c = k -th largest value of p , counting repeats c = k\text{-th largest value of } p \text{, counting repeats} c = k -th largest value of p , counting repeats Zero everything below the cut-off Show code: Zero everything below the cut-off #Link to equation: Zero everything below the cut-off Copy LaTeX p i ′ = { p i p i ≥ c 0 otherwise p'_i = \begin{cases} p_i & p_i \geq c \\ 0 & \text{otherwise} \end{cases} p i ′ = { p i 0 p i ≥ c otherwise Renormalise Show code: Renormalise #Link to equation: Renormalise Copy LaTeX p i ′ ′ = p i ′ ∑ j p j ′ p''_i = \frac{p'_i}{\sum_j p'_j} p i ′′ = ∑ j p j ′ p i ′ File and function generate.pysample_nextTensor shapes Top-k changes values, not the shape. Removed tokens stay as zeros p before (V,)p after (V,)Links Explanation No numerical example generate.py:7-26 Top-p Derived from implementation Direction Neither (not differentiated)Equation Sort Show code: Sort #Link to equation: Sort Copy LaTeX π = token indices sorted by p , largest first \pi = \text{token indices sorted by } p \text{, largest first} π = token indices sorted by p , largest first Cumulative sum Show code: Cumulative sum #Link to equation: Cumulative sum Copy LaTeX C m = ∑ j = 0 m p π j C_m = \sum_{j=0}^{m} p_{\pi_j} C m = j = 0 ∑ m p π j The smallest prefix that reaches top_p Show code: The smallest prefix that reaches top_p #Link to equation: The smallest prefix that reaches top_p Copy LaTeX s = min { m : C m ≥ ρ } , keep π 0 , … , π s s = \min \{ m : C_m \geq \rho \}, \qquad \text{keep } \pi_0, \dots, \pi_s s = min { m : C m ≥ ρ } , keep π 0 , … , π s Renormalise Show code: Renormalise #Link to equation: Renormalise Copy LaTeX q i = p i ∑ m = 0 s p π m for kept i , q i = 0 otherwise q_i = \frac{p_i}{\sum_{m=0}^{s} p_{\pi_m}} \text{ for kept } i, \qquad q_i = 0 \text{ otherwise} q i = ∑ m = 0 s p π m p i for kept i , q i = 0 otherwise File and function generate.pysample_nextTensor shapes Top-p sorts and sums over all V tokens. Only keep is shorter, and p keeps its shape order (V,)cum (V,)p (V,)Links Explanation No numerical example generate.py:7-26 Encode and decode Implemented Direction Neither (not differentiated)Equation The id of a token Show code: The id of a token #Link to equation: The id of a token Copy LaTeX id ( τ ) = { stoi [ τ ] if τ ∈ stoi stoi [ <unk> ] otherwise \text{id}(\tau) = \begin{cases} \text{stoi}[\tau] & \text{if } \tau \in \text{stoi} \\ \text{stoi}[\,\texttt{<unk>}\,] & \text{otherwise} \end{cases} id ( τ ) = { stoi [ τ ] stoi [ <unk> ] if τ ∈ stoi otherwise Decoding joins the strings Show code: Decoding joins the strings #Link to equation: Decoding joins the strings Copy LaTeX decode ( i 1 , … , i n ) = itos [ i 1 ] ∥ ⋯ ∥ itos [ i n ] \text{decode}(i_1, \dots, i_n) = \text{itos}[i_1] \,\Vert\, \cdots \,\Vert\, \text{itos}[i_n] decode ( i 1 , … , i n ) = itos [ i 1 ] ∥ ⋯ ∥ itos [ i n ] File and function tokenizer.pyencodetokenizer.pydecodetokenizer.pyapply_mergesTensor shapes The ids of the corpus, as a list and then as an array tok.encode(text) (N,)ids (N,)Links Explanation Numerical example tokenizer.py:86-90 tokenizer.py:92-96 tokenizer.py:63-73 The vocabulary Implemented Direction Neither (not differentiated)Equation The order of the vocabulary Show code: The order of the vocabulary #Link to equation: The order of the vocabulary Copy LaTeX itos = [ <unk> ] + [ characters, first seen first ] + [ new merged tokens, in merge order ] \text{itos} = [\,\texttt{<unk>}\,] + [\,\text{characters, first seen first}\,] + [\,\text{new merged tokens, in merge order}\,] itos = [ <unk> ] + [ characters, first seen first ] + [ new merged tokens, in merge order ] File and function tokenizer.pybuild_tokenizerTensor shapes None declared Links Explanation Numerical example tokenizer.py:99-119 Pre-tokens and word counts Implemented Direction Neither (not differentiated)Equation How often the pre-token w occurs Show code: How often the pre-token w occurs #Link to equation: How often the pre-token w occurs Copy LaTeX n w = ∣ { i : p i = w } ∣ n_w = \left|\{\, i : p_i = w \,\}\right| n w = ∣ { i : p i = w } ∣ File and function tokenizer.pycount_wordsTensor shapes None declared Links Explanation Numerical example tokenizer.py:6-10 Applying a merge Implemented Direction Neither (not differentiated)Equation Replace an adjacent pair with one token Show code: Replace an adjacent pair with one token #Link to equation: Replace an adjacent pair with one token Copy LaTeX [ … , a ∗ , b ∗ , … ] → [ … , a ∗ b ∗ , … ] [\,\dots,\, a^*,\, b^*,\, \dots\,] \;\to\; [\,\dots,\, a^* b^*,\, \dots\,] [ … , a ∗ , b ∗ , … ] → [ … , a ∗ b ∗ , … ] File and function tokenizer.pymerge_tokensTensor shapes None declared Links Explanation Numerical example tokenizer.py:12-22 Pair frequency and the best pair Implemented Direction Neither (not differentiated)Equation Pair frequency, weighted by word count Show code: Pair frequency, weighted by word count #Link to equation: Pair frequency, weighted by word count Copy LaTeX f ( a , b ) = ∑ w n w ∣ { i : s i w = a , s i + 1 w = b } ∣ f(a,b) = \sum_{w} n_w \,\bigl|\{\, i : s^{w}_{i} = a,\; s^{w}_{i+1} = b \,\}\bigr| f ( a , b ) = w ∑ n w { i : s i w = a , s i + 1 w = b } The best pair Show code: The best pair #Link to equation: The best pair Copy LaTeX ( a ∗ , b ∗ ) = arg max ( a , b ) f ( a , b ) (a^*, b^*) = \operatorname*{arg\,max}_{(a,b)} f(a,b) ( a ∗ , b ∗ ) = ( a , b ) arg max f ( a , b ) File and function tokenizer.pytrain_bpetokenizer.pytrain_bpe_on_textTensor shapes None declared Links Explanation Numerical example tokenizer.py:24-61 tokenizer.py:76-77 Mini-batch loss and gradient Derived from implementation Direction Neither (not differentiated)Equation The batch size, read from this batch Show code: The batch size, read from this batch #Link to equation: The batch size, read from this batch Copy LaTeX B = number of rows of x b B = \text{number of rows of } x_b B = number of rows of x b How the code accumulates the loss Show code: How the code accumulates the loss #Link to equation: How the code accumulates the loss Copy LaTeX L ← L + L i B i = 1 , … , B , starting from L = 0 L \leftarrow L + \frac{L_i}{B} \qquad i = 1, \dots, B, \quad \text{starting from } L = 0 L ← L + B L i i = 1 , … , B , starting from L = 0 How the code accumulates each gradient Show code: How the code accumulates each gradient #Link to equation: How the code accumulates each gradient Copy LaTeX g ˉ k ← g ˉ k + g k , i B for every key k , starting from zeros \bar g_k \leftarrow \bar g_k + \frac{g_{k,i}}{B} \qquad \text{for every key } k, \quad \text{starting from zeros} g ˉ k ← g ˉ k + B g k , i for every key k , starting from zeros The result, the batch mean loss Show code: The result, the batch mean loss #Link to equation: The result, the batch mean loss Copy LaTeX L = 1 B ∑ i = 1 B L i L = \frac{1}{B} \sum_{i=1}^{B} L_i L = B 1 i = 1 ∑ B L i The result, the mean gradient Show code: The result, the mean gradient #Link to equation: The result, the mean gradient Copy LaTeX g ˉ k = 1 B ∑ i = 1 B g k , i \bar g_k = \frac{1}{B} \sum_{i=1}^{B} g_{k,i} g ˉ k = B 1 i = 1 ∑ B g k , i File and function train.pybatch_loss_and_gradsTensor shapes A batch goes through the model one window at a time xb (B, T)(16, 16)x (T,)(16,)logits inside loss_and_grads (T, V)(16, V)One key as an example. Averaging does not change a gradient's shape g["head_w"] (d, V)(32, V)grads["head_w"] (d, V)(32, V)Links Explanation Numerical example train.py:12-24 Step counts Derived from implementation Direction Neither (not differentiated)Equation Steps per epoch Show code: Steps per epoch #Link to equation: Steps per epoch Copy LaTeX S epoch = ⌈ W B ⌉ S_{\text{epoch}} = \left\lceil \frac{W}{B} \right\rceil S epoch = ⌈ B W ⌉ Total steps Show code: Total steps #Link to equation: Total steps Copy LaTeX S total = E ⋅ S epoch S_{\text{total}} = E \cdot S_{\text{epoch}} S total = E ⋅ S epoch File and function train.pymainTensor shapes None declared Links Explanation Numerical example train.py:36-81 Validation loss Derived from implementation Direction Neither (not differentiated)Equation Mean cross-entropy over validation windows Show code: Mean cross-entropy over validation windows #Link to equation: Mean cross-entropy over validation windows Copy LaTeX v = 1 W v ∑ j = 1 W v L j v = \frac{1}{W_v} \sum_{j=1}^{W_v} L_j v = W v 1 j = 1 ∑ W v L j File and function train.pyvalidation_lossTensor shapes Validation windows go through the model one at a time inputs (W, T)(W, 16)x (T,)(16,)logits (T, V)(16, V)Links Explanation Numerical example train.py:26-34 Attention backward Implemented Direction BackwardEquation Gradient on the output weights Show code: Gradient on the output weights #Link to equation: Gradient on the output weights Copy LaTeX d W o = O ⊤ d O u t dW_o = O^\top\, d\mathrm{Out} d W o = O ⊤ d Out Gradient on the output bias Show code: Gradient on the output bias #Link to equation: Gradient on the output bias Copy LaTeX d b o = 1 ⊤ d O u t db_o = \mathbf{1}^\top d\mathrm{Out} d b o = 1 ⊤ d Out Through the output projection, then the head split Show code: Through the output projection, then the head split #Link to equation: Through the output projection, then the head split Copy LaTeX d O h = split ( d O u t W o ⊤ , H ) dO_h = \operatorname{split}(d\mathrm{Out}\, W_o^\top, H) d O h = split ( d Out W o ⊤ , H ) Gradient on the values Show code: Gradient on the values #Link to equation: Gradient on the values Copy LaTeX d V h = A ⊤ d O h dV_h = A^\top dO_h d V h = A ⊤ d O h Gradient on the weights Show code: Gradient on the weights #Link to equation: Gradient on the weights Copy LaTeX d A = d O h V h ⊤ dA = dO_h V_h^\top d A = d O h V h ⊤ Through softmax, softmax_backward(A, dA) Show code: Through softmax, softmax_backward(A, dA) #Link to equation: Through softmax, softmax_backward(A, dA) Copy LaTeX d S i j = A i j ( d A i j − ∑ k A i k d A i k ) dS_{ij} = A_{ij}\Big(dA_{ij} - \sum_{k} A_{ik}\, dA_{ik}\Big) d S ij = A ij ( d A ij − k ∑ A ik d A ik ) A masked cell gets no gradient Derived from implementation Show code: A masked cell gets no gradient #Link to equation: A masked cell gets no gradient Copy LaTeX A i j = 0 ⇒ d S i j = 0 A_{ij} = 0 \quad\Rightarrow\quad dS_{ij} = 0 A ij = 0 ⇒ d S ij = 0 Gradient on the queries Show code: Gradient on the queries #Link to equation: Gradient on the queries Copy LaTeX d Q h = d S K h d h dQ_h = \frac{dS\, K_h}{\sqrt{d_h}} d Q h = d h d S K h Gradient on the keys Show code: Gradient on the keys #Link to equation: Gradient on the keys Copy LaTeX d K h = d S ⊤ Q h d h dK_h = \frac{dS^\top Q_h}{\sqrt{d_h}} d K h = d h d S ⊤ Q h Heads side by side again Derived from implementation Show code: Heads side by side again #Link to equation: Heads side by side again Copy LaTeX d Q , d K , d V = merge ( d Q h ) , merge ( d K h ) , merge ( d V h ) dQ,\ dK,\ dV = \operatorname{merge}(dQ_h),\ \operatorname{merge}(dK_h),\ \operatorname{merge}(dV_h) d Q , d K , d V = merge ( d Q h ) , merge ( d K h ) , merge ( d V h ) Gradient on the three weight matrices Show code: Gradient on the three weight matrices #Link to equation: Gradient on the three weight matrices Copy LaTeX d W n = X n ⊤ d M n , n ∈ { q , k , v } dW_n = X_n^\top dM_n, \qquad n \in \{q, k, v\} d W n = X n ⊤ d M n , n ∈ { q , k , v } Gradient on the three biases Show code: Gradient on the three biases #Link to equation: Gradient on the three biases Copy LaTeX d b n = 1 ⊤ d M n , n ∈ { q , k , v } db_n = \mathbf{1}^\top dM_n, \qquad n \in \{q, k, v\} d b n = 1 ⊤ d M n , n ∈ { q , k , v } Gradient on the input, summed over q, k and v Show code: Gradient on the input, summed over q, k and v #Link to equation: Gradient on the input, summed over q, k and v Copy LaTeX d X n = d Q W q ⊤ + d K W k ⊤ + d V W v ⊤ dX_n = dQ\, W_q^\top + dK\, W_k^\top + dV\, W_v^\top d X n = d Q W q ⊤ + d K W k ⊤ + d V W v ⊤ The key bias has a gradient of zero Show code: The key bias has a gradient of zero #Link to equation: The key bias has a gradient of zero Copy LaTeX d b k = 1 ⊤ d K = 0 db_k = \mathbf{1}^\top dK = 0 d b k = 1 ⊤ d K = 0 File and function transformer.pyattention_backwardTensor shapes From dOut to the values dOut (T, d)(16, 32)dOh (H, T, dh)(2, 16, 16)dVh (H, T, dh)(2, 16, 16)dV (T, d)(16, 32)From dOh through the weights and scores to the queries. Keys go the same way dOh (H, T, dh)(2, 16, 16)dA (H, T, T)(2, 16, 16)dS (H, T, T)(2, 16, 16)dQh (H, T, dh)(2, 16, 16)dQ (T, d)(16, 32)Parameter gradient and input gradient, for one projection dM (T, d)(16, 32)grads wn (d, d)(32, 32)dXn (T, d)(16, 32)Links Explanation Numerical example transformer.py:92-109 Residual connections, backward Implemented Direction BackwardEquation Through the feed-forward branch Show code: Through the feed-forward branch #Link to equation: Through the feed-forward branch Copy LaTeX d X n 2 = FFN back ( d Y ) dX_{n2} = \operatorname{FFN}^{\text{back}}(dY) d X n 2 = FFN back ( d Y ) Through the second LayerNorm Show code: Through the second LayerNorm #Link to equation: Through the second LayerNorm Copy LaTeX d X 1 branch = LN 2 back ( d X n 2 ) dX_1^{\text{branch}} = \operatorname{LN}_2^{\text{back}}(dX_{n2}) d X 1 branch = LN 2 back ( d X n 2 ) Second add, reversed Show code: Second add, reversed #Link to equation: Second add, reversed Copy LaTeX d X 1 = d Y + d X 1 branch dX_1 = dY + dX_1^{\text{branch}} d X 1 = d Y + d X 1 branch Through the attention branch Show code: Through the attention branch #Link to equation: Through the attention branch Copy LaTeX d X n = Attention back ( d X 1 ) dX_n = \operatorname{Attention}^{\text{back}}(dX_1) d X n = Attention back ( d X 1 ) Through the first LayerNorm Show code: Through the first LayerNorm #Link to equation: Through the first LayerNorm Copy LaTeX d X branch = LN 1 back ( d X n ) dX^{\text{branch}} = \operatorname{LN}_1^{\text{back}}(dX_n) d X branch = LN 1 back ( d X n ) First add, reversed Show code: First add, reversed #Link to equation: First add, reversed Copy LaTeX d X = d X 1 + d X branch dX = dX_1 + dX^{\text{branch}} d X = d X 1 + d X branch What each branch receives Derived from implementation Show code: What each branch receives #Link to equation: What each branch receives Copy LaTeX ∂ Y ∂ F f = I ⇒ d F f = d Y , ∂ X 1 ∂ A t t = I ⇒ d A t t = d X 1 \frac{\partial Y}{\partial \mathrm{Ff}} = I \;\Rightarrow\; d\mathrm{Ff} = dY, \qquad \frac{\partial X_1}{\partial \mathrm{Att}} = I \;\Rightarrow\; d\mathrm{Att} = dX_1 ∂ Ff ∂ Y = I ⇒ d Ff = d Y , ∂ Att ∂ X 1 = I ⇒ d Att = d X 1 Why the identity path carries dY unchanged Derived from implementation Show code: Why the identity path carries dY unchanged #Link to equation: Why the identity path carries dY unchanged Copy LaTeX Y = X 1 + F f ( X 1 ) ⇒ d X 1 = d Y ( I + ∂ F f ∂ X 1 ) = d Y + d Y ∂ F f ∂ X 1 Y = X_1 + \mathrm{Ff}(X_1) \;\Rightarrow\; dX_1 = dY \left( I + \frac{\partial \mathrm{Ff}}{\partial X_1} \right) = dY + dY \frac{\partial \mathrm{Ff}}{\partial X_1} Y = X 1 + Ff ( X 1 ) ⇒ d X 1 = d Y ( I + ∂ X 1 ∂ Ff ) = d Y + d Y ∂ X 1 ∂ Ff File and function transformer.pyblock_backwardTensor shapes Every gradient in the block has the same shape dY (T, d)(16, 32)dXn2 (T, d)(16, 32)dX1_branch (T, d)(16, 32)dX1 (T, d)(16, 32)dXn (T, d)(16, 32)dX_branch (T, d)(16, 32)dX (T, d)(16, 32)Links Explanation Numerical example transformer.py:141-152 Pre-LN block, forward Implemented Direction ForwardEquation First LayerNorm, on a copy Show code: First LayerNorm, on a copy #Link to equation: First LayerNorm, on a copy Copy LaTeX X n = LN 1 ( X ) X_n = \operatorname{LN}_1(X) X n = LN 1 ( X ) Attention reads the copy Show code: Attention reads the copy #Link to equation: Attention reads the copy Copy LaTeX A t t = Attention ( X n ) \mathrm{Att} = \operatorname{Attention}(X_n) Att = Attention ( X n ) First residual add Show code: First residual add #Link to equation: First residual add Copy LaTeX X 1 = X + A t t X_1 = X + \mathrm{Att} X 1 = X + Att Second LayerNorm, on a copy Show code: Second LayerNorm, on a copy #Link to equation: Second LayerNorm, on a copy Copy LaTeX X n 2 = LN 2 ( X 1 ) X_{n2} = \operatorname{LN}_2(X_1) X n 2 = LN 2 ( X 1 ) The feed-forward layer reads the copy Show code: The feed-forward layer reads the copy #Link to equation: The feed-forward layer reads the copy Copy LaTeX F f = FFN ( X n 2 ) \mathrm{Ff} = \operatorname{FFN}(X_{n2}) Ff = FFN ( X n 2 ) Second residual add Show code: Second residual add #Link to equation: Second residual add Copy LaTeX Y = X 1 + F f Y = X_1 + \mathrm{Ff} Y = X 1 + Ff The block on one line Show code: The block on one line #Link to equation: The block on one line Copy LaTeX X 1 = X + Attention ( LN 1 ( X ) ) , Y = X 1 + FFN ( LN 2 ( X 1 ) ) X_1 = X + \operatorname{Attention}(\operatorname{LN}_1(X)), \qquad Y = X_1 + \operatorname{FFN}(\operatorname{LN}_2(X_1)) X 1 = X + Attention ( LN 1 ( X )) , Y = X 1 + FFN ( LN 2 ( X 1 )) File and function transformer.pyblock_forwardTensor shapes Every array in the block has the same shape X (T, d)(16, 32)Xn (T, d)(16, 32)att (T, d)(16, 32)X1 (T, d)(16, 32)Xn2 (T, d)(16, 32)ff (T, d)(16, 32)Y (T, d)(16, 32)Links Explanation Numerical example transformer.py:131-139 Parameter count Implemented Direction Neither (not differentiated)Equation Scalars in one block Show code: Scalars in one block #Link to equation: Scalars in one block Copy LaTeX N block = 4 d 2 + 2 d F + F + 9 d N_{\text{block}} = 4d^2 + 2dF + F + 9d N block = 4 d 2 + 2 d F + F + 9 d Scalars in the whole model Show code: Scalars in the whole model #Link to equation: Scalars in the whole model Copy LaTeX N = ( 2 d + 1 ) V + ( T max + 2 ) d + L ( 4 d 2 + 2 d F + F + 9 d ) N = (2d+1)V + (T_{\max}+2)d + L\left(4d^2 + 2dF + F + 9d\right) N = ( 2 d + 1 ) V + ( T m a x + 2 ) d + L ( 4 d 2 + 2 d F + F + 9 d ) File and function transformer.pycount_paramstransformer.pyinit_paramsTensor shapes The two feed-forward weight matrices of one block ff1_w (d, F)(32, 128)ff2_w (F, d)(128, 32)Links Explanation Numerical example transformer.py:44-45 transformer.py:17-42 Embedding backward Implemented Direction BackwardEquation The chain rule for one entry of the token table Derived from implementation Show code: The chain rule for one entry of the token table #Link to equation: The chain rule for one entry of the token table Copy LaTeX ∂ L ∂ E v j = ∑ t d X t j ( 0 ) ∂ X t j ( 0 ) ∂ E v j , ∂ X t j ( 0 ) ∂ E v j = { 1 i t = v 0 otherwise \frac{\partial L}{\partial E_{vj}} = \sum_{t} dX^{(0)}_{tj}\, \frac{\partial X^{(0)}_{tj}}{\partial E_{vj}}, \qquad \frac{\partial X^{(0)}_{tj}}{\partial E_{vj}} = \begin{cases} 1 & i_t = v \\ 0 & \text{otherwise} \end{cases} ∂ E v j ∂ L = t ∑ d X t j ( 0 ) ∂ E v j ∂ X t j ( 0 ) , ∂ E v j ∂ X t j ( 0 ) = { 1 0 i t = v otherwise Gradient accumulation per token row Show code: Gradient accumulation per token row #Link to equation: Gradient accumulation per token row Copy LaTeX d E v = ∑ t : i t = v d X t ( 0 ) dE_v = \sum_{t\,:\, i_t = v} dX^{(0)}_t d E v = t : i t = v ∑ d X t ( 0 ) Gradient on the position table Show code: Gradient on the position table #Link to equation: Gradient on the position table Copy LaTeX d P t = { d X t ( 0 ) t < T 0 t ≥ T dP_t = \begin{cases} dX^{(0)}_t & t < T \\ 0 & t \geq T \end{cases} d P t = { d X t ( 0 ) 0 t < T t ≥ T File and function transformer.pyembed_backwardTensor shapes Into the token table. Rows of ids not in this window stay 0 dX0 (T, d)(16, 32)grads["tok_emb"] (V, d)(V, 32)Into the position table. Rows from T onward stay 0 dX0 (T, d)(16, 32)grads["pos_emb"] (Tmax, d)(32, 32)Links Explanation Numerical example transformer.py:55-57 Embedding sum Implemented Direction ForwardEquation One row of the sum Show code: One row of the sum #Link to equation: One row of the sum Copy LaTeX X t ( 0 ) = E i t + P t X^{(0)}_t = E_{i_t} + P_t X t ( 0 ) = E i t + P t The whole sum, as the code writes it Show code: The whole sum, as the code writes it #Link to equation: The whole sum, as the code writes it Copy LaTeX X ( 0 ) = E [ ids ] + P [ : T ] X^{(0)} = E[\text{ids}] + P[:T] X ( 0 ) = E [ ids ] + P [ : T ] File and function transformer.pyembed_forwardTensor shapes Both terms have the same shape, and so does their sum tok_emb[ids] (T, d)(16, 32)pos_emb[:T] (T, d)(16, 32)X0 (T, d)(16, 32)Links Explanation Numerical example transformer.py:52-53 Feed-forward network, backward Implemented Direction BackwardEquation Gradient on the ReLU output Show code: Gradient on the ReLU output #Link to equation: Gradient on the ReLU output Copy LaTeX d A = d O u t W 2 ⊤ dA = d\mathrm{Out}\, W_2^\top d A = d Out W 2 ⊤ Gradient on the second weights Show code: Gradient on the second weights #Link to equation: Gradient on the second weights Copy LaTeX d W 2 = A ⊤ d O u t dW_2 = A^\top d\mathrm{Out} d W 2 = A ⊤ d Out Gradient on the second bias Show code: Gradient on the second bias #Link to equation: Gradient on the second bias Copy LaTeX d b 2 = 1 ⊤ d O u t db_2 = \mathbf{1}^\top d\mathrm{Out} d b 2 = 1 ⊤ d Out Through ReLU Show code: Through ReLU #Link to equation: Through ReLU Copy LaTeX d Z = d A ⊙ 1 [ Z > 0 ] dZ = dA \odot \mathbb{1}[Z > 0] d Z = d A ⊙ 1 [ Z > 0 ] Gradient on the input Show code: Gradient on the input #Link to equation: Gradient on the input Copy LaTeX d X n = d Z W 1 ⊤ dX_n = dZ\, W_1^\top d X n = d Z W 1 ⊤ Gradient on the first weights Show code: Gradient on the first weights #Link to equation: Gradient on the first weights Copy LaTeX d W 1 = X n ⊤ d Z dW_1 = X_n^\top dZ d W 1 = X n ⊤ d Z Gradient on the first bias Show code: Gradient on the first bias #Link to equation: Gradient on the first bias Copy LaTeX d b 1 = 1 ⊤ d Z db_1 = \mathbf{1}^\top dZ d b 1 = 1 ⊤ d Z File and function transformer.pyffn_backwardTensor shapes Back through the three steps dOut (T, d)(16, 32)dA (T, F)(16, 128)dZ (T, F)(16, 128)dXn (T, d)(16, 32)Parameter gradients of the second layer dW2 (F, d)(128, 32)db2 (d,)(32,)Parameter gradients of the first layer dW1 (d, F)(32, 128)db1 (F,)(128,)Links Explanation Numerical example transformer.py:118-128 Feed-forward network, forward Implemented Direction ForwardEquation Widen Show code: Widen #Link to equation: Widen Copy LaTeX Z = X n W 1 + 1 b 1 ⊤ Z = X_n W_1 + \mathbf{1} b_1^\top Z = X n W 1 + 1 b 1 ⊤ ReLU, element by element Show code: ReLU, element by element #Link to equation: ReLU, element by element Copy LaTeX A = ReLU ( Z ) A = \operatorname{ReLU}(Z) A = ReLU ( Z ) Narrow Show code: Narrow #Link to equation: Narrow Copy LaTeX F f = A W 2 + 1 b 2 ⊤ \mathrm{Ff} = A W_2 + \mathbf{1} b_2^\top Ff = A W 2 + 1 b 2 ⊤ The same three steps for one row x Show code: The same three steps for one row x #Link to equation: The same three steps for one row x Copy LaTeX FFN ( x ) = ReLU ( x W 1 + b 1 ) W 2 + b 2 \operatorname{FFN}(x) = \operatorname{ReLU}(x W_1 + b_1)\, W_2 + b_2 FFN ( x ) = ReLU ( x W 1 + b 1 ) W 2 + b 2 File and function transformer.pyffn_forwardTensor shapes Widen, ReLU, narrow Xn (T, d)(16, 32)Z (T, F)(16, 128)A (T, F)(16, 128)out (T, d)(16, 32)Links Explanation Numerical example transformer.py:112-116 Whole model, forward Implemented Direction ForwardEquation The sequence must fit the position table Show code: The sequence must fit the position table #Link to equation: The sequence must fit the position table Copy LaTeX T ≤ T max T \le T_{\max} T ≤ T m a x Blocks in sequence Show code: Blocks in sequence #Link to equation: Blocks in sequence Copy LaTeX X ( l + 1 ) = Block l ( X ( l ) ) , l = 0 , … , L − 1 X^{(l+1)} = \operatorname{Block}_l\!\big(X^{(l)}\big), \qquad l = 0, \dots, L - 1 X ( l + 1 ) = Block l ( X ( l ) ) , l = 0 , … , L − 1 Final LayerNorm Show code: Final LayerNorm #Link to equation: Final LayerNorm Copy LaTeX X f = LN f ( X ( L ) ) X_f = \operatorname{LN}_f\!\big(X^{(L)}\big) X f = LN f ( X ( L ) ) The head Show code: The head #Link to equation: The head Copy LaTeX Z = X f W h + 1 b h ⊤ Z = X_f W_h + \mathbf{1} b_h^\top Z = X f W h + 1 b h ⊤ The last two steps together Show code: The last two steps together #Link to equation: The last two steps together Copy LaTeX L o g i t s = LN f ( X ( L ) ) W h + b h \mathrm{Logits} = \operatorname{LN}_f\!\big(X^{(L)}\big)\, W_h + b_h Logits = LN f ( X ( L ) ) W h + b h File and function transformer.pyforwardTensor shapes One sequence through the model ids (T,)(16,)X (T, d)(16, 32)Xf (T, d)(16, 32)logits (T, V)(16, V)Links Explanation Numerical example transformer.py:155-168 Starting values of the parameters Implemented Direction Neither (not differentiated)Equation Starting values, by kind of parameter Show code: Starting values, by kind of parameter #Link to equation: Starting values, by kind of parameter Copy LaTeX W ∼ Xavier uniform , b = 0 , γ = 1 , β = 0 W \sim \text{Xavier uniform},\quad b = \mathbf{0},\quad \gamma = \mathbf{1},\quad \beta = \mathbf{0} W ∼ Xavier uniform , b = 0 , γ = 1 , β = 0 File and function transformer.pyinit_paramsTensor shapes The head's weights and bias head_w (d, V)(32, V)head_b (V,)Links Explanation Numerical example transformer.py:17-42 Whole model, backward Implemented Direction BackwardEquation Order of calls, the head Show code: Order of calls, the head #Link to equation: Order of calls, the head Copy LaTeX d X f , d W h , d b h = Head back ( d Z ) dX_f,\ dW_h,\ db_h = \operatorname{Head}^{\text{back}}(dZ) d X f , d W h , d b h = Head back ( d Z ) Order of calls, the final LayerNorm Show code: Order of calls, the final LayerNorm #Link to equation: Order of calls, the final LayerNorm Copy LaTeX d X ( L ) , d γ f , d β f = LN f back ( d X f ) dX^{(L)},\ d\gamma_f,\ d\beta_f = \operatorname{LN}_f^{\text{back}}(dX_f) d X ( L ) , d γ f , d β f = LN f back ( d X f ) Order of calls, the blocks in reverse Show code: Order of calls, the blocks in reverse #Link to equation: Order of calls, the blocks in reverse Copy LaTeX d X ( l ) = Block l back ( d X ( l + 1 ) ) , l = L − 1 , … , 0 dX^{(l)} = \operatorname{Block}_l^{\text{back}}\!\big(dX^{(l+1)}\big), \qquad l = L - 1, \dots, 0 d X ( l ) = Block l back ( d X ( l + 1 ) ) , l = L − 1 , … , 0 Order of calls, the embeddings Show code: Order of calls, the embeddings #Link to equation: Order of calls, the embeddings Copy LaTeX d t o k _ e m b , d p o s _ e m b = Embed back ( d X ( 0 ) ) d\,\mathrm{tok\_emb},\ d\,\mathrm{pos\_emb} = \operatorname{Embed}^{\text{back}}\!\big(dX^{(0)}\big) d tok_emb , d pos_emb = Embed back ( d X ( 0 ) ) File and function transformer.pyloss_and_gradsTensor shapes The gradient from the logits to the first block's input dZ (T, V)(16, V)dXf (T, d)(16, 32)dX (T, d)(16, 32)Parameter gradients of the head head_w (d, V)(32, V)head_b (V,)Parameter gradients of the embedding tables tok_emb (V, d)(V, 32)pos_emb (Tmax, d)(32, 32)Links Explanation Numerical example transformer.py:171-190