Backprop Through a Dense Layer

This is the single most reused derivation in deep learning. Every framework's Linear layer runs exactly these three lines in its backward pass.

The forward pass is Y = X @ W + b with

• X of shape (n, d_in)
• W of shape (d_in, d_out)
• b of shape (d_out,), broadcast over the batch
• Y of shape (n, d_out)

You are handed grad_out = dL/dY of shape (n, d_out). Return (grad_x, grad_w, grad_b) matching the shapes of X, W and b.

A useful sanity check while deriving: each gradient must have the same shape as the thing it is a gradient of, and there is only one way to arrange each matrix product to make the shapes work.

Because b is broadcast across the batch, its gradient is a sum over the batch, not a mean.

Function signature

def dense_backward(x, w, b, grad_out):

Example

Input

x = np.array([[1.0, 2.0]])            # (1, 2)
w = np.array([[1.0], [3.0]])          # (2, 1)
b = np.array([0.0])
grad_out = np.array([[1.0]])
dense_backward(x, w, b, grad_out)

Output

(array([[1., 3.]]), array([[1.], [2.]]), array([1.]))

Reasoning

grad_x = grad_out @ w.T = [[1, 3]], grad_w = x.T @ grad_out = [[1], [2]], and grad_b sums grad_out over the batch, which gives [1]. Each result has the shape of the thing it is the gradient of.

Constraints and notes

• grad_b is a sum over the batch, not a mean.
• Each gradient has the same shape as its input.