Batched backprop: dense layers on a minibatch

Worked solutions

Problem 1

1. $XW$ is $(N \times d)(d \times m) = N \times m$.A product needs the inner dimensions to agree, so $W$ has $d$ rows, one per input feature, and its $m$ columns are the output features.
2. $\mathbf{1}b^\top$ is $(N \times 1)(1 \times m) = N \times m$, and every row of it is $b^\top$.$(\mathbf{1}b^\top)_{nj} = 1 \cdot b_j$. This is the matrix form of adding $b$ to every row, which array code does by broadcasting.
3. Row $n$ of $XW$ is $x^{(n)\top}W$.Row $n$ of a product is row $n$ of the left factor times the right factor.
4. $dA$ has the shape of $A$ for every $A$.It has one entry $\partial L/\partial A_{\dots}$ per entry of $A$, so that $A - \eta\,dA$ is defined.
5. $X: N\times d$, $W: d\times m$, $b: m$, $Y: N\times m$, and each $dA$ has the shape of $A$; row $n$ of $Y$ is $x^{(n)\top}W + b^\top$Row $n$ uses only row $n$ of $X$ and the shared $W$ and $b$: the batch is $N$ copies of the single-example layer $y = W^\top x + b$, one per row, sharing the parameters.

Problem 2

1. $Y_{nj} = \sum_i X_{ni}W_{ij} + b_j$.Index form of $Y = XW + \mathbf{1}b^\top$.
2. $W_{ij}$ appears in $Y_{nj}$ for every $n$, with coefficient $X_{ni}$, and in no entry outside column $j$.Column $j$ of $W$ builds only column $j$ of $Y$, and every example uses the same $W$.
3. $\dfrac{\partial L}{\partial W_{ij}} = \sum_n \dfrac{\partial L}{\partial Y_{nj}}\,X_{ni} = \sum_n X_{ni}\,dY_{nj}$.The chain rule sums over every entry of $Y$ that contains $W_{ij}$; by step 2 that is one entry per example.
4. $\sum_n X_{ni}\,dY_{nj} = (X^\top dY)_{ij}$.$(X^\top)_{in} = X_{ni}$, and a sum over the shared index $n$ is the definition of the matrix product.
5. $dW = X^\top dY$ ($d\times m$)$(d \times N)(N \times m) = d \times m$, the shape of $W$. The batch index is the one summed over, which is why $N$ is the inner dimension of the product.

Problem 3

1. $b_j$ appears in $Y_{nj}$ for every $n$, with coefficient $1$.$\mathbf{1}b^\top$ copies $b^\top$ into every row (Problem 1, step 2).
2. $\partial L/\partial b_j = \sum_n dY_{nj}$.The chain rule sums over the $N$ entries that contain $b_j$, each with derivative $1$.
3. $db = dY^\top\mathbf{1}$ ($m$)$(m \times N)(N \times 1) = m \times 1$, the shape of $b$; multiplying by $\mathbf{1}$ sums each column of $dY$, which array code writes as a sum over axis $0$. A bias added to every row has its gradient summed over every row.

Problem 4

1. $X_{ni}$ appears in $Y_{nj}$ for every $j$, with coefficient $W_{ij}$, and in no other row of $Y$.Example $n$ feeds only row $n$: in a dense layer the examples do not interact.
2. $\dfrac{\partial L}{\partial X_{ni}} = \sum_j dY_{nj}\,W_{ij} = (dY\,W^\top)_{ni}$.The chain rule sums over the entries of row $n$ that contain $X_{ni}$, one per output feature; $(W^\top)_{ji} = W_{ij}$.
3. $dX = dY\,W^\top$ ($N\times d$)$(N \times m)(m \times d) = N \times d$, the shape of $X$. The sum is over output features, not over the batch: each example gets its own input gradient, and row $n$ is $(W\,dy^{(n)})^\top$, the single-example rule for $y = W^\top x + b$.

Problem 5

1. For one example, $y = W^\top x + b$ has $y_j = \sum_i x_i W_{ij} + b_j$, so its gradient with respect to $W$ has entries $x_i\,dy_j$: it is $x\,dy^\top$, $d \times m$.$W_{ij}$ appears only in $y_j$, with coefficient $x_i$; an outer product $ab^\top$ has entries $a_i b_j$.
2. $dW = \sum_n x^{(n)}\,dy^{(n)\top}$.Every example uses the same $W$, so the loss reaches $W$ through all $N$ rows and the contributions add.
3. For matrices $A$ and $C$ with rows $a^{(n)\top}$ and $c^{(n)\top}$, $\sum_n a^{(n)}c^{(n)\top} = A^\top C$.$(A^\top C)_{ij} = \sum_n A_{ni}C_{nj} = \sum_n a^{(n)}_i c^{(n)}_j$, the $(i,j)$ entry of the sum of outer products.
4. Flatten $Y$ to $Nm$ entries and $W$ to $dm$; the Jacobian is $Nm \times dm$, and entry $(Y_{nj}, W_{ij'})$ is $X_{ni}$ if $j' = j$ and $0$ otherwise.$Y_{nj}$ depends only on column $j$ of $W$ (Problem 2, step 2), so for each of the $Nm$ outputs only $d$ of the $dm$ weights give a nonzero entry: $Nm \cdot d$ in total.
5. $\sum_n x^{(n)}\,dy^{(n)\top} = X^\top dY$; the Jacobian of $Y$ with respect to $W$ has $Nm \times dm$ entries, only $Ndm$ of them nonzero, and the matrix product never forms itStep 3 with $A = X$, $C = dY$. $X^\top dY$ costs $Ndm$ multiply-adds, the same as the forward product $XW$, while the full Jacobian has $Ndm^2$ entries: at $N = 512$ and $d = m = 4096$, about $3.5 \times 10^{13}$. The matrix product is the vector-Jacobian product of the Jacobians page, done for the whole batch at once.

Problem 6

1. $L = \tfrac1N\sum_n \ell_n$ with $\ell_n = \tfrac12\|y^{(n)} - t^{(n)}\|^2$, where $t^{(n)\top}$ is row $n$ of $T$.The squared Frobenius norm is the sum of the squared entries, which is the sum of the squared row norms.
2. $\partial L/\partial Y_{nj} = \tfrac1N(Y_{nj} - T_{nj})$.Only $\ell_n$ contains row $n$ of $Y$, it carries the factor $\tfrac1N$, and $\tfrac12u^2$ differentiates to $u$.
3. $dY = \tfrac1N(Y - T)$, $N \times m$.Step 2 for every entry.
4. $dW = X^\top dY$ and $db = dY^\top\mathbf{1}$.Problems 2 and 3 hold for any $dY$; they never used where it came from.
5. $dY = \tfrac1N(Y - T)$; $dW = \tfrac1N X^\top(Y - T)$; $db = \tfrac1N(Y - T)^\top\mathbf{1}$The $\tfrac1N$ enters once, in $dY$, and every later gradient is linear in $dY$, so it carries the factor without applying it again. $dW$ is the average of the single-example gradients; with a summed loss it would be $N$ times larger.

Problem 7

1. $H_{nj}$ depends only on $Z_{nj}$.$\tanh$ is applied to each entry on its own, in every row.
2. $\partial L/\partial Z_{nj} = dH_{nj}\,\tanh'(Z_{nj})$.$Z_{nj}$ reaches $L$ only through $H_{nj}$, so the chain rule has one term.
3. $\tanh'(u) = 1 - \tanh^2(u)$, so $\tanh'(Z_{nj}) = 1 - H_{nj}^2$.Differentiate $\tanh = \sinh/\cosh$ with the quotient rule and use $\cosh^2 - \sinh^2 = 1$; the forward pass already saved $H$.
4. $dZ = (1 - H\odot H)\odot dH$ ($N\times m$)$1 - H\odot H$ means $1$ minus each entry of $H\odot H$. The Jacobian of $H$ with respect to $Z$ is $Nm \times Nm$ and diagonal, with these entries on the diagonal, so multiplying by it is an elementwise product: the batch changes the shape, not the rule.

Problem 8

1. Let $M = \sum_n \mu_n = \mathbf{1}^\top\mu$, the number of real rows, so $L = \tfrac1M\sum_n \mu_n\ell_n$.$M \ge 1$ by assumption, and it does not depend on $Z$, so it is a constant for the derivative.
2. Row $n$ of $Z$ feeds only $\ell_n$, and $\nabla_{z^{(n)}}\ell_n = s^{(n)} - t^{(n)}$, with $z^{(n)\top}$, $s^{(n)\top}$, $t^{(n)\top}$ row $n$ of $Z$, $S$, $T$.The softmax is taken row by row; the gradient is the softmax page's Problem 5 for example $n$.
3. Row $n$ of $dZ$ is $\tfrac{\mu_n}{M}\,(s^{(n)} - t^{(n)})^\top$.$\ell_n$ enters $L$ with coefficient $\mu_n/M$.
4. Multiplying row $n$ of a matrix by $\mu_n$, for every $n$, is multiplying it on the left by $\operatorname{diag}(\mu)$.$(\operatorname{diag}(\mu)A)_{ni} = \mu_n A_{ni}$.
5. $dZ = \tfrac{1}{\mathbf{1}^\top\mu}\operatorname{diag}(\mu)(S - T)$; with every $\mu_n = 1$ it is $\tfrac1N(S - T)$With every row real, $\operatorname{diag}(\mu) = I$ and $\mathbf{1}^\top\mu = N$: the softmax page's batched result. Padding rows get a zero gradient whatever their logits, and dividing by the number of real rows, not by $N$, keeps the loss on the same scale however much padding a batch carries. In code $\operatorname{diag}(\mu)$ is never built: the mask multiplies each row.

Problem 9

1. $L = \tfrac1N\sum_n \ell_n = \tfrac1N\sum_j \sum_{n \in \text{batch } j}\ell_n$.Group the sum by micro-batch; each example is in exactly one.
2. $\sum_{n \in \text{batch } j}\ell_n = B\,L^{(j)}$.$L^{(j)}$ is the mean of the $B$ losses in micro-batch $j$.
3. $L = \tfrac BN\sum_j L^{(j)} = \tfrac1K\sum_j L^{(j)}$.$N = KB$.
4. $\nabla_W L = \tfrac1K\sum_j \nabla_W L^{(j)}$.The gradient of a sum is the sum of the gradients, and $\tfrac1K$ is a constant.
5. $L = \tfrac1K\sum_j L^{(j)}$, so $\nabla_W L = \tfrac1K\sum_j \nabla_W L^{(j)}$Accumulating the $K$ micro-batch gradients and dividing once by $K$ reproduces the full-batch gradient, up to rounding. This needs the examples not to interact, which holds for every layer on this page; batch normalisation, which mixes examples, is the exception. With sizes $B_j$ the same steps give weights $B_j/N$ in place of $\tfrac1K$.

Problem 10

1. $dW = \sum_{n,t} x^{(n,t)}\,dy^{(n,t)\top}$ and $db = \sum_{n,t} dy^{(n,t)}$.Every position of every sequence uses the same $W$ and $b$, so each is an example in the sense of Problems 3 and 5, and their contributions add.
2. Let $\tilde X \in \mathbb{R}^{NT\times d}$ have one row $x^{(n,t)\top}$ per pair $(n,t)$, and $d\tilde Y \in \mathbb{R}^{NT\times m}$ the rows $dy^{(n,t)\top}$ in the same order.Stacking the pairs as rows turns the double sum into a single sum over the $NT$ rows; any order works if both use the same one.
3. $dW = \tilde X^\top d\tilde Y$ and $db = d\tilde Y^\top\mathbf{1}$, where $\tilde X$ ($NT\times d$) and $d\tilde Y$ ($NT\times m$) stack every position of every sequence as a rowProblem 5, step 3 and Problem 3 applied to the $NT$ rows: $(d \times NT)(NT \times m) = d \times m$. A position-wise layer treats positions exactly like examples, so a transformer's feed-forward and projection layers see a batch of $NT$ rows. If the loss is a mean over all positions, the factor in $d\tilde Y$ is $\tfrac1{NT}$, not $\tfrac1N$.