Jacobians and the chain rule

Worked solutions

Problem 1

1. $h_i = \sigma(z_i)$ for $i = 1, \dots, n$.Elementwise means output $i$ is computed from input $i$ alone.
2. $\partial h_i/\partial z_j = \sigma'(z_i)$ if $i = j$, and $0$ if $i \neq j$.$h_i$ contains no $z_j$ with $j \neq i$, so those partials vanish; the remaining one is the ordinary derivative of $\sigma$ at $z_i$.
3. The Jacobian has $\sigma'(z_1), \dots, \sigma'(z_n)$ on the diagonal and zeros elsewhere.Entry $(i,j)$ of $\partial h/\partial z$ is $\partial h_i/\partial z_j$, so step 2 fills it entry by entry.
4. $\partial h/\partial z = \operatorname{diag}(\sigma'(z))$, $n \times n$One row per $h_i$ and one column per $z_j$, and both have $n$ entries.

Problem 2

1. $z_i = \sum_j W_{ij} x_j + b_i$.Writing entry $i$ out shows which terms contain each $x_j$ and each $b_k$.
2. $\partial z_i/\partial x_j = W_{ij}$.Only term $j$ of the sum contains $x_j$, with coefficient $W_{ij}$; $b_i$ does not depend on $x$.
3. $\partial z_i/\partial b_k = 1$ if $i = k$, and $0$ otherwise.$b_k$ appears only in $z_k$, with coefficient 1, and the sum does not depend on $b$.
4. $\partial z/\partial x = W$ ($m \times n$); $\partial z/\partial b = I_m$Step 2 fills the $m \times n$ Jacobian with the entries of $W$; step 3 fills the $m \times m$ Jacobian with ones on the diagonal. The previous page's $\partial(Ax)/\partial x = A$ is the case $b = 0$.

Problem 3

1. Let $z = Wx + b$, so $h = \sigma(z)$.Naming the pre-activation splits $h$ into two maps whose Jacobians are Problems 1 and 2.
2. $\partial h/\partial x = (\partial h/\partial z)(\partial z/\partial x)$.Chain rule in numerator layout: the Jacobian of the outer map, $\sigma$, goes on the left.
3. $\partial h/\partial z = \operatorname{diag}(\sigma'(z))$, $m \times m$; $\partial z/\partial x = W$, $m \times n$.Problem 1 with $m$ in place of $n$, since $z$ has $m$ entries, and Problem 2.
4. $(m \times m)(m \times n) = m \times n$.The inner dimensions agree, and the result has one row per $h_i$ and one column per $x_j$, the shape $\partial h/\partial x$ must have.
5. $\partial h/\partial x = \operatorname{diag}(\sigma'(z))\,W$, $m \times n$Row $i$ of this product is row $i$ of $W$ scaled by $\sigma'(z_i)$: $h_i$ depends on $x$ only through $z_i$, and $z_i$ depends on $x$ through row $i$ of $W$.

Problem 4

1. Let $u = Ax$, so $y = Bu$.Naming the intermediate vector makes $y$ a composition of two linear maps, which is the form the chain rule takes.
2. $\partial u/\partial x = A$, $m \times n$; $\partial y/\partial u = B$, $p \times m$.The Jacobian of a linear map is its matrix (Problem 2 with $b = 0$).
3. $\partial y/\partial x = (\partial y/\partial u)(\partial u/\partial x) = BA$.The outer map is $B$, so its Jacobian goes on the left.
4. $(p \times m)(m \times n) = p \times n$. The other order, $(m \times n)(p \times m)$, needs $n = p$.Only $BA$ has inner dimensions that agree for every $m, n, p$, and its result is $p \times n$, one row per $y_i$ and one column per $x_j$. The shapes force this order.
5. $\partial y/\partial x = BA$, $(p \times m)(m \times n) = p \times n$; $AB$ is not defined unless $p = n$, and is wrong even thenDirectly, $y = (BA)x$, whose Jacobian is $BA$. When $p = n$, $AB$ is $m \times m$, not $p \times n$; when all three sizes are equal it has the right shape, but $AB \neq BA$ in general.

Problem 5

1. Let $z = Wx + b$, so $h = \sigma(z)$ and $L = \tfrac12 (h - y)^\top (h - y)$.The squared norm is the inner product of $h - y$ with itself, which the product rule can differentiate.
2. $dL = (h - y)^\top dh$, so $\partial L/\partial h = (h - y)^\top$, $1 \times m$.Product rule gives two equal scalar terms, and the $\tfrac12$ cancels the 2; $y$ is constant.
3. $\partial h/\partial x = \operatorname{diag}(\sigma'(z))\,W$, $m \times n$.Problem 3.
4. $\partial L/\partial x = (h - y)^\top \operatorname{diag}(\sigma'(z))\,W$, with shapes $(1 \times m)(m \times m)(m \times n) = 1 \times n$.Chain rule with the outer Jacobian, that of $L$ with respect to $h$, on the left.
5. $\nabla_x L = W^\top \operatorname{diag}(\sigma'(z))\,(h - y)$, with shapes $(n \times m)(m \times m)(m \times 1) = n \times 1$.The gradient is the transpose of the row; transposing a product reverses it, and a diagonal matrix is its own transpose.
6. $\operatorname{diag}(\sigma'(z))\,(h - y) = \sigma'(z) \odot (h - y)$.A diagonal matrix times a vector scales entry $i$ of the vector by diagonal entry $i$, which is exactly the elementwise product.
7. $\nabla_x L = W^\top\big(\sigma'(z) \odot (h - y)\big)$Substitute step 6 into step 5. Shape check: $(n \times m)(m \times 1) = n \times 1$, the shape of $x$.

Problem 6

1. $\partial L/\partial x = (\partial L/\partial h)(\partial h/\partial x)$, with shapes $(1 \times m)(m \times n) = 1 \times n$.$L$ depends on $x$ only through $h$, so the chain rule applies with $L$ as the outer map.
2. $\nabla_x L = (\partial h/\partial x)^\top (\partial L/\partial h)^\top = (\partial h/\partial x)^\top \nabla_h L$.Transpose both sides: the gradient is the transposed row, a transposed product reverses its order, and $(\partial L/\partial h)^\top$ is by definition $\nabla_h L$.
3. Let $z = Wx + b$. Then $(\partial h/\partial x)^\top = \big(\operatorname{diag}(\sigma'(z))\,W\big)^\top = W^\top \operatorname{diag}(\sigma'(z))$.Problem 3, transposed; a diagonal matrix is its own transpose.
4. $\nabla_x L = W^\top\big(\operatorname{diag}(\sigma'(z))\,\nabla_h L\big)$.Multiplying right to left means the $m \times m$ matrix only ever acts on the vector $\nabla_h L$, so the matrix $\operatorname{diag}(\sigma'(z))\,W$ is never needed.
5. $\operatorname{diag}(\sigma'(z))\,\nabla_h L = \sigma'(z) \odot \nabla_h L$.A diagonal matrix acting on a vector is an elementwise product: $m$ multiplications, and only the $m$ numbers $\sigma'(z)$ are stored.
6. $\nabla_x L = W^\top(\sigma'(z) \odot \nabla_h L)$; the $\operatorname{diag}$ never appears, so the elementwise layer costs $O(m)$ storage instead of the $O(m^2)$ of its JacobianAt $m = 4096$ the Jacobian of $\sigma$ has $4096^2 \approx 16.8$ million entries, about 67 MB in 32-bit floats, per example and per layer, and all but 4096 of them are zero. The elementwise form stores 4096 numbers, and the remaining work is the product with $W^\top$, which the backward pass pays anyway.

Problem 7

1. Let $z_1 = W_1 x$, so $h_1 = \sigma(z_1)$. Take $x \in \mathbb{R}^n$, $W_1 \in \mathbb{R}^{k \times n}$, $W_2 \in \mathbb{R}^{p \times k}$, so $z_1, h_1 \in \mathbb{R}^k$ and $h_2, y \in \mathbb{R}^p$.$\sigma'$ has to be evaluated at the pre-activation, so it needs a name; the sizes let every product below be checked.
2. $\nabla_{h_2} L = h_2 - y$, $p \times 1$.The previous page's $\nabla_x\|Ax - b\|^2 = 2A^\top(Ax - b)$ with $A = I$, halved.
3. $\nabla_{h_1} L = W_2^\top (h_2 - y)$, with shapes $(k \times p)(p \times 1) = k \times 1$.Problem 6's rule: the gradient with respect to a map's input is its transposed Jacobian times the gradient with respect to its output, and $\partial h_2/\partial h_1 = W_2$ (Problem 2).
4. $\nabla_{z_1} L = \sigma'(z_1) \odot \nabla_{h_1} L$, $k \times 1$.The same rule through $\sigma$: the transposed Jacobian is $\operatorname{diag}(\sigma'(z_1))$, and it acts on a vector as an elementwise product.
5. $\nabla_x L = W_1^\top \nabla_{z_1} L$, with shapes $(n \times k)(k \times 1) = n \times 1$.The same rule once more, with $\partial z_1/\partial x = W_1$.
6. $\nabla_x L = W_1^\top\big(\sigma'(z_1) \odot W_2^\top (h_2 - y)\big)$Substitute steps 3 and 4 into step 5. Evaluated from the inside out, this is backpropagation: the last layer's transpose acts first and the first layer's acts last.

Problem 8

1. Let $r = Wx - y$. Then $L = \tfrac12 \sum_i r_i^2$ with $r_i = \sum_j W_{ij} x_j - y_i$, so $L = \tfrac12 \sum_i \big(\sum_j W_{ij}x_j - y_i\big)^2$.A matrix gradient is defined entry by entry, so write $L$ in terms of the entries $W_{ij}$.
2. $\partial L/\partial W_{ij} = r_i\,\partial r_i/\partial W_{ij}$.$W_{ij}$ sits in row $i$ of $W$, which only builds $r_i$, so only the $i$-th square depends on it; its derivative is $r_i$ times the derivative of $r_i$.
3. $\partial r_i/\partial W_{ij} = x_j$.In $\sum_{j'} W_{ij'} x_{j'}$, only the $j' = j$ term contains $W_{ij}$, with coefficient $x_j$.
4. $(\nabla_W L)_{ij} = (Wx - y)_i\,x_j$.Entry $(i,j)$ of $\nabla_W L$ is $\partial L/\partial W_{ij}$ (the previous page's convention).
5. A matrix with entries $a_i b_j$ is the outer product $a b^\top$.$(a b^\top)_{ij} = a_i b_j$: $a$ supplies the row index and $b$ the column index.
6. $\nabla_W L = (Wx - y)\,x^\top$, $m \times n$Step 5 with $a = Wx - y$ and $b = x$. Shape check: $(m \times 1)(1 \times n) = m \times n$, the shape of $W$, so $W - \eta\,\nabla_W L$ makes sense.

Problem 9

1. $u = x\,\|x\|^{-1}$.A vector times a scalar function of $x$, which the product rule handles.
2. $du = \|x\|^{-1}\,dx + x\,d\big(\|x\|^{-1}\big)$.Product rule; both factors depend on $x$.
3. $d\big(\|x\|^{-1}\big) = -\|x\|^{-2}\,d\|x\| = -\|x\|^{-3}\,x^\top dx$.Scalar chain rule for $t \mapsto 1/t$, then $d\|x\| = x^\top dx/\|x\|$, which is the previous page's $\nabla_x\|x\| = x/\|x\|$ written as a differential.
4. $du = \|x\|^{-1}\,dx - \|x\|^{-3}\,x\,x^\top dx$.Substitute step 3; $x\,(x^\top dx) = (x x^\top)\,dx$ by associativity.
5. $\partial u/\partial x = \dfrac{1}{\|x\|}\Big(I - \dfrac{x x^\top}{\|x\|^2}\Big)$, $n \times n$.The Jacobian is the matrix that multiplies $dx$; factor out $1/\|x\|$.
6. Let $\hat{x} = x/\|x\|$, the unit vector along $x$ (it equals $u$). Then $x x^\top/\|x\|^2 = \hat{x}\hat{x}^\top$.Split $\|x\|^2$ into one $\|x\|$ for each factor of $x$.
7. $\partial u/\partial x = \dfrac{1}{\|x\|}\big(I - \hat{x}\hat{x}^\top\big)$, $\hat{x} = x/\|x\|$Sanity check: scaling $x$ does not change $u$, so the Jacobian must send the direction of $x$ to zero, and $(I - \hat{x}\hat{x}^\top)\hat{x} = \hat{x} - \hat{x} = 0$.

Problem 10

1. Let $z = Wx$ and $h = \sigma(z)$, so the function is $f = \|h\|^2 = h^\top h$.Naming the pre-activation and the output splits $f$ into a layer (Problem 3) and a squared norm (previous page).
2. $\nabla_h f = 2h = 2\sigma(z)$.The previous page's $\nabla_x(x^\top A x) = (A + A^\top)x$ with $A = I$.
3. $\nabla_x f = (\partial h/\partial x)^\top \nabla_h f$.Problem 6: the gradient with respect to $x$ is the transposed Jacobian of $h$ times the gradient with respect to $h$.
4. $(\partial h/\partial x)^\top = W^\top \operatorname{diag}(\sigma'(z))$.Problem 3 with $b = 0$, transposed.
5. $\nabla_x f = W^\top\big(\sigma'(z) \odot 2\sigma(z)\big)$, with shapes $(n \times m)(m \times 1) = n \times 1$.The diagonal matrix acts on $2\sigma(z)$ as an elementwise product (Problem 6).
6. $\nabla_x \|\sigma(Wx)\|^2 = 2W^\top\big(\sigma'(z) \odot \sigma(z)\big)$, $z = Wx$The scalar 2 moves out of the elementwise product and the matrix product. For the sigmoid, $\sigma'(z) \odot \sigma(z) = h \odot h \odot (1 - h)$, which is how it is computed from the saved forward output.