Backprop through time: RNN and LSTM gradients

Worked solutions

Problem 1

1. $W h_{t-1} \in \mathbb{R}^m$ for $h_{t-1} \in \mathbb{R}^m$, so $W$ is $m\times m$; $U x_t \in \mathbb{R}^m$ for $x_t \in \mathbb{R}^d$, so $U$ is $m\times d$; $b \in \mathbb{R}^m$.All three terms of $a_t$ are added to give the $m$ entries of $a_t$, so each must be an $m$-vector.
2. $V$ is $k\times m$ and $c \in \mathbb{R}^k$.$z_t = Vh_t + c$ maps the $m$-vector $h_t$ to the $k$-vector $z_t$.
3. The count is $m^2 + md + m + km + k$.The sizes of the five arrays. None of them has a time index, because the same arrays are reused at every step.
4. $h_t = \tanh(W h_{t-1} + U x_t + b)$ and, unrolling, $h_{t-1}$ depends on $h_0, x_1, \dots, x_{t-1}$.Each step adds one new input to what the previous state already carried; by induction from $h_1 = \tanh(Wh_0 + Ux_1 + b)$.
5. $W: m\times m$, $U: m\times d$, $b: m$, $V: k\times m$, $c: k$; $m^2 + md + m + km + k$ parameters, whatever $T$ is; $h_t$ depends on $h_0$ and $x_1, \dots, x_t$, and on no later inputThe parameter count does not grow with the sequence length, which is the point of sharing the weights; the price is that every weight gradient is a sum over steps (Problem 4).

Problem 2

1. $\partial a_t/\partial h_{t-1} = W$ and $\partial a_t/\partial x_t = U$.$a_t$ is affine in each of $h_{t-1}$ and $x_t$, and the Jacobian of $Mv + \text{const}$ is $M$.
2. $\partial h_t/\partial a_t = \operatorname{diag}(1 - h_t\odot h_t)$.$\tanh$ is elementwise, so its Jacobian is diagonal, with $\tanh'(a_{t,i}) = 1 - h_{t,i}^2$ on the diagonal; the forward pass has already computed $h_t$.
3. $\partial h_t/\partial h_{t-1} = \operatorname{diag}(1 - h_t\odot h_t)\,W$ and $\partial h_t/\partial x_t = \operatorname{diag}(1 - h_t\odot h_t)\,U$The chain rule in numerator layout puts the outer Jacobian on the left: $(m\times m)(m\times m)$ and $(m\times m)(m\times d)$, the shapes of an $m$-vector differentiated by an $m$-vector and by a $d$-vector.

Problem 3

1. $h_T$ reaches $L$ only through $z_T$, so $g_T = (\partial z_T/\partial h_T)^\top e_T = V^\top e_T$.There is no step $T + 1$, and the later losses are the only other way $h_T$ could matter.
2. For $t < T$, $h_t$ reaches $L$ through $z_t$ and through $a_{t+1}$, and through nothing else.$h_t$ is used in exactly two places in the forward pass: the output of step $t$ and the pre-activation of step $t + 1$.
3. $g_t = V^\top e_t + (\partial a_{t+1}/\partial h_t)^\top\delta_{t+1} = V^\top e_t + W^\top\delta_{t+1}$.The chain rule adds the two routes of step 2; $\delta_{t+1}$ is the total gradient at $a_{t+1}$, so it already carries every step after $t + 1$.
4. $\delta_t = (\partial h_t/\partial a_t)^\top g_t = (1 - h_t\odot h_t)\odot g_t$.$a_t$ reaches $L$ only through $h_t$, and a diagonal Jacobian acts as an elementwise product (Problem 2, step 2).
5. $g_T = V^\top e_T$; $g_t = V^\top e_t + W^\top\delta_{t+1}$ for $t < T$; $\delta_t = (1 - h_t\odot h_t)\odot g_t$The recursion runs backwards from $t = T$ to $t = 1$, one step per forward step, so the backward pass costs about as much as the forward pass; it needs every $h_t$, which is why training stores the whole sequence of states.

Problem 4

1. Give step $t$ its own copy $W^{(t)}$ of $W$, used only in $a_t$; then $\nabla_W L = \sum_t \nabla_{W^{(t)}}L$ evaluated at $W^{(t)} = W$.$L$ depends on $W$ only through its $T$ uses, and the chain rule adds the contributions of every use: this is the unrolled network with its shared weights untied.
2. $W^{(t)}_{ij}$ appears only in $a_{t,i}$, with coefficient $h_{t-1,j}$, so $\partial L/\partial W^{(t)}_{ij} = \delta_{t,i}\,h_{t-1,j}$, that is $\nabla_{W^{(t)}}L = \delta_t h_{t-1}^\top$.The copy $W^{(t)}$ reaches $L$ only through $a_t$, and $\delta_t$ is the total gradient there; $h_{t-1}$ is an input to step $t$, so it is not differentiated here, and its own dependence on $W$ is the other copies' contribution.
3. In the same way $\nabla_{U^{(t)}}L = \delta_t x_t^\top$ and $\nabla_{b^{(t)}}L = \delta_t$.$U^{(t)}_{ij}$ multiplies $x_{t,j}$ into $a_{t,i}$, and $b^{(t)}$ is added to $a_t$ with coefficient $1$.
4. $\nabla_{V^{(t)}}L = e_t h_t^\top$ and $\nabla_{c^{(t)}}L = e_t$.$V^{(t)}$ and $c^{(t)}$ appear only in $z_t$, which reaches $L$ only through $\ell_t$, so the gradient at $z_t$ is $e_t$; the pattern is the outer product of the one-hidden-layer page.
5. $\nabla_W L = \sum_t \delta_t h_{t-1}^\top$, $\nabla_U L = \sum_t \delta_t x_t^\top$, $\nabla_b L = \sum_t \delta_t$, $\nabla_V L = \sum_t e_t h_t^\top$, $\nabla_c L = \sum_t e_t$Sums of the copies' gradients, steps 1 to 4, each with the shape of its parameter: $\delta_t h_{t-1}^\top$ is $(m\times1)(1\times m)$. The output weights see only the local $e_t$; the recurrent weights see the $\delta_t$, which carry the future.

Problem 5

1. $h_1 = \tfrac12\cdot0 + 1\cdot1 = 1$ and $h_2 = \tfrac12\cdot1 + 1\cdot2 = \tfrac52$.The forward pass, which the backward pass needs.
2. $\delta_2 = g_2 = h_2 = \tfrac52$.$\tfrac{d}{dh}\tfrac12h^2 = h$, and without a $\tanh$ the factor $1 - h\odot h$ becomes $1$, so $\delta_t = g_t$.
3. $\delta_1 = g_1 = w\,\delta_2 = \tfrac54$.Problem 3, step 3, with no loss at step $1$ ($e_1 = 0$) and $W^\top = w$.
4. $\partial L/\partial w = \delta_2 h_1 + \delta_1 h_0 = \tfrac52\cdot1 + \tfrac54\cdot0 = \tfrac52$.Problem 4: one term per step, each pairing $\delta_t$ with the state that came into step $t$.
5. $\partial L/\partial u = \delta_2 x_2 + \delta_1 x_1 = \tfrac52\cdot2 + \tfrac54\cdot1 = \tfrac{25}4$.Problem 4 with $x_t$ in place of $h_{t-1}$.
6. Directly: $h_2 = w\,u x_1 + u x_2 = u(w + 2)$, so $L = \tfrac12u^2(w + 2)^2$, $\partial L/\partial w = u^2(w + 2) = \tfrac52$ and $\partial L/\partial u = u(w + 2)^2 = \tfrac{25}4$.Substituting $h_1$ into $h_2$ removes the recursion; the two routes agree.
7. $\partial L/\partial w = \tfrac52$ and $\partial L/\partial u = \tfrac{25}4$The $\tfrac54$ in $\partial L/\partial u$ is the contribution of $x_1$ through two steps; a gradient that stopped at the last step would be $5$.

Problem 6

1. $\dfrac{\partial h_T}{\partial h_t} = \dfrac{\partial h_T}{\partial h_{T-1}}\,\dfrac{\partial h_{T-1}}{\partial h_{T-2}}\cdots\dfrac{\partial h_{t+1}}{\partial h_t}$.$h_T$ depends on $h_t$ only through the chain $h_{t+1}, \dots, h_{T-1}$, with the inputs $x_s$ fixed; numerator layout puts the latest step on the left.
2. Each factor is $D_s W$ with $D_s = \operatorname{diag}(1 - h_s\odot h_s)$, $s = t+1, \dots, T$.Problem 2.
3. $\|D_s\|_2 \le 1$.$D_s$ is diagonal with entries $1 - h_{s,i}^2$, and $h_{s,i} = \tanh(a_{s,i}) \in (-1, 1)$ puts every entry in $(0, 1]$.
4. $\big\|\prod_s D_s W\big\|_2 \le \prod_s \|D_s\|_2\,\|W\|_2 \le \|W\|_2^{\,T-t}$.Submultiplicativity applied to the $2(T - t)$ factors, then step 3.
5. $\dfrac{\partial h_T}{\partial h_t} = D_T W\,D_{T-1} W\cdots D_{t+1} W$ with $D_s = \operatorname{diag}(1 - h_s\odot h_s)$, and $\big\|\partial h_T/\partial h_t\big\|_2 \le \|W\|_2^{\,T-t}$The loss at step $T$ reaches $g_t$ as $(\partial h_T/\partial h_t)^\top V^\top e_T$, so with $\|W\|_2 < 1$ its size shrinks at least geometrically with the distance $T - t$: the vanishing gradient, and saturated units ($D_s$ near $0$) shrink it further. With $\|W\|_2 > 1$ the bound allows the product to grow geometrically, the exploding gradient, which is why RNN training clips gradient norms. The bound is one-sided: it guarantees vanishing, it does not guarantee explosion.

Problem 7

1. $e_t = 0$ for $t < T$ and $e_T = s - y$.Steps before $T$ have no loss of their own; the softmax page gives the gradient of cross-entropy at the logits.
2. $g_T = V^\top(s - y)$ and $g_t = W^\top\delta_{t+1}$ for $t < T$.Problem 3 with step 1: only the route through the next step is left.
3. $\delta_T = (1 - h_T\odot h_T)\odot V^\top(s - y)$ and $\delta_t = (1 - h_t\odot h_t)\odot W^\top\delta_{t+1}$.Problem 3, step 4.
4. $\nabla_V L = \sum_t e_t h_t^\top = (s - y)h_T^\top$.Problem 4; every term with $t < T$ is $0$.
5. $\delta_T = (1 - h_T\odot h_T)\odot V^\top(s - y)$, $\delta_t = (1 - h_t\odot h_t)\odot W^\top\delta_{t+1}$ for $t < T$; $\nabla_W L = \sum_t \delta_t h_{t-1}^\top$ and $\nabla_V L = (s - y)\,h_T^\top$The whole training signal for the early steps arrives through Problem 6's product, so this setup is where vanishing gradients bite hardest: the first inputs of a long sequence barely move the weights.

Problem 8

1. $x_t$ is used only in $a_t$, so $\nabla_{x_t}L = (\partial a_t/\partial x_t)^\top\delta_t = U^\top\delta_t$.$x_t$ appears in no other step, and $\delta_t$ is the total gradient at $a_t$.
2. $h_0$ is used only in $a_1$, so $\nabla_{h_0}L = (\partial a_1/\partial h_0)^\top\delta_1 = W^\top\delta_1$.There is no output $z_0$, so the route through a step's own output is absent: Problem 3, step 3, with only the second term.
3. $\nabla_{x_t}L = U^\top\delta_t$ and $\nabla_{h_0}L = W^\top\delta_1$$\nabla_{x_t}L$ continues the backward pass into whatever produced $x_t$: an embedding table or a lower RNN layer. $\nabla_{h_0}L$ is needed when $h_0$ is learned, or when it is another network's output, as in an encoder–decoder where the decoder starts from the encoder's last state; with $h_0$ fixed at zero it is discarded.

Problem 9

1. $c_t$ reaches $L$ through $h_t$ and through $c_{t+1}$, so $dc_t = \bar c_t + dh_t\odot o_t\odot\big(1 - \tanh(c_t)\odot\tanh(c_t)\big)$.Two routes, added: $\bar c_t$ is the second by definition, and the first goes through $h_{t,i} = o_{t,i}\tanh(c_{t,i})$, whose derivative in $c_{t,i}$ is $o_{t,i}(1 - \tanh^2 c_{t,i})$.
2. $do_t = \tanh(c_t)\odot dh_t$.$o_t$ appears only in $h_t$, entry by entry, with coefficient $\tanh(c_t)$.
3. $c_{t,j} = f_{t,j}c_{t-1,j} + i_{t,j}\tilde c_{t,j}$ for each $j$, so the gradient at each of the four factors is the other factor of its product times $dc_t$.Every term is a product of two entries with the same index, and $c_t$ is the only place they appear in this cell.
4. $dc_t = \bar c_t + dh_t\odot o_t\odot(1 - \tanh(c_t)\odot\tanh(c_t))$; $f_t\odot dc_t$ to $c_{t-1}$; $df_t = c_{t-1}\odot dc_t$, $di_t = \tilde c_t\odot dc_t$, $d\tilde c_t = i_t\odot dc_t$, $do_t = \tanh(c_t)\odot dh_t$$f_t\odot dc_t$ is the $\bar c_{t-1}$ of the previous step, so the cell gradients run backwards by the recursion $\bar c_{t-1} = f_t\odot dc_t$. The gate gradients then go through each gate's sigmoid or $\tanh$ into the weights, as in Problem 4, and into $dh_{t-1}$.

Problem 10

1. With the gates fixed, $\partial c_{s}/\partial c_{s-1} = \operatorname{diag}(f_s)$.$c_s = f_s\odot c_{s-1} + i_s\odot\tilde c_s$ is elementwise in $c_{s-1}$ with coefficient $f_s$, and the second term does not involve $c_{s-1}$.
2. $\dfrac{\partial c_T}{\partial c_t} = \operatorname{diag}(f_T)\operatorname{diag}(f_{T-1})\cdots\operatorname{diag}(f_{t+1})$.The chain rule along $c_{t+1}, \dots, c_T$, as in Problem 6, step 1.
3. A product of diagonal matrices is diagonal, with the products of the entries on its diagonal.Diagonal matrices multiply entry by entry.
4. $\partial c_T/\partial c_t = \operatorname{diag}(f_{t+1}\odot f_{t+2}\odot\cdots\odot f_T)$Unlike Problem 6 there is no $W$ and no $\tanh'$ in the product, only gates the network sets for itself: where it keeps $f \approx 1$, the gradient reaches step $t$ almost undiminished however large $T - t$ is. This path is why LSTMs learn long-range dependencies that vanilla RNNs cannot; the full gradient also has paths through the gates and $h$, which can still vanish.