\( \newcommand{\vect}[1]{\boldsymbol{\mathbf{#1}}} \newcommand{\mat}[1]{\boldsymbol{\mathbf{#1}}} \newcommand{\R}{\mathbb{R}} \newcommand{\E}{\mathbb{E}} \newcommand{\Prob}{\mathbb{P}} \newcommand{\T}{^{\top}} \newcommand{\norm}[1]{\lVert #1 \rVert} \newcommand{\abs}[1]{\lvert #1 \rvert} \newcommand{\ihat}{\mkern1mu\hat{\boldsymbol{\imath}}\mkern2mu} \newcommand{\jhat}{\mkern1mu\hat{\boldsymbol{\jmath}}\mkern2mu} \newcommand{\pd}[2]{\frac{\partial #1}{\partial #2}} \newcommand{\argmax}{\operatorname*{arg\,max}} \newcommand{\argmin}{\operatorname*{arg\,min}} \newcommand{\cdata}[1]{\class{c-data}{#1}} \newcommand{\cmodel}[1]{\class{c-model}{#1}} \newcommand{\ctruth}[1]{\class{c-truth}{#1}} \newcommand{\cgrad}[1]{\class{c-grad}{#1}} \newcommand{\closs}[1]{\class{c-loss}{#1}} \newcommand{\cparam}[1]{\class{c-param}{#1}} \newcommand{\csecond}[1]{\class{c-second}{#1}} \renewcommand{\,}{\mkern3mu} \)

Matrix products, reductions and broadcasting

The backward rules that make a network possible: matrix products read off from shapes, sums that spread the gradient back, reshapes and transposes undone, broadcasting summed away by an unbroadcast helper, and a fused cross-entropy.

DL01.6.2L2 Intermediate25 min

In review (what this means)

The maths and code on this page are checked by automated tests on every change, but an independent reader has not signed it off yet. Spot a mistake or a confusing step? Tell us.

Before this lesson: DL01.6.1 From scalars to tensors

What you will be able to do

• Derive the backward rules of a matrix product, G B-transpose and A-transpose G, and check them against the shapes.
• Explain why the backward pass of broadcasting sums over every stretched axis, and implement it as an unbroadcast helper.
• Write the backward rules of sum, mean, reshape and transpose.
• Build a stable log-softmax from the engine’s operations and a fused cross-entropy, and say why fusing is worth doing.
• Check every rule against central differences on random arrays, with the shapes asserted first.

The question

A layer computes \({Z = XW\T + \vect{1}\vect{b}\T}\): one matrix product and one bias added to every row. In NumPy that is X @ W.T + b, and the bias of shape (3,) is quietly copied onto every row of a (64, 3) result. The engine of DL01.6.1 can do neither: its operations demand equal shapes and keep every entry where it is. What are the backward rules of operations that mix entries and change shapes?

From the last lesson

1. What does the backward closure of an elementwise operation \({v = g(u)}\) add to u.grad?

Show answer

\({g'(u) \odot \delta_v}\): the vector-Jacobian product with a diagonal Jacobian, computed entry by entry.

2. Why are the closures called in reverse topological order?

Show answer

So that a node passes its gradient back only after every node that uses it has reported; then what it passes back is complete.

Intuition

Broadcasting is copying. Adding a column of shape (3, 1) to a row of shape (1, 4) behaves as if the column were copied into four columns and the row into three rows, then the two (3, 4) grids added. A copied entry is a value used several times, and since M05.8.3 we know what that means backward: its uses add up.

DL01.6.2-F1

Broadcasting forward and backward. Each entry of \(a\) was copied into four places, so it collects four contributions; each entry of \(b\) was copied into three. The gradient ends with the shape of the array it belongs to.

So the backward rule of broadcasting is a sum over every axis that was stretched. The same habit of thought, “where did each entry go, and how many times?”, gives every other rule of this lesson.

Build it up

Step 1. A matrix product mixes entries. For \({C = AB}\) with \(A\) of shape \({(n, k)}\) and \(B\) of shape \({(k, m)}\), the gradient \(G\) arriving at \(C\) has shape \({(n, m)}\). A’s gradient must have shape \({(n, k)}\), and the only product of \(G\) with \(B\) or \({B\T}\) that has it is \({GB\T}\). Likewise \({A\T G}\) is the only one with the shape of \(B\). Lesson M05.7.3 derived both; the derivation is repeated below.

DL01.6.2-F2

The two backward rules of a matrix product, drawn by shape. In each line the inner sizes match and the outer sizes are those of the array whose gradient is wanted.

Step 2. A sum spreads, a transpose swaps back. If \({s = \sum_i x_i}\), every \(x_i\) moves \(s\) one for one, so each gets the same gradient: the incoming one, copied. A transpose or a reshape moves entries without changing them, so the gradient just goes back to where each entry came from.

DL01.6.2-F3

Left: summing over an axis copies the gradient back along that axis. Right: a transpose sends the gradient back transposed, so every number returns to the entry it came from.

Step 3. Broadcasting is undone by summing. Every operation of two arguments may now broadcast. Its closure computes the gradient as if both arguments already had the result’s shape, then hands each argument the sum over the axes that were stretched for it. A small helper, unbroadcast, does that last step.

The formal statement

NoteThe backward rules for shapes

With \(G\) the gradient arriving at the result:

• Matrix product \({C = AB}\): \(A\) receives \({GB\T}\), \(B\) receives \({A\T G}\).
• Sum over an axis: \(x\) receives \(G\) copied back along that axis. Mean: the same, divided by the number of entries averaged.
• Reshape and transpose: \(x\) receives \(G\) reshaped or transposed back.
• Broadcasting: an argument of shape \(s\) receives the gradient summed over every axis that was added in front of \(s\) or stretched from length 1.

Derivation of the matrix rules. Entry \({c_{ij} = \sum_l a_{il} b_{lj}}\), so \(a_{il}\) appears in every entry of row \(i\) of \(C\), with coefficient \(b_{lj}\) in entry \((i, j)\). The adjoint recurrence adds up those uses:

\[ \begin{aligned} \pd{L}{a_{il}} &= \sum_j \pd{L}{c_{ij}}\,\pd{c_{ij}}{a_{il}} \\ &= \sum_j g_{ij}\, b_{lj} = (GB\T)_{il} . \end{aligned} \]

In the same way \(b_{lj}\) appears in every entry of column \(j\), with coefficient \(a_{il}\), so \({\partial L / \partial b_{lj}} = {\sum_i a_{il}\, g_{ij}} = {(A\T G)_{lj}}\). Goodfellow et al. (section 6.5.6) state exactly these two rules for the matrix product’s backward method, and the CS231n notes read them off from the shapes.

Derivation of the broadcasting rule. Broadcasting along an axis is the same as copying the array along it and then running the elementwise operation. The copy is a map \({x \mapsto (x, x, \dots, x)}\), and every copy \(x^{(r)}\) moves with \(x\) one for one, so by the adjoint recurrence

\[ \pd{L}{x} = \sum_{r} \pd{L}{x^{(r)}} \cdot 1 , \]

the sum of the gradients of the copies: a sum over the stretched axis. Axes added in front count as stretched from length 1, which is how NumPy treats a missing axis.

The layer in this course’s convention. For a batch with rows as examples, a layer computes \({Z = XW\T + \vect{1}\vect{b}\T}\) with \(W\) of shape (out, in). The product rules applied to \(X\) and \({W\T}\), and the transpose rule applied to \(W\), give

\[ \begin{aligned} \nabla_X L &= GW, \\ \nabla_W L &= (X\T G)\T = G\T X, \end{aligned} \]

and the bias, broadcast down the batch, gets the column sums of \(G\). These are the batch gradients of the course conventions. M05.7.3 wrote the same layer as \({XW + \vect{1}\vect{b}\T}\) with \(W\) of shape (in, out), so its \({X\T G}\) is the transpose of ours.

A stable log-softmax, and a fused cross-entropy. For logits \(z\) of shape \({(n, K)}\), the log-softmax \({z_k - \log \sum_j e^{z_j}}\) is unchanged when a constant is added to a row (DL01.4.3). So subtract each row’s maximum \(m\) first, treated as a constant: \({s = z - m}\) and \({\log\operatorname{softmax}(z)} = {s} - {\log \sum_j e^{s_j}}\). Every exponent is then at most 0, so nothing overflows. Built from the engine’s operations this takes a chain of nodes; the mean cross-entropy after it takes more.

DL01.6.2-F6

The mean cross-entropy built from ten small operations, and the same function written as one operation whose backward rule is the known result \({(P - Y)/n}\).

A fused operation computes the whole loss in one forward step and has its own backward rule, \({(P - Y)/n}\), where \(P\) holds the softmax probabilities and \(Y\) the one-hot labels (M05.7.3, with the factor \({1/n}\) from the mean). It is worth doing for three reasons: one node instead of ten, so the graph keeps one array of shape \({(n, K)}\) instead of four; one short formula for the gradient instead of a chain of exp, sum and log steps; and stability built in rather than left to the order of the user’s operations.

Worked example

A layer’s backward pass by hand. Take a batch of two examples with two features and a layer with three outputs:

\[ X = \begin{pmatrix} 1 & 2 \\ 0 & -1 \end{pmatrix}, \quad W = \begin{pmatrix} 1 & 0 \\ -1 & 1 \\ 2 & 1 \end{pmatrix}, \]

and \({b = (0, 1, -1)}\). Row 1 of \({XW\T}\) is \({(1, 1, 4)}\) and row 2 is \({(0, -1, -1)}\). Adding \(b\) to each row gives \(Z\), and suppose the loss sends back \(G\):

\[ \begin{aligned} Z &= \begin{pmatrix} 1 & 2 & 3 \\ 0 & 0 & -2 \end{pmatrix}, \\ G &= \begin{pmatrix} 1 & 0 & -1 \\ 2 & 1 & 0 \end{pmatrix}. \end{aligned} \]

DL01.6.2-F4

The engine’s backward pass through \({Z = XW\T + b}\), one closure at a time. The yellow rows are the three gradients the layer needs.

• The closure of the addition passes \(G\) to \({XW\T}\) unchanged, and passes \(b\) its column sums, \({(1 + 2,\ 0 + 1,\ -1 + 0)} = {(3, 1, -1)}\), because \(b\) was copied onto both rows.
• The closure of the product sends \({G(W\T)\T = GW}\) to \(X\). Row 1 of \(G\) is \({(1, 0, -1)}\), so row 1 of \(GW\) is \({1 \cdot (1, 0)} - {1 \cdot (2, 1)} = {(-1, -1)}\); row 2 is \({2 \cdot (1, 0)} + {1 \cdot (-1, 1)} = {(1, 1)}\).
• It sends \({X\T G}\), shape \({(2, 3)}\), to the node \({W\T}\), and the transpose’s closure turns it into \({G\T X}\) for \(W\). Row \(j\) of \({G\T X}\) is \({\sum_i g_{ij}\,\vect{x}_i\T}\): row 1 is \({1 \cdot (1, 2)} + {2 \cdot (0, -1)} = {(1, 0)}\), row 2 is \({(0, -1)}\) and row 3 is \({(-1, -2)}\).

Every gradient has the shape of its array: \({(2, 2)}\) for \(X\), \({(3, 2)}\) for \(W\), \({(3,)}\) for \(b\).

Code

The class of DL01.6.1 is folded here, without its three same-shape operations, which are replaced below.

import numpy as np

class Tensor:
    def __init__(self, data):
        self.data = np.asarray(data, dtype=float)
        self.grad = np.zeros_like(self.data)
        self.parents = ()
        self.backward_fn = None

    def record(self, parents, backward_fn):
        self.parents, self.backward_fn = parents, backward_fn
        return self

    @property
    def shape(self):
        return self.data.shape

    def __neg__(self):
        out = Tensor(-self.data)
        def backward_fn():
            self.grad += -out.grad
        return out.record((self,), backward_fn)

    def __pow__(self, k):
        out = Tensor(self.data ** k)
        def backward_fn():
            self.grad += k * self.data ** (k - 1) * out.grad
        return out.record((self,), backward_fn)

    def exp(self):
        out = Tensor(np.exp(self.data))
        def backward_fn():
            self.grad += out.data * out.grad
        return out.record((self,), backward_fn)

    def log(self):
        out = Tensor(np.log(self.data))
        def backward_fn():
            self.grad += out.grad / self.data
        return out.record((self,), backward_fn)

    def sin(self):
        out = Tensor(np.sin(self.data))
        def backward_fn():
            self.grad += np.cos(self.data) * out.grad
        return out.record((self,), backward_fn)

    def tanh(self):
        out = Tensor(np.tanh(self.data))
        def backward_fn():
            self.grad += (1 - out.data ** 2) * out.grad
        return out.record((self,), backward_fn)

    def relu(self):
        out = Tensor(np.maximum(self.data, 0))
        def backward_fn():
            self.grad += (self.data > 0) * out.grad
        return out.record((self,), backward_fn)

    def topo_order(self):
        order, seen, stack = [], set(), [(self, False)]
        while stack:
            node, finished = stack.pop()
            if finished:
                order.append(node)
            elif id(node) not in seen:
                seen.add(id(node))
                stack.append((node, True))
                stack.extend((p, False) for p in node.parents if id(p) not in seen)
        return order

    def backward(self):
        order = self.topo_order()
        for node in order:
            if node.parents:
                node.grad = np.zeros_like(node.data)
        self.grad = np.ones_like(self.data)
        for node in reversed(order):
            if node.backward_fn:
                node.backward_fn()

First the helper, then addition, subtraction and multiplication that broadcast. A plain number or array is lifted to a tensor with no parents: a constant.

def unbroadcast(grad, shape):
    """Sum grad over the axes that broadcasting stretched, so that it has the given shape."""
    while grad.ndim > len(shape):              # axes that broadcasting added in front
        grad = grad.sum(axis=0)
    for axis, size in enumerate(shape):        # axes stretched from length 1
        if size == 1 and grad.shape[axis] != 1:
            grad = grad.sum(axis=axis, keepdims=True)
    return grad

def lift(x):
    return x if isinstance(x, Tensor) else Tensor(x)

def add(self, other):
    other = lift(other)
    out = Tensor(self.data + other.data)
    def backward_fn():
        self.grad += unbroadcast(out.grad, self.shape)
        other.grad += unbroadcast(out.grad, other.shape)
    return out.record((self, other), backward_fn)

def sub(self, other):
    other = lift(other)
    out = Tensor(self.data - other.data)
    def backward_fn():
        self.grad += unbroadcast(out.grad, self.shape)
        other.grad += unbroadcast(-out.grad, other.shape)
    return out.record((self, other), backward_fn)

def mul(self, other):
    other = lift(other)
    out = Tensor(self.data * other.data)
    def backward_fn():
        self.grad += unbroadcast(other.data * out.grad, self.shape)
        other.grad += unbroadcast(self.data * out.grad, other.shape)
    return out.record((self, other), backward_fn)

Tensor.__add__, Tensor.__sub__, Tensor.__mul__ = add, sub, mul
Tensor.__radd__, Tensor.__rmul__ = add, mul

Then the operations that change shapes.

def matmul(self, other):
    out = Tensor(self.data @ other.data)
    def backward_fn():
        self.grad += out.grad @ other.data.T       # G B^T, shape of A
        other.grad += self.data.T @ out.grad       # A^T G, shape of B
    return out.record((self, other), backward_fn)

def sum_(self, axis=None, keepdims=False):
    out = Tensor(self.data.sum(axis=axis, keepdims=keepdims))
    def backward_fn():
        g = out.grad if keepdims or axis is None else np.expand_dims(out.grad, axis)
        self.grad += np.broadcast_to(g, self.shape)    # copy back along the summed axis
    return out.record((self,), backward_fn)

def mean(self, axis=None, keepdims=False):
    total = self.sum(axis=axis, keepdims=keepdims)
    return total * (total.data.size / self.data.size)

def reshape(self, *shape):
    out = Tensor(self.data.reshape(*shape))
    def backward_fn():
        self.grad += out.grad.reshape(self.shape)
    return out.record((self,), backward_fn)

def transpose(self):
    out = Tensor(self.data.T)
    def backward_fn():
        self.grad += out.grad.T
    return out.record((self,), backward_fn)

Tensor.__matmul__, Tensor.sum, Tensor.mean, Tensor.reshape = matmul, sum_, mean, reshape
Tensor.T = property(transpose)

The worked example, by the engine:

X = Tensor([[1.0, 2.0], [0.0, -1.0]])
W = Tensor([[1.0, 0.0], [-1.0, 1.0], [2.0, 1.0]])
b = Tensor([0.0, 1.0, -1.0])
G = np.array([[1.0, 0.0, -1.0], [2.0, 1.0, 0.0]])

Z = X @ W.T + b
(Z * G).sum().backward()          # weighting by G makes G the gradient that arrives at Z
print("Z =", Z.data.tolist())
print("X.grad =", X.grad.tolist(), " W.grad =", W.grad.tolist(), " b.grad =", b.grad.tolist())

Z = [[1.0, 2.0, 3.0], [0.0, 0.0, -2.0]]
X.grad = [[-1.0, -1.0], [1.0, 1.0]]  W.grad = [[1.0, 0.0], [0.0, -1.0], [-1.0, -2.0]]  b.grad = [3.0, 1.0, -1.0]

Now every rule against central differences. The check weights the output with a random array, so that the gradient arriving at the operation is not all ones, and it asserts each gradient’s shape before comparing a single number.

rng = np.random.default_rng(0)

def check(f, *inputs, h=1e-6):
    """Largest gap between backward() and central differences, over every entry of every input."""
    weights = rng.normal(size=f(*inputs).shape)
    loss = lambda: (f(*inputs) * weights).sum()
    for x in inputs:
        x.grad = np.zeros_like(x.data)
    loss().backward()
    gap = 0.0
    for x in inputs:
        assert x.grad.shape == x.shape, (x.grad.shape, x.shape)
        for i in np.ndindex(x.shape):
            old = x.data[i]
            x.data[i] = old + h
            up = loss().data
            x.data[i] = old - h
            down = loss().data
            x.data[i] = old
            gap = max(gap, abs((up - down) / (2 * h) - x.grad[i]))
    return gap

T = lambda *shape: Tensor(rng.normal(size=shape))
cases = {
    "broadcast add (3,1)+(1,4)": (lambda a, b: a + b, T(3, 1), T(1, 4)),
    "broadcast mul (2,3)*(3,)":  (lambda a, b: a * b, T(2, 3), T(3)),
    "matmul (4,3)@(3,2)":        (lambda a, b: a @ b, T(4, 3), T(3, 2)),
    "sum over axis 0":           (lambda a: a.sum(axis=0), T(2, 3)),
    "mean over axis 1, kept":    (lambda a: a.mean(axis=1, keepdims=True), T(2, 3)),
    "reshape (2,3)->(3,2)":      (lambda a: a.reshape(3, 2), T(2, 3)),
    "transpose":                 (lambda a: a.T, T(2, 3)),
    "layer X @ W.T + b":         (lambda x, w, b: (x @ w.T + b).tanh(), T(5, 3), T(4, 3), T(4)),
}
gaps = {name: check(f, *xs) for name, (f, *xs) in cases.items()}
for name, gap in gaps.items():
    print(f"{name:28s} largest gap {gap:.1e}")

broadcast add (3,1)+(1,4)    largest gap 1.5e-10
broadcast mul (2,3)*(3,)     largest gap 2.9e-10
matmul (4,3)@(3,2)           largest gap 1.0e-10
sum over axis 0              largest gap 1.5e-10
mean over axis 1, kept       largest gap 5.7e-11
reshape (2,3)->(3,2)         largest gap 2.9e-10
transpose                    largest gap 1.2e-10
layer X @ W.T + b            largest gap 4.9e-10

Last, the log-softmax built from these operations, the mean cross-entropy built on top of it, and the fused version with its own backward rule.

def log_softmax(z):
    s = z - z.data.max(axis=1, keepdims=True)    # the row maxima enter as a constant
    return s - s.exp().sum(axis=1, keepdims=True).log()

def cross_entropy(z, labels):
    """Mean cross-entropy of logits z (n, K) against integer labels, as one operation."""
    n = z.shape[0]
    shifted = z.data - z.data.max(axis=1, keepdims=True)
    log_p = shifted - np.log(np.exp(shifted).sum(axis=1, keepdims=True))
    out = Tensor(-log_p[np.arange(n), labels].mean())
    def backward_fn():
        g = np.exp(log_p)                       # P
        g[np.arange(n), labels] -= 1            # P - Y
        z.grad += g / n * out.grad
    return out.record((z,), backward_fn)

labels = np.array([0, 2, 1, 1, 0])
Y = np.eye(3)[labels]
z = T(5, 3)
composed = -(log_softmax(z) * Y).sum(axis=1).mean()
composed.backward()
g_composed = z.grad.copy()
z.grad = np.zeros_like(z.data)
fused = cross_entropy(z, labels)
fused.backward()
ops = sum(1 for node in composed.topo_order() if node.parents)
print(f"losses {composed.data:.6f} and {fused.data:.6f}; gradient gap {np.abs(z.grad - g_composed).max():.0e}")
print("operations in the composed graph:", ops, "  in the fused one:", len(fused.topo_order()) - 1)
print("logits (1000, 0, -1000), label 1:", cross_entropy(Tensor([[1000.0, 0.0, -1000.0]]), np.array([1])).data)

losses 1.451307 and 1.451307; gradient gap 3e-17
operations in the composed graph: 10   in the fused one: 1
logits (1000, 0, -1000), label 1: 1000.0

The two losses and gradients agree to rounding, and the logits \({(1000, 0, -1000)}\) give the loss 1000 without overflow: the true log-probability of class 1 is \({-1000} - {\ln(1 + e^{-1000} + e^{-2000})}\), which rounds to \(-1000\).

Explore

DL01.6-X1

Broadcast and backward

Choose the shapes of \(A\) and \(B\) and an operation: add, multiply, matrix product, or a sum of \(A\) over an axis. Step forward to see the result as a grid, then send back a gradient of ones and step again to watch it summed back to each array’s shape.

Three things to try. With \(A\) of shape (3, 1) and \(B\) of shape (1, 4), add, and watch each gradient collect 4 and 3: the number of copies. Switch to multiply: \(A\)’s gradient becomes the sum of \(B\)’s row, and \(B\)’s the sum of \(A\)’s column. Then choose shapes that clash, such as (3, 2) and (1, 4), and read why NumPy refuses.

Pitfalls

DL01.6.2-F5

Wrong beside right. A bias broadcast down a batch must have its gradient summed down the batch. Passed on as it is, the gradient has the batch’s shape, not the bias’s.

“Broadcasting is free, so its backward pass is free too.”

Forward, broadcasting copies nothing. Backward, the copies have to be paid for: every stretched axis is summed. Forgetting it gives a bias gradient of shape (n, out). With += NumPy then refuses, which is lucky; written as b.grad = b.grad + G it broadcasts silently, and the next update turns the single bias into one bias per example.

“The weight gradient is \({X\T G}\).”

Only when the layer is written \({XW}\) with \(W\) of shape (in, out), as in M05.7.3. With the weight stored as (out, in), as here and in PyTorch, it is \({G\T X}\). Check the shape: \({X\T G}\) has shape (in, out), which is not the shape of \(W\). When the shapes are square, the wrong formula fits the array and is silently wrong, which is why the check above uses a rectangular weight.

“Log-softmax is log(softmax(z)), so build it that way.”

With logits \({(1000, 0, -1000)}\), \({e^{1000}}\) overflows to infinity and the ratio becomes nan. Subtracting the row maximum first keeps every exponent at or below 0. Because the shift does not change the result, it can be treated as a constant, and its gradient is never needed.

Check yourself

Recall

1. What does \(A\) receive from the backward step of \({C = AB}\), and what does \(B\) receive?

Show answer

\(A\) receives \({GB\T}\) and \(B\) receives \({A\T G}\), where \(G\) is the gradient arriving at \(C\).

2. What is the backward rule of broadcasting?

Show answer

Sum the gradient over every axis that was stretched from length 1 or added in front, so that it ends with the shape of the array.

3. What does a sum over axis 0 send back to its input?

Show answer

The incoming gradient copied along axis 0: every term of the sum gets the same gradient.

4. In the course convention \({Z = XW\T + \vect{1}\vect{b}\T}\), what are the gradients of \(W\) and \(\vect{b}\)?

Show answer

\({G\T X}\), of shape (out, in), and the column sums of \(G\).

5. Give two reasons to fuse softmax and cross-entropy into one operation.

Show answer

One node instead of about ten, so the graph keeps far fewer arrays; and a short, exact gradient \({(P - Y)/n}\) with the stable form built in. (Speed is a third.)

Exercises

1. ★ Arrays of shapes (5, 1, 3) and (4, 1) are added. What is the shape of the result, and what does unbroadcast do to the gradient for each array?

Show solution

Aligning from the right, (4, 1) is treated as (1, 4, 1). Axis by axis: 5 and 1 give 5, 1 and 4 give 4, 3 and 1 give 3, so the result has shape (5, 4, 3). For the first array, axis 1 was stretched from 1 to 4, so the gradient is summed over axis 1 with the axis kept: (5, 1, 3). For the second, the added front axis is summed away first, giving (4, 3), and then axis 1 (length 1 in the array, 3 in the gradient) is summed with the axis kept: (4, 1).

2. ★★ Derive the backward rule of \({v = \vect{x}\T\vect{w}}\) for a single example from the matrix rules, and check it with \({\vect{x} = (1, 2)}\), \({\vect{w} = (3, -1)}\) and an incoming gradient \({g = 2}\).

Show solution

Write \(\vect{x}\T\) as \(A\) of shape (1, 2) and \(\vect{w}\) as \(B\) of shape (2, 1), so \(C\) and \(G\) are (1, 1) and \({G = (g)}\). Then \(A\) receives \({GB\T = g\,\vect{w}\T}\), so \(\vect{x}\) receives \({g\,\vect{w} = (6, -2)}\), and \(B\) receives \({A\T G = g\,\vect{x} = (2, 4)}\). Directly: \({v = 3 \cdot 1 + (-1) \cdot 2 = 1}\), \({\partial v / \partial x_1} = {w_1} = {3}\) and so on, times \({g = 2}\): the same.

3. ★★★ The mean over axis 1 of an array of shape \({(n, K)}\) is computed as the sum over that axis times \({1/K}\). Show that its backward pass sends \({G/K}\) copied along the axis, and that the fused cross-entropy’s gradient \({(P - Y)/n}\) carries the \({1/n}\) of the mean over the batch for the same reason.

Show solution

Let \({m_i = \frac{1}{K}\sum_k x_{ik}}\). Then \({\partial m_i / \partial x_{ik}} = {1/K}\) and \(m_i\) does not depend on other rows, so \({\partial L / \partial x_{ik}} = {g_i / K}\): the incoming \(g_i\) copied along the row and divided by \(K\). In the engine, the multiplication by the constant \({1/K}\) sends \({g/K}\) to the sum, and the sum copies it along the axis. For the loss \({L = \frac1n\sum_i \ell_i}\), the same argument over the batch gives \({\partial L / \partial \ell_i} = {1/n}\), and each example’s \({\partial \ell_i / \partial z_i = p_i - y_i}\) (M05.7.3). Multiplying, row \(i\) of the gradient is \({(p_i - y_i)/n}\).

Summary card

Every backward rule sends each entry the sum of what its uses received.

DL01.6.2-F1

Copy forward, sum backward.

• \({C = AB}\): \(A\) gets \({GB\T}\), \(B\) gets \({A\T G}\).
• Sum over an axis: copy \(G\) back along it; mean: also divide.
• Reshape, transpose: undo them on \(G\).
• Broadcasting: sum over the stretched axes (unbroadcast).
• Fused cross-entropy: \({(P - Y)/n}\).

Sources and next steps

Sources. Each was opened and checked against this lesson on 15 November 2026.

• I. Goodfellow, Y. Bengio and A. Courville, Deep Learning, section 6.5.6: the matrix product \({C = AB}\) defines two backward rules, \({GB\T}\) for \(A\) and \({A\T G}\) for \(B\), and each operation’s bprop returns a sum of Jacobian-gradient products (equation 6.54).
• Stanford CS231n, Backpropagation, Intuitions, section “Gradients for vectorized operations”: the matrix-matrix multiply gradient written in code, and the advice to re-derive it by dimension analysis; also “gradients add up at forks”.
• NumPy, Broadcasting: shapes are compared from the trailing dimension, two dimensions are compatible when equal or when one is 1, missing dimensions count as 1, and the stretching makes no copies.
• A. G. Baydin and others, Automatic Differentiation in Machine Learning: a Survey, JMLR 18(153), 2018, section 3.2: reverse mode computes the transposed Jacobian-vector product without forming the Jacobian.

Next. DL01.6.3 Parameters, modules and an optimizer builds a small PyTorch-shaped library on top of the engine.

This lesson unlocks

DL01.6.3 Parameters, modules and an optimizer

Your progress is saved in this browser only. No account, no tracking.