Samples a (fan_in, fan_out) matrix from a uniform distribution
- Reached from
train.main- Used by, in the pipeline
transformer.init_params- Used by, in tests
tests.test_ops.test_xavier_bound- Mathematics
- Xavier uniform initialisation
Differentiable operations with hand-written backward functions
11 of 12 functions, classes and methods here are reachable from train.main, and 3 from talk.main.
| Pipeline stage | Step | Explained on |
|---|---|---|
| Blocks | LayerNorm 1 | layernorm |
| Blocks | LayerNorm 2 | layernorm |
| Output | Final LayerNorm | full model |
| Loss and update | Softmax and cross-entropy | softmax cross entropy |
| Loss and update | Backward | how the model learns |
numpygenerate.py, tests/test_ops.py, train.py, transformer.pygenerate.py, train.py, transformer.pytests/test_ops.py| Concept | What this file does | Lines | Status |
|---|---|---|---|
| Xavier uniform initialisation | Computes it | 4, 5 | Implemented |
| Matrix multiplication | Computes it | 8 | Derived from implementation |
| Linear layer | Computes it | 8 | Implemented |
| Linear layer, backward | Computes it | 11, 12, 13 | Implemented |
| ReLU | Computes it | 17 | Implemented |
| ReLU, backward | Computes it | 20 | Implemented |
| Softmax per row | Computes it | 23, 24 | Implemented |
| Softmax | Computes it | 23, 24 | Implemented |
| Temperature | Calls it | 24 | Implemented |
| Softmax, backward | Computes it | 27 | Implemented |
| Attention backward | Computes it | 27 | Implemented |
| Cross-entropy | Computes it | 30, 31, 33, 34 | Implemented |
| Cross-entropy, backward | Computes it | 38, 39, 40 | Implemented |
| Deriving the backward formula | Computes it | 43, 44, 45, 46, 51, 52, 53, 55, 56 | Derived from implementation |
| LayerNorm | Computes it | 43, 44, 45, 46, 47 | Implemented |
| LayerNorm, backward | Computes it | 51, 52, 53, 55, 56 | Implemented |
| Perplexity | Computes it | 60 | Implemented |
Samples a (fan_in, fan_out) matrix from a uniform distribution
train.maintransformer.init_paramstests.test_ops.test_xavier_boundReturns X @ W + b
tests.test_ops.test_linear_worked_exampleReturns dX, dW, db
train.maintransformer.ffn_backward, transformer.loss_and_gradstests.test_ops.test_linear_worked_exampledef relu_forward(X)Lines 16-17GitHub: ops.relu_forwardElement-wise max(0, X)
train.main, talk.maintransformer.ffn_forwardtests.test_ops.test_reluKeeps dA where Z > 0
train.maintransformer.ffn_backwardtests.test_ops.test_reluSoftmax on the last axis, with the maximum subtracted first
train.main, talk.maingenerate.sample_next, transformer.attention_forwardtests.test_ops.test_softmax_backward_rows_sum_to_zero, tests.test_ops.test_softmax_values_and_stabilitySoftmax backward product
train.maintransformer.attention_backwardtests.test_ops.test_softmax_backward_rows_sum_to_zeroReturns the mean cross-entropy and the probabilities
train.maintrain.validation_loss, transformer.loss_and_gradstests.test_ops.test_cross_entropy_extreme_logits_no_nan, tests.test_ops.test_cross_entropy_is_mean_over_positions, tests.test_ops.test_cross_entropy_worked_exampleReturns the gradient with respect to the logits
train.maintransformer.loss_and_gradstests.test_ops.test_cross_entropy_is_mean_over_positions, tests.test_ops.test_cross_entropy_worked_exampleLayer norm; returns the output and cache (Xhat, s)
train.main, talk.maintransformer.block_forward, transformer.forwardtests.test_ops.test_layer_norm_worked_exampleReturns dX, dgamma, dbeta
train.maintransformer.block_backward, transformer.loss_and_gradstests.test_ops.test_layer_norm_worked_exampledef perplexity(loss)Lines 59-60GitHub: ops.perplexityReturns exp(loss)
train.maintrain.maintests.test_ops.test_perplexity_of_uniform_is_vocab_sizeimport numpy as npdef xavier_uniform(fan_in, fan_out, rng): a = np.sqrt(6.0 / (fan_in + fan_out)) return rng.uniform(-a, a, size=(fan_in, fan_out))def linear_forward(X, W, b): return X @ W + bdef linear_backward(X, W, dY): dW = X.T @ dY db = dY.sum(axis=0) dX = dY @ W.T return dX, dW, dbdef relu_forward(X): return np.maximum(0.0, X)def relu_backward(Z, dA): return dA * (Z > 0)def softmax_forward(x): e = np.exp(x - np.max(x, axis=-1, keepdims=True)) return e / e.sum(axis=-1, keepdims=True)def softmax_backward(y, dY): return y * (dY - np.sum(y * dY, axis=-1, keepdims=True))def cross_entropy_forward(logits, targets): z = logits - logits.max(axis=-1, keepdims=True) logp = z - np.log(np.exp(z).sum(axis=-1, keepdims=True)) T = len(targets) loss = -logp[np.arange(T), targets].mean() return float(loss), np.exp(logp) def cross_entropy_backward(probs, targets): T = len(targets) dZ = probs.copy() dZ[np.arange(T), targets] -= 1.0 return dZ / Tdef layer_norm_forward(X, gamma, beta, epsilon = 1e-5): mu = X.mean(axis=-1, keepdims=True) var = ((X - mu) ** 2).mean(axis=-1, keepdims=True) s = 1.0 / np.sqrt(var + epsilon) Xhat = (X - mu) * s return gamma * Xhat + beta, (Xhat, s)def layer_norm_backward(dY, cache, gamma): Xhat, s = cache dXhat = dY * gamma dX = s * (dXhat - dXhat.mean(axis=-1, keepdims=True) - Xhat * (dXhat * Xhat).mean(axis=-1, keepdims=True)) dgamma = (dY * Xhat).sum(axis=0) dbeta = dY.sum(axis=0) return dX, dgamma, dbetadef perplexity(loss): return float(np.exp(loss))