Gradients for linear layers and softmax

Real layers process a batch of vectors at once, so backprop works with matrices. A linear layer has a simple matrix form for its gradients. A softmax followed by cross-entropy is even simpler: the gradient is the predicted probabilities minus the one-hot label.

The linear layer, with shapes

For a batch X of shape (B, D), weights W of shape (H, D) and bias b of shape (H):

Z = X Wᵀ + b          # (B, H)

Given the incoming gradient G = ∂L/∂Z, the three gradients are:

∂L/∂W = Gᵀ X          # (H, D), same shape as W
∂L/∂b = sum(G, axis=0) # (H), sum over the batch
∂L/∂X = G W           # (B, D), same shape as X

Each formula is the chain rule applied to z_bh = Σ_d x_bd · W_hd.

X = [[1, 2, 0],      W = [[0.5,  0, -1],
     [0, 1, 1]]           [1,   -1,  0]]

The shapes are the main check. A gradient always matches the shape of the thing it is the gradient of.

Softmax with cross-entropy

For one example with logits z and true class y:

p_i = exp(z_i) / Σ_j exp(z_j)
L   = −log p_y

Applying the chain rule through softmax and the log collapses to one line:

∂L/∂z_i = p_i − 1[i = y]

The gradient is the probabilities minus the one-hot label. The correct class is pushed up by p_y − 1, and every other class is pushed down by its probability.

import numpy as np

def softmax_ce_backward(logits, y):
    z = logits - logits.max()           # subtract the max for stability
    p = np.exp(z)
    p /= p.sum()
    loss = -np.log(p[y])
    grad = p.copy()
    grad[y] -= 1                        # p minus the one-hot label
    return loss, grad

loss, grad = softmax_ce_backward(np.array([2.0, 1.0, 0.1]), y=0)
print("loss", round(loss, 3), "grad", np.round(grad, 3))

Output:

loss 0.417 grad [-0.341  0.242  0.099]

The correct class has a negative gradient, so it moves up. The other two move down. Frameworks fuse softmax and cross-entropy for this reason: the forward is numerically stable, and the backward is trivial.

Common mistakes

• Forgetting to sum the bias gradient over the batch. The bias is broadcast across rows in the forward pass, so its gradient must be summed over those rows. If you skip the sum, the bias gradient has the wrong shape.
• Swapping Gᵀ X and G W. The weight gradient uses the transpose of G on the left. The input gradient uses G as is.
• Computing softmax and the log separately. The result is correct, but the numbers can overflow. Use a fused, stable cross-entropy from your framework.

Check yourself

1. For W of shape (H, D), what shape is ∂L/∂W? (H, D), the same as W.
2. The model gives the correct class probability 0.99. About how big is the gradient on the correct logit? About -0.01, because p − 1 = 0.99 − 1. The model has little left to learn from that example.

Key takeaways

• For a linear layer, ∂L/∂W = Gᵀ X, ∂L/∂X = G W and ∂L/∂b sums G over the batch.
• Gradients always match the shape of the thing they differentiate.
• Softmax with cross-entropy has the gradient p − onehot.
• Use fused, stable loss functions rather than building softmax and log by hand.

Further learning

See all of this at work on one small network: a worked training step.