Batch norm backward

Worked solutions

Problem 1

1. $\mu = \tfrac1N X^\top\mathbf{1} \in \mathbb{R}^d$, and likewise $\sigma^2 \in \mathbb{R}^d$.Each is a sum over the $N$ examples, which leaves one number per feature; $(d\times N)(N\times 1) = d\times 1$.
2. $X - \mathbf{1}\mu^\top$ is $N\times d$ and subtracts $\mu_j$ from every entry of column $j$; multiplying on the right by the diagonal $D$ scales column $j$ by $(\sigma_j^2+\epsilon)^{-1/2}$, so $\hat X$ is $N\times d$.$(\mathbf{1}\mu^\top)_{nj} = \mu_j$, and $(AD)_{nj} = A_{nj}D_{jj}$ for a diagonal $D$.
3. $Y = \hat X\operatorname{diag}(\gamma) + \mathbf{1}\beta^\top$ is $N\times d$, with $\gamma, \beta \in \mathbb{R}^d$.One scale and one shift per feature, shared by every example, as the bias of the previous page's dense layer is.
4. $\hat X_{nj}$ uses $X_{nj}$, $\mu_j$ and $\sigma_j^2$, and $\mu_j$ and $\sigma_j^2$ are built from all $N$ entries of column $j$ and nothing else.The definitions of $\mu_j$ and $\sigma_j^2$ sum over $n$ with $j$ fixed.
5. $\mu, \sigma^2, \gamma, \beta \in \mathbb{R}^d$; $\hat X, Y \in \mathbb{R}^{N\times d}$; $Y_{nj}$ depends on every entry of column $j$ of $X$ and on no other columnFeatures never mix and examples always do, the reverse of the dense layer on the previous page, where row $n$ of the output used only row $n$ of the input. So everything below can be done one column at a time, and this is the layer the previous page's Problem 9 excluded from gradient accumulation.

Problem 2

1. $Y_{nj} = \gamma_j\hat X_{nj} + \beta_j$.Index form of $Y = \hat X\operatorname{diag}(\gamma) + \mathbf{1}\beta^\top$; $\hat X$ does not depend on $\gamma$ or $\beta$.
2. $\partial L/\partial\gamma_j = \sum_n dY_{nj}\hat X_{nj}$ and $\partial L/\partial\beta_j = \sum_n dY_{nj}$.$\gamma_j$ and $\beta_j$ appear in $Y_{nj}$ for every example $n$, with coefficients $\hat X_{nj}$ and $1$, and the chain rule sums over every entry that contains them.
3. $\partial L/\partial\hat X_{nj} = \gamma_j\,dY_{nj}$.$\hat X_{nj}$ appears only in $Y_{nj}$, with coefficient $\gamma_j$.
4. For a matrix $M$ with $N$ rows, $M^\top\mathbf{1}$ adds up its rows, and $M\operatorname{diag}(\gamma)$ scales column $j$ by $\gamma_j$.$(M^\top\mathbf{1})_j = \sum_n M_{nj}$ and $(M\operatorname{diag}(\gamma))_{nj} = M_{nj}\gamma_j$.
5. $d\gamma = (dY\odot\hat X)^\top\mathbf{1}$; $d\beta = dY^\top\mathbf{1}$; $d\hat X = dY\operatorname{diag}(\gamma)$Shapes $d$, $d$ and $N\times d$. The sums over the batch appear because $\gamma$ and $\beta$ are shared by every example. From here on the input gradient needs only $d\hat X$, one column at a time: for feature $j$, $d\hat x = \gamma_j$ times column $j$ of $dY$.

Problem 3

1. $\partial\mu/\partial x_m = \tfrac1N$ for every $m$.$\mu = \tfrac1N\sum_n x_n$ contains $x_m$ once, with coefficient $\tfrac1N$.
2. $\partial\sigma^2/\partial x_m = \tfrac2N\sum_n (x_n-\mu)\big(\partial x_n/\partial x_m - \partial\mu/\partial x_m\big)$.Chain rule on each square $(x_n-\mu)^2$, with $\mu$ kept as the function of $x$ that it is.
3. $\partial\sigma^2/\partial x_m = \tfrac2N(x_m-\mu) - \tfrac2{N^2}\sum_n (x_n-\mu)$.$\partial x_n/\partial x_m$ is $1$ for $n = m$ and $0$ otherwise, which picks out one term; step 1 gives $\partial\mu/\partial x_m = \tfrac1N$, the same for every $n$.
4. $\sum_n (x_n - \mu) = \sum_n x_n - N\mu = 0$.$N\mu = \sum_n x_n$ by the definition of $\mu$.
5. $\nabla_x\mu = \tfrac1N\mathbf{1}$; $\nabla_x\sigma^2 = \tfrac2N(x - \mu\mathbf{1})$Both $N\times 1$, the shape of $x$. The $\tfrac1N$ in $\nabla_x\sigma^2$ is there because the forward pass uses the biased variance; it is the layer-norm page's Problem 2 with the batch in place of the features.

Problem 4

1. $\sigma^2$ feeds only the $\hat x_n$, and $\partial\hat x_n/\partial\sigma^2 = -\tfrac12(x_n-\mu)(\sigma^2+\epsilon)^{-3/2}$.In the graph each node's inputs are held fixed when it is varied; power rule on $(\sigma^2+\epsilon)^{-1/2}$.
2. $d\sigma^2 = \sum_n d\hat x_n\cdot\big(-\tfrac12\big)(x_n-\mu)(\sigma^2+\epsilon)^{-3/2} = -\dfrac{1}{2(\sigma^2+\epsilon)}\sum_n d\hat x_n\hat x_n$.The chain rule sums over the $N$ children of $\sigma^2$, and $(x_n-\mu)(\sigma^2+\epsilon)^{-1/2} = \hat x_n$ leaves one factor $(\sigma^2+\epsilon)^{-1}$.
3. $\mu$ feeds every $\hat x_n$, with $\partial\hat x_n/\partial\mu = -1/\sqrt{\sigma^2+\epsilon}$, and feeds $\sigma^2$, with $\partial\sigma^2/\partial\mu = -\tfrac2N\sum_n (x_n-\mu)$.Differentiate each child of $\mu$ with its other inputs held fixed.
4. $\partial\sigma^2/\partial\mu = 0$.Deviations from the mean sum to zero (Problem 3, step 4).
5. $d\mu = -\dfrac{1}{\sqrt{\sigma^2+\epsilon}}\sum_n d\hat x_n + d\sigma^2\cdot 0$.The chain rule sums over both kinds of child, the $\hat x_n$ and $\sigma^2$.
6. $d\sigma^2 = -\dfrac{1}{2(\sigma^2+\epsilon)}\textstyle\sum_n d\hat x_n\hat x_n$; $d\mu = -\dfrac{1}{\sqrt{\sigma^2+\epsilon}}\textstyle\sum_n d\hat x_n$Two scalars per feature. The path from $\mu$ through $\sigma^2$ is real but carries nothing, for the same reason that $\nabla_x\sigma^2$ in Problem 3 has no $\mu$ term: $\mu$ minimises the mean squared deviation, so a small change in it moves $\sigma^2$ only to second order.

Problem 5

1. $x_n$ feeds $\hat x_n$ with $\partial\hat x_n/\partial x_n = 1/\sqrt{\sigma^2+\epsilon}$, feeds $\sigma^2$ with $\partial\sigma^2/\partial x_n = \tfrac2N(x_n-\mu)$, and feeds $\mu$ with $\partial\mu/\partial x_n = \tfrac1N$.These are the children of $x_n$ in the graph of Problem 4, each differentiated with its other inputs held fixed; $x_n$ reaches $\hat x_m$ for $m \ne n$ only through $\mu$ and $\sigma^2$.
2. $dx_n = \dfrac{d\hat x_n}{\sqrt{\sigma^2+\epsilon}} + d\sigma^2\cdot\dfrac2N(x_n-\mu) + \dfrac{d\mu}{N}$.The chain rule sums over the three children.
3. $d\sigma^2\cdot\dfrac2N(x_n-\mu) = -\dfrac{1}{N(\sigma^2+\epsilon)}(x_n-\mu)\sum_m d\hat x_m\hat x_m = -\dfrac{1}{N\sqrt{\sigma^2+\epsilon}}\,\hat x_n\sum_m d\hat x_m\hat x_m$.Problem 4 for $d\sigma^2$, with the summation index renamed $m$; then $(x_n-\mu)/\sqrt{\sigma^2+\epsilon} = \hat x_n$.
4. $\dfrac{d\mu}{N} = -\dfrac{1}{N\sqrt{\sigma^2+\epsilon}}\sum_m d\hat x_m$.Problem 4 for $d\mu$.
5. $dx_n = \dfrac{1}{N\sqrt{\sigma^2+\epsilon}}\Big(N\,d\hat x_n - \sum_m d\hat x_m - \hat x_n\sum_m d\hat x_m\hat x_m\Big)$.Put the three terms over the common factor $1/(N\sqrt{\sigma^2+\epsilon})$; the first term becomes $N\,d\hat x_n$.
6. $dx = \dfrac{1}{N\sqrt{\sigma^2+\epsilon}}\Big(N\,d\hat x - \big(\textstyle\sum_n d\hat x_n\big)\mathbf{1} - \hat x\textstyle\sum_n d\hat x_n\hat x_n\Big)$ for each feature$N\times 1$. For all features at once, $dX = \tfrac1N\big(N\,d\hat X - \mathbf{1}\mathbf{1}^\top d\hat X - \hat X\odot\mathbf{1}\mathbf{1}^\top(d\hat X\odot\hat X)\big)D$, where $\mathbf{1}\mathbf{1}^\top M$ puts the column sums of $M$ in every row. Two sums per feature and elementwise work: $O(Nd)$, with $\hat X$ and $\sigma^2$ saved from the forward pass.

Problem 6

1. $\mathbf{1}^\top\hat x = \sum_n (x_n-\mu)/\sqrt{\sigma^2+\epsilon} = 0$.Deviations from the mean sum to zero, whatever $\epsilon$ is.
2. $\hat x^\top\hat x = \dfrac{\sum_n (x_n-\mu)^2}{\sigma^2+\epsilon} = \dfrac{N\sigma^2}{\sigma^2+\epsilon}$.$\sum_n (x_n-\mu)^2 = N\sigma^2$ by the definition of the biased variance.
3. $\mathbf{1}^\top dx = \dfrac{1}{N\sqrt{\sigma^2+\epsilon}}\Big(N\sum_n d\hat x_n - N\sum_n d\hat x_n - (\mathbf{1}^\top\hat x)\sum_n d\hat x_n\hat x_n\Big) = 0$.Problem 5 with $\mathbf{1}^\top\mathbf{1} = N$, and step 1 for the last term.
4. $\hat x^\top dx = \dfrac{1}{N\sqrt{\sigma^2+\epsilon}}\Big(N - \dfrac{N\sigma^2}{\sigma^2+\epsilon}\Big)\sum_n d\hat x_n\hat x_n$.Problem 5 again: $\hat x^\top d\hat x = \sum_n d\hat x_n\hat x_n$, the middle term vanishes by step 1, and step 2 gives the last.
5. $\mathbf{1}^\top dx = 0$; $\hat x^\top dx = \dfrac{1}{\sqrt{\sigma^2+\epsilon}}\,\dfrac{\epsilon}{\sigma^2+\epsilon}\textstyle\sum_n d\hat x_n\hat x_n$, which is $0$ when $\epsilon = 0$$N - N\sigma^2/(\sigma^2+\epsilon) = N\epsilon/(\sigma^2+\epsilon)$. Adding a constant to one feature across the batch leaves $\hat X$ unchanged, and so (up to $\epsilon$) does scaling it, so the loss cannot push either way. It is the layer-norm page's Problem 8 with columns in place of rows. One consequence: a bias added to a feature just before batch norm gets gradient $\mathbf{1}^\top dx = 0$ and never moves, which is why layers followed by batch norm usually drop their bias.

Problem 7

1. Row $j$ of $X^\top$ is column $j$ of $X$: the $N$ values of feature $j$.Transposing swaps the roles of rows and columns.
2. Row $j$ of $\operatorname{LN}(X^\top)$ is $\big((x - \mu_j\mathbf{1})/\sqrt{\sigma_j^2+\epsilon}\big)^\top$, which is column $j$ of $\hat X$ as a row.Layer norm's row mean and biased variance, taken over the $N$ entries of that row, are $\mu_j$ and $\sigma_j^2$.
3. So $\hat X = \operatorname{LN}(X^\top)^\top$, and $L$ depends on $X$ through $X^\top$, then $\operatorname{LN}$, then a transpose.Step 2 for every $j$.
4. A transpose only relabels entries, so its backward pass is a transpose: $dX$ is the transpose of the gradient at $X^\top$, and the upstream gradient at $\operatorname{LN}(X^\top)$ is $d\hat X^\top$.Each entry of $X$ is one entry of $X^\top$ with derivative $1$, and likewise for $\hat X$.
5. The layer-norm page's Problem 7 on row $j$ of $X^\top$, with upstream $d\hat x^\top$ and $\gamma = \mathbf{1}$, gives $\tfrac{1}{\sqrt{\sigma_j^2+\epsilon}}\big(d\hat x - \tfrac1N(\sum_n d\hat x_n)\mathbf{1} - \hat x\,\tfrac1N\sum_n d\hat x_n\hat x_n\big)$, as a row.Its $\operatorname{mean}$ is $\tfrac1N\sum_n$ here; factoring out $\tfrac1N$ gives Problem 5 exactly.
6. $\hat X = \operatorname{LN}(X^\top)^\top$, and $dX$ is the transpose of layer norm's input gradient at $X^\top$ with upstream $d\hat X^\top$: batch norm takes its means down each column, over the $N$ examples (axis 0), and layer norm along each row, over the $d$ features (axis 1)The algebra is the same; what changes is which entries share statistics. In layer norm row $n$ of $dX$ depends only on row $n$ of $d\hat X$, so examples stay independent. In batch norm column $j$ of $dX$ depends on all of column $j$ of $d\hat X$, so each example's gradient depends on the rest of the batch.

Problem 8

1. $\hat X_{nj} = (X_{nj} - \bar\mu_j)/\sqrt{\bar\sigma_j^2+\epsilon}$, so $\hat X_{nj}$ depends only on $X_{nj}$, with slope $(\bar\sigma_j^2+\epsilon)^{-1/2}$.$\bar\mu$ and $\bar\sigma^2$ were accumulated from earlier batches and do not depend on the current $X$.
2. $Y_{nj} = \gamma_j\hat X_{nj} + \beta_j$ depends only on $X_{nj}$, with slope $\gamma_j(\bar\sigma_j^2+\epsilon)^{-1/2}$.Step 1 and the chain rule through one scalar.
3. $\partial L/\partial X_{nj} = \gamma_j\,dY_{nj}/\sqrt{\bar\sigma_j^2+\epsilon}$, and over all entries that is $dY\operatorname{diag}(\gamma)\bar D$.$X_{nj}$ reaches $L$ only through $Y_{nj}$; right-multiplying by the diagonal matrices scales column $j$ by $\gamma_j$ and by $(\bar\sigma_j^2+\epsilon)^{-1/2}$.
4. $d\gamma$ and $d\beta$ come from Problem 2 unchanged, with this $\hat X$.Problem 2 used only $Y = \hat X\operatorname{diag}(\gamma) + \mathbf{1}\beta^\top$, not where $\hat X$ came from.
5. $dX = dY\operatorname{diag}(\gamma)\bar D$, that is $dX_{nj} = \gamma_j\,dY_{nj}/\sqrt{\bar\sigma_j^2+\epsilon}$; $d\gamma = (dY\odot\hat X)^\top\mathbf{1}$ and $d\beta = dY^\top\mathbf{1}$ with $\hat X = (X - \mathbf{1}\bar\mu^\top)\bar D$$N\times d$, $d$ and $d$. With fixed statistics the layer is an elementwise affine map, $Y = X\bar D\operatorname{diag}(\gamma) + \mathbf{1}\big(\beta - \operatorname{diag}(\gamma)\bar D\bar\mu\big)^\top$, so examples no longer interact, its gradient has no centring terms, and it can be folded into the weights and bias of the layer before it. This is the gradient used when a network is fine-tuned with batch norm frozen.

Problem 9

1. $\mu = X^\top$, the single row as a column, and $\sigma^2 = 0$.Each column has one entry, which is its own mean and has zero deviation from it.
2. $\hat X = (X - \mathbf{1}\mu^\top)D = 0\cdot D = 0$, and $D = \epsilon^{-1/2}I$ is finite.$X - \mathbf{1}\mu^\top = X - X = 0$; with $\epsilon = 0$ every entry would be $0/0$.
3. $Y = 0\cdot\operatorname{diag}(\gamma) + \mathbf{1}\beta^\top = \beta^\top$.$\mathbf{1}$ has one entry.
4. $Y$ does not depend on $X$ or $\gamma$ at all, so $dX = 0$ and $d\gamma = (dY\odot 0)^\top\mathbf{1} = 0$, while $d\beta = dY^\top\mathbf{1} = dY^\top$.Problem 2. Problem 5 agrees: with $N = 1$ and $\hat x = 0$ it gives $\tfrac{1}{\sqrt\epsilon}(d\hat x - d\hat x - 0) = 0$.
5. With $N = 1$: $\mu = X^\top$, $\sigma^2 = 0$, $\hat X = 0$, $Y = \beta^\top$; $dX = 0$, $d\gamma = 0$, $d\beta = dY^\top$The layer outputs $\beta$ whatever the input, so nothing below it receives a gradient and $\gamma$ never trains. Small batches degrade the same way more gently: the statistics of a few examples are noisy, and every example's gradient is entangled with them.

Problem 10

1. $\mu$ equals that common value, so $x - \mu\mathbf{1} = 0$ and $\sigma^2 = 0$.The mean of equal numbers is that number, and every deviation is zero.
2. $\hat x = 0/\sqrt{0+\epsilon} = 0$.$\epsilon > 0$ makes the denominator $\sqrt\epsilon$; with $\epsilon = 0$ it would be $0/0$.
3. Problem 5 holds at $\sigma^2 = 0$: $dx = \tfrac{1}{N\sqrt\epsilon}\big(N\,d\hat x - (\sum_n d\hat x_n)\mathbf{1} - 0\big)$.Its derivation needed only $\sigma^2+\epsilon > 0$, so the map is smooth here; the last term carries the factor $\hat x = 0$.
4. On a constant feature $\hat x = 0$ and $dx = \dfrac{1}{\sqrt\epsilon}\Big(d\hat x - \dfrac1N\big(\textstyle\sum_n d\hat x_n\big)\mathbf{1}\Big)$; without $\epsilon$ the forward pass would divide $0$ by $0$Distribute the $\tfrac1N$ over the bracket. The gain is $1/\sqrt\epsilon$: about $316$ for $\epsilon = 10^{-5}$, so the centred upstream gradient comes back several hundred times larger, and a tiny spread in a nearly constant feature is blown up to unit scale in the forward pass. For a feature with $\sigma^2 \gg \epsilon$, $\epsilon$ changes $\hat x$ only by a relative $\epsilon/(2\sigma^2)$, and the $\hat x$ component of $dx$ in Problem 6 is of the same small order.