Softmax & Cross-Entropy

Softmax and the loss look like one operation, but in this code they are separate. softmax_forward serves attention and sampling. cross_entropy_forward computes its own log-probabilities from the logits, and cross_entropy_backward has its own gradient.

Softmax

What is it?

SoftmaxImplemented

Basic operations

Turns a vector of scores into probabilities along the last axis. The code subtracts the largest score first, so the exponentials cannot overflow.

Why does the model need it?

Attention uses it to turn scores into weights, and sampling uses it to turn logits into a distribution over tokens.

Mathematics

Shift by the maximum, then exponentiate #

e_i = \exp\left(x_i - \max_k x_k\right)

Normalize #

p_i = \frac{e_i}{\sum_k e_k}

Subtracting the maximum does not change the result, because the same factor appears above and below the line.

Why the shift is safe. The factor from any constant c cancels #

\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}}

So the code is free to choose c. It chooses the largest score, which makes the largest exponent exactly 0 and every other exponent negative. No exponential can then overflow.

What does each symbol mean?

Symbol	Meaning	In code	Shape
$x$	Input scores	`x`	any; the last axis holds the scores
$c$	The constant subtracted. The code uses the largest score	`np.max(x, axis=-1, keepdims=True)`	same as x, with 1 on the last axis
$e$	Shifted exponentials	`e`	same as x
$p$	Probabilities, which sum to 1 along the last axis	not named in the code	same as x

Tiny numerical example

Illustrative ExampleHand-picked input. Values computed by running ops.softmax_forward; they are not learned values.

Input x: [2.0000, 1.0000, 0.0000]
Shifted exponentials e: [1.0000, 0.3679, 0.1353]
Probabilities: [0.6652, 0.2447, 0.0900]

Moving all scores together changes nothing, and large scores stay finiteIllustrative ExampleHand-picked scores. The last column computes the exponential of the largest score with no shift: at 1002 it does not fit in a 64-bit float, so normalising would give NaN.

Scores x	softmax_forward(x)	exp(max x), with no shift
`[2, 1, 0]`	`[0.6652, 0.2447, 0.0900]`	`7.3891`
`[-3, -4, -5]`	`[0.6652, 0.2447, 0.0900]`	`0.0498`
`[0, 1, 2]`	`[0.0900, 0.2447, 0.6652]`	`7.3891`
`[1000, 1001, 1002]`	`[0.0900, 0.2447, 0.6652]`	`Infinity`

Tensor shapes

Attention weightstransformer.py:83

S(H, T, T)(2, 16, 16)
A(H, T, T)(2, 16, 16)

H = num_heads = 2. T = seq_len = 16. T is at most seq_len: generation uses the last seq_len ids, or fewer for a short prompt.

Samplinggenerate.py:11

logits / temperature(V,)
p(V,)

V depends on the corpus. It is not a value in config.py.

Actual implementation

ops.pysoftmax_forwardGitHub

22def softmax_forward(x):23    e = np.exp(x - np.max(x, axis=-1, keepdims=True))24    return e / e.sum(axis=-1, keepdims=True)

Code and math connection

Equation	Role	Lines	Code
Shift by the maximum, then exponentiate	Computes it	ops.py:23	`e = np.exp(x - np.max(x, axis=-1, keepdims=True))`
Normalize	Computes it	ops.py:24	`return e / e.sum(axis=-1, keepdims=True)`
Normalize	Calls it	transformer.py:83	`A = softmax_forward(S)`
Normalize	Calls it	generate.py:11	`p = softmax_forward(logits / temperature)`
Why the shift is safe. The factor from any constant c cancels	Computes it	ops.py:23	`e = np.exp(x - np.max(x, axis=-1, keepdims=True))`

Where is it used?

Source

ops.py, lines 22-24, softmax_forward

PDF reference

Section 4.3, Eq. 23, p. 7

Used by

Called by: generate.sample_next, transformer.attention_forward

Tested by: tests.test_ops.test_softmax_backward_rows_sum_to_zero, tests.test_ops.test_softmax_values_and_stability

Written inline in ops.py, line 30. The same shift by the maximum, written inline before the log-softmax.

Calls it: transformer.py, line 83 (attention_forward), for Normalize
Calls it: generate.py, line 11 (sample_next), for Normalize

Status

Implemented

Implemented in

Equation	Role	File	Function	Lines	GitHub
Shift by the maximum, then exponentiate	Computes it	`ops.py`	`softmax_forward`	23	GitHub: ops.py, line 23
Normalize	Computes it	`ops.py`	`softmax_forward`	24	GitHub: ops.py, line 24
Why the shift is safe. The factor from any constant c cancels	Computes it	`ops.py`	`softmax_forward`	23	GitHub: ops.py, line 23

Softmax backward

What is it?

Softmax, backwardImplemented

Basic operations

Given dY, the gradient on softmax's output y, this computes the gradient on its input scores. It is one line, along the last axis.

Why does the model need it?

Attention weights come from softmax, so the backward pass must cross it to reach the scores, and from there the query and key parameters.

Mathematics

Each output depends on every input, so the derivative is a matrix.

The derivative of softmax. The code never builds this matrix #

\frac{\partial y_i}{\partial x_j} = y_i\left(\delta_{ij} - y_j\right)

The code never forms that matrix. Multiplying dY by it, and writing the sum once, gives the single line that the function holds.

The backward line, which multiplies dY by that derivative #

dx_j = y_j\left(dY_j - \sum_i y_i\, dY_i\right)

What does each symbol mean?

Symbol	Meaning	In code	Shape
$y$	Softmax output from the forward pass	`y`	same as x; the last axis holds the probabilities
$dY$	Gradient of the loss on y	`dY`	same as y
$dx$	Gradient on the input scores. The function returns it without naming it	`dS in attention_backward`	same as y
$\delta_{ij}$	1 when i equals j, otherwise 0	not named in the code	scalar

Tiny numerical example

Softmax backwardIllustrative ExampleHand-picked numbers from test_softmax_backward_rows_sum_to_zero, with two rows. Each row is computed on its own, because the sum is along the last axis. dY is chosen by hand.

Step	Code	Value
Softmax output	`y`	`[[0.6652, 0.2447, 0.0900], [0.0420, 0.8438, 0.1142]]`
Gradient on the output	`dY`	`[[1.0000, 0.0000, 0.0000], [0.0000, 2.0000, -1.0000]]`
Sum over each row	`np.sum(y * dY, axis=-1, keepdims=True)`	`[[0.6652], [1.5734]]`
Gradient on the scores	`y * (dY - ...)`	`[[0.2227, -0.1628, -0.0599], [-0.0661, 0.3600, -0.2939]]`

Tensor shapes

Attentiontransformer.py:100

dA(H, T, T)(2, 16, 16)
dS(H, T, T)(2, 16, 16)

H = num_heads = 2. T = seq_len = 16. T is at most seq_len: generation uses the last seq_len ids, or fewer for a short prompt.

Actual implementation

ops.pysoftmax_backwardGitHub

26def softmax_backward(y, dY):27    return y * (dY - np.sum(y * dY, axis=-1, keepdims=True))

Code and math connection

Equation	Role	Lines	Code
The derivative of softmax. The code never builds this matrix	Computes it	ops.py:27	`return y * (dY - np.sum(y * dY, axis=-1, keepdims=True))`
The backward line, which multiplies dY by that derivative	Computes it	ops.py:27	`return y * (dY - np.sum(y * dY, axis=-1, keepdims=True))`
The backward line, which multiplies dY by that derivative	Calls it	transformer.py:100	`dS = softmax_backward(c["A"], dA)`
The backward line, which multiplies dY by that derivative	Tests it	tests/test_ops.py:40	`dz = softmax_backward(p, dp)`

Where is it used?

Only attention calls it. The loss and sampling do not.

Source

ops.py, lines 26-27, softmax_backward

PDF reference

Section 4.3, Eq. 24, 25, p. 7

Used by

Called by: transformer.attention_backward

Tested by: tests.test_ops.test_softmax_backward_rows_sum_to_zero

Calls it: transformer.py, line 100 (attention_backward), for The backward line, which multiplies dY by that derivative
Tests it: tests/test_ops.py, line 40 (test_softmax_backward_rows_sum_to_zero), for The backward line, which multiplies dY by that derivative

Status

Implemented

Implemented in

Equation	Role	File	Function	Lines	GitHub
The derivative of softmax. The code never builds this matrix	Computes it	`ops.py`	`softmax_backward`	27	GitHub: ops.py, line 27
The backward line, which multiplies dY by that derivative	Computes it	`ops.py`	`softmax_backward`	27	GitHub: ops.py, line 27

How this reaches a parameter update. Softmax has no parameters, so nothing is stored here. dS (transformer.py:100) goes on to dQh and dKh (transformer.py:101), and from those to the gradients of the query and key weights (transformer.py:106-108). batch_loss_and_grads averages them (train.py:61), clip_grad_norm may scale them (train.py:62), and opt.step uses them (train.py:66). Attention covers those steps.

Cross-entropy

What is it?

Cross-entropyImplemented

Loss

The loss. For each position it takes the log-probability the model gave to the correct next token, negates it, and averages over positions. It also returns the probabilities, which the backward pass reuses.

Why does the model need it?

Training needs one number to reduce. This loss falls when the model gives more probability to the correct token. The code works in log-probabilities, so it never takes the log of a number that has underflowed to 0.

Mathematics

The same shift as softmax, written inline #

z_{tv} = \text{logits}_{tv} - \max_{k} \text{logits}_{tk}

Log-softmax #

\log p_{tv} = z_{tv} - \log \sum_{k} e^{z_{tk}}

Mean over positions #

L = -\frac{1}{T} \sum_{t} \log p_{t,\, y_t}

The probabilities it returns #

P_{tv} = e^{\log p_{tv}}

The mean is over the positions of one sequence. batch_loss_and_grads then averages over the windows of a batch (train.py:17).

What does each symbol mean?

Symbol	Meaning	In code	Shape
$\text{logits}$	The head's output	`logits`	(T, V)
$y_t$	The correct next token id at position t	`targets`	(T,)
$z$	Logits with each row's largest value subtracted	`z`	(T, V)
$\log p$	Log-probabilities	`logp`	(T, V)
$L$	The loss	`loss`	scalar
$P$	Probabilities	`probs in the caller`	(T, V)
$T$	Positions in this sequence	`T`	scalar

Tiny numerical example

Two positions, three tokensIllustrative ExampleHand-picked logits and targets. They are not model output.

Step	Code	Value
Logits	`logits`	`[[2.0000, 1.0000, 0.0000], [0.5000, 0.5000, 3.0000]]`
Correct token at each position	`targets`	`[0.0000, 2.0000]`
Shifted	`logits - logits.max(axis=-1, keepdims=True)`	`[[0.0000, -1.0000, -2.0000], [-2.5000, -2.5000, 0.0000]]`
Log-probabilities	`z - np.log(np.exp(z).sum(axis=-1, keepdims=True))`	`[[-0.4076, -1.4076, -2.4076], [-2.6520, -2.6520, -0.1520]]`
Picked at the targets	`logp[np.arange(T), targets]`	`[-0.4076, -0.1520]`
Loss	`-(...).mean()`	`0.2798`
Probabilities returned	`np.exp(logp)`	`[[0.6652, 0.2447, 0.0900], [0.0705, 0.0705, 0.8590]]`

Why the code keeps log-probabilitiesIllustrative ExampleHand-picked logits from test_cross_entropy_extreme_logits_no_nan. The target's probability underflows to exactly 0, and the log of 0 is minus infinity. The code never takes that log, so its loss is finite.

Step	Code	Value
Logits	`logits`	`[[1000.0000, 0.0000, -1000.0000]]`
Correct token	`targets`	`[2.0000]`
Probabilities	`np.exp(logp)`	`[[1.0000, 0.0000, 0.0000]]`
Loss	`loss`	`2000.0000`

Tensor shapes

targets has shape (T,), and the loss is one number.

From logits to probabilities. The loss is a single numbertransformer.py:173

logits(T, V)(16, V)
logp(T, V)(16, V)
probs(T, V)(16, V)

T = seq_len = 16. T is at most seq_len: generation uses the last seq_len ids, or fewer for a short prompt. V depends on the corpus. It is not a value in config.py.

Actual implementation

ops.pycross_entropy_forwardGitHub

29def cross_entropy_forward(logits, targets):30    z = logits - logits.max(axis=-1, keepdims=True)31    logp = z - np.log(np.exp(z).sum(axis=-1, keepdims=True))32    T = len(targets)33    loss = -logp[np.arange(T), targets].mean()34    return float(loss), np.exp(logp)

Code and math connection

Equation	Role	Lines	Code
The same shift as softmax, written inline	Computes it	ops.py:30	`z = logits - logits.max(axis=-1, keepdims=True)`
Log-softmax	Computes it	ops.py:31	`logp = z - np.log(np.exp(z).sum(axis=-1, keepdims=True))`
Mean over positions	Computes it	ops.py:33	`loss = -logp[np.arange(T), targets].mean()`
Mean over positions	Calls it	transformer.py:173	`loss, probs = cross_entropy_forward(logits, targets)`
Mean over positions	Calls it	train.py:31	`loss, _ = cross_entropy_forward(logits, y)`
Mean over positions	Tests it	tests/test_ops.py:48	`loss, probs = cross_entropy_forward(logits, targets)`
The probabilities it returns	Computes it	ops.py:34	`return float(loss), np.exp(logp)`

Where is it used?

loss_and_grads keeps both results. validation_loss keeps only the loss.

Source

ops.py, lines 29-34, cross_entropy_forward

PDF reference

Section 4.3, Eq. 26, 28, p. 7

Used by

Called by: train.validation_loss, transformer.loss_and_grads

Tested by: tests.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_example

Calls it: transformer.py, line 173 (loss_and_grads), for Mean over positions
Calls it: train.py, line 31 (validation_loss), for Mean over positions
Tests it: tests/test_ops.py, line 48 (test_cross_entropy_worked_example), for Mean over positions

Status

Implemented

Implemented in

Equation	Role	File	Function	Lines	GitHub
The same shift as softmax, written inline	Computes it	`ops.py`	`cross_entropy_forward`	30	GitHub: ops.py, line 30
Log-softmax	Computes it	`ops.py`	`cross_entropy_forward`	31	GitHub: ops.py, line 31
Mean over positions	Computes it	`ops.py`	`cross_entropy_forward`	33	GitHub: ops.py, line 33
The probabilities it returns	Computes it	`ops.py`	`cross_entropy_forward`	34	GitHub: ops.py, line 34

Cross-entropy backward

What is it?

Cross-entropy, backwardImplemented

Loss

Computes dZ, the gradient of the loss on the logits. It copies the probabilities, subtracts 1 at each position's correct token, and divides by T.

Why does the model need it?

This is where every gradient starts. dZ goes straight into the backward pass of the head, and the rest of the backward pass follows from there.

Mathematics

The derivative of the loss at one position #

\frac{\partial \ell_t}{\partial \text{logits}_{tv}} = P_{tv} - \mathbf{1}[v = y_t]

The loss is the mean over T positions, and each position depends only on its own row of logits. So each row is divided by T.

The combined backward, where E_y is one-hot at each correct token #

dZ = \frac{P - E_y}{T}

What does each symbol mean?

Symbol	Meaning	In code	Shape
$P$	The probabilities cross_entropy_forward returned	`probs`	(T, V)
$E_y$	One-hot rows, with 1 at each position's correct token. The code subtracts in place and never builds this array	not named in the code	(T, V)
$\ell_t$	The loss at one position, before the mean	not named in the code	scalar
$dZ$	Gradient of the loss on the logits	`dZ`	(T, V)
$T$	Positions in this sequence	`T`	scalar

Tiny numerical example

A worked backwardIllustrative ExampleThe same hand-picked logits and targets as the cross-entropy example above.

Step	Code	Value
Probabilities	`dZ = probs.copy()`	`[[0.6652, 0.2447, 0.0900], [0.0705, 0.0705, 0.8590]]`
Subtract 1 at each correct token	`dZ[np.arange(T), targets] -= 1.0`	`[[-0.3348, 0.2447, 0.0900], [0.0705, 0.0705, -0.1410]]`
Divide by T	`return dZ / T`	`[[-0.1674, 0.1224, 0.0450], [0.0353, 0.0353, -0.0705]]`

By hand, at position 0 and its correct token: (0.6652 - 1) / 2 = -0.1674. In this row the entry for the correct token is negative and the others are positive: raising the correct logit lowers the loss.

Checking this backward numericallyIllustrative ExampleA test-time idea, not part of the pipeline: gradcheck.numerical_grad is only called by tests. It nudges each number up and down by h = 1e-5 and measures how the loss changes. Tests compare the central differences here with that real function.

Gradient	From the backward formula	From central differences	Largest gap
`dZ`	`[-0.1674, 0.1224, 0.0450][0.0353, 0.0353, -0.0705]`	`[-0.1674, 0.1224, 0.0450][0.0353, 0.0353, -0.0705]`	`2.0e-12`

Tensor shapes

The start of the backward passtransformer.py:175

probs(T, V)(16, V)
dZ(T, V)(16, V)

T = seq_len = 16. T is at most seq_len: generation uses the last seq_len ids, or fewer for a short prompt. V depends on the corpus. It is not a value in config.py.

Actual implementation

ops.pycross_entropy_backwardGitHub

36def cross_entropy_backward(probs, targets):37    T = len(targets)38    dZ = probs.copy()39    dZ[np.arange(T), targets] -= 1.040    return dZ / T

Code and math connection

Equation	Role	Lines	Code
The derivative of the loss at one position	Computes it	ops.py:39	`dZ[np.arange(T), targets] -= 1.0`
The combined backward, where E_y is one-hot at each correct token	Computes it	ops.py:38	`dZ = probs.copy()`
The combined backward, where E_y is one-hot at each correct token	Computes it	ops.py:39	`dZ[np.arange(T), targets] -= 1.0`
The combined backward, where E_y is one-hot at each correct token	Computes it	ops.py:40	`return dZ / T`
The combined backward, where E_y is one-hot at each correct token	Calls it	transformer.py:175	`dZ = cross_entropy_backward(probs, targets)`
The combined backward, where E_y is one-hot at each correct token	Tests it	tests/test_ops.py:59	`dZ = cross_entropy_backward(probs, targets)`

Where is it used?

Source

ops.py, lines 36-40, cross_entropy_backward

PDF reference

Section 4.3, Eq. 27, 28, p. 7

Used by

Called by: transformer.loss_and_grads

Tested by: tests.test_ops.test_cross_entropy_is_mean_over_positions, tests.test_ops.test_cross_entropy_worked_example

Calls it: transformer.py, line 175 (loss_and_grads), for The combined backward, where E_y is one-hot at each correct token
Tests it: tests/test_ops.py, line 59 (test_cross_entropy_is_mean_over_positions), for The combined backward, where E_y is one-hot at each correct token

Status

Implemented

Implemented in

Equation	Role	File	Function	Lines	GitHub
The derivative of the loss at one position	Computes it	`ops.py`	`cross_entropy_backward`	39	GitHub: ops.py, line 39
The combined backward, where E_y is one-hot at each correct token	Computes it	`ops.py`	`cross_entropy_backward`	38	GitHub: ops.py, line 38
The combined backward, where E_y is one-hot at each correct token	Computes it	`ops.py`	`cross_entropy_backward`	39	GitHub: ops.py, line 39
The combined backward, where E_y is one-hot at each correct token	Computes it	`ops.py`	`cross_entropy_backward`	40	GitHub: ops.py, line 40

How this reaches a parameter update. The loss has no parameters, so nothing is stored here. dZ (transformer.py:175) goes straight into linear_backward for the head (transformer.py:176), whose results fill head_w and head_b (transformer.py:178). The rest of the backward pass follows from dXf. batch_loss_and_grads averages the gradients (train.py:61), clip_grad_norm may scale them (train.py:62), and opt.step uses them (train.py:66).

Softmax lab

Edit the logits, or choose another correct token. Sampling covers temperature, which this lab leaves out.

Logits to probabilities, loss and gradientIllustrative ExampleHand-picked logits that you can edit. They are not model output. With one position, T is 1, so dZ is P minus the one-hot row.

logits
	token 0	token 1	token 2
t0

Correct token

Loss: 0.4076

P = exp(logp), shape (1, 3)
0.6652	0.2447	0.09

dZ = (P - one-hot) / T, shape (1, 3)
-0.3348	0.2447	0.09

Probabilities from these logitsIllustrative ExampleThe correct token is token 0. The loss falls as its bar grows.

token 00.6652

token 10.2447

token 20.0900

Data table

Label	Probability
token 0	0.6652
token 1	0.2447
token 2	0.0900

Perplexity

What is it?

PerplexityImplemented

Loss

The exponential of the loss. train.py prints it once per epoch, next to the validation loss, and nothing else reads it.

Why does the model need it?

The model does not need it: it is not part of the loss or the update. It is a printed number that is easier to read than a log value. In general, it can be read as the number of equally likely tokens that would give the same loss.

Mathematics

Perplexity #

\text{perplexity} = e^{L}

What does each symbol mean?

Symbol	Meaning	In code	Shape
$L$	A loss. train.main passes the validation loss	`loss in perplexity, v in train.main`	scalar

Tiny numerical example

Loss and perplexityIllustrative ExampleHand-picked losses, not losses from a training run.

Loss	Where it comes from	Perplexity
0.2798	The loss of the cross-entropy example above	1.3229
1.6094	ln 5, the loss of uniform probabilities over 5 tokens	5.0000
5.7038	ln 300, the value test_perplexity_of_uniform_is_vocab_size uses	300.0000

Tensor shapes

It takes one number and returns one number, so there are no tensors.

Actual implementation

ops.pyperplexityGitHub

59def perplexity(loss):60    return float(np.exp(loss))

Code and math connection

Equation	Role	Lines	Code
Perplexity	Computes it	ops.py:60	`return float(np.exp(loss))`
Perplexity	Calls it	train.py:71	`print(f"epoch {epoch + 1}: val loss {v:.3f} ppl {perplexity(v):.1f}")`
Perplexity	Tests it	tests/test_ops.py:88	`assert np.isclose(perplexity(np.log(300)), 300)`

Where is it used?

train.main computes the validation loss after each epoch, and prints it with its perplexity:

train.pymainGitHub

70        v = validation_loss(params, cfg, va_x, va_y)71        print(f"epoch {epoch + 1}: val loss {v:.3f}  ppl {perplexity(v):.1f}")

Source

ops.py, lines 59-60, perplexity

PDF reference

Section 4.3, p. 7

Used by

Called by: train.main

Tested by: tests.test_ops.test_perplexity_of_uniform_is_vocab_size

Calls it: train.py, line 71 (main), for Perplexity
Tests it: tests/test_ops.py, line 88 (test_perplexity_of_uniform_is_vocab_size), for Perplexity

Status

Implemented

Implemented in

Equation	Role	File	Function	Lines	GitHub
Perplexity	Computes it	`ops.py`	`perplexity`	60	GitHub: ops.py, line 60