Understanding Backpropagation Through Time (BPTT)
BPTT backpropagates gradients through temporal dependencies in an RNN. This post explains how BPTT works and how to compute its gradients.
An RNN uses the hidden state from the previous time step to compute the current hidden state.
\[a_t=Ux_t+Vh_{t-1}+b\] \[h_t=\tanh(a_t)\]Here, $U,V,b$ are learnable parameters shared across all time steps.
BPTT (Backpropagation Through Time) unrolls an RNN through time and backpropagates gradients across its time steps.
1. Unrolling an RNN through time
Consider an RNN over three time steps.
Each hidden state is used in the computation at the next time step.
\[h_0\rightarrow h_1\rightarrow h_2\rightarrow h_3\]In tasks such as language modeling, where the model makes a prediction at every time step, we can define the total loss as
\[L=L_1+L_2+L_3\]The important point is that $V$ appears at all three time steps.
To compute the total gradient with respect to $V$, we must sum its contributions to the loss across all time steps.
2. The gradient we want to compute
Our goal is to compute
\[\frac{\partial L}{\partial V}\]$V$ is reused in the computations of $a_1,a_2,a_3$.
Adding the contributions from each time step gives
\[\frac{\partial L}{\partial V} = \left.\frac{\partial L}{\partial V}\right|_{t=1} + \left.\frac{\partial L}{\partial V}\right|_{t=2} + \left.\frac{\partial L}{\partial V}\right|_{t=3}\]Each term represents the contribution to the total loss from $V$ as used at that time step.
To compute these contributions, we first need the gradient of the total loss with respect to each $a_t$.
\[\frac{\partial L}{\partial a_1}, \qquad \frac{\partial L}{\partial a_2}, \qquad \frac{\partial L}{\partial a_3}\]Why $a_t$?
Because $V$ is used directly in the operation
\[a_t=Ux_t+Vh_{t-1}+b\]Once we know the gradient backpropagated to $a_t$, we can use the chain rule to compute the contribution to the gradient of $V$ at that time step.
Let’s compute these gradients, starting at the final time step.
3. Backpropagating from the final time step
Start with the gradient with respect to $a_3$.
\[\frac{\partial L}{\partial a_3}\]$a_3$ affects $L_3$ through $h_3$. It has no effect on the earlier losses $L_1,L_2$.
By the chain rule,
\[\frac{\partial L}{\partial a_3} = \frac{\partial L_3}{\partial h_3} \frac{\partial h_3}{\partial a_3}\]Since $h_3=\tanh(a_3)$,
\[\frac{\partial h_3}{\partial a_3}=1-h_3^2\]For clarity, these equations describe a scalar RNN. For a vector RNN, we apply the derivatives elementwise.
We therefore obtain
\[\frac{\partial L}{\partial a_3} = \frac{\partial L_3}{\partial h_3}(1-h_3^2)\]This is the same process as ordinary backpropagation.
Now consider $a_2$
We choose $a_2$ because it is an intermediate time step that affects both the current loss and a future loss.
Through $h_2$, $a_2$ affects the loss along two paths.
The first path leads to the loss at the current time step.
\[a_2\rightarrow h_2\rightarrow L_2\]The second path leads to the loss at a future time step.
\[a_2\rightarrow h_2\rightarrow a_3\rightarrow h_3\rightarrow L_3\]We must add the gradients from these two paths.
\[\begin{aligned} \frac{\partial L}{\partial a_2} ={}& \frac{\partial L_2}{\partial h_2} \frac{\partial h_2}{\partial a_2}\\ &+ \frac{\partial L_3}{\partial h_3} \frac{\partial h_3}{\partial a_3} \frac{\partial a_3}{\partial h_2} \frac{\partial h_2}{\partial a_2} \end{aligned}\]Factoring out $\partial h_2/\partial a_2$, which appears in both terms, gives
\[\frac{\partial L}{\partial a_2} = \left( \frac{\partial L_2}{\partial h_2} + \frac{\partial L_3}{\partial h_3} \frac{\partial h_3}{\partial a_3} \frac{\partial a_3}{\partial h_2} \right) \frac{\partial h_2}{\partial a_2}\]We can reuse a result we have already computed:
\[\frac{\partial L_3}{\partial h_3} \frac{\partial h_3}{\partial a_3} = \frac{\partial L}{\partial a_3}\]For a scalar RNN,
\[\frac{\partial a_3}{\partial h_2}=v\]and
\[\frac{\partial h_2}{\partial a_2}=1-h_2^2\]Thus,
\[\boxed{ \frac{\partial L}{\partial a_2} = \left( \frac{\partial L_2}{\partial h_2} + v\frac{\partial L}{\partial a_3} \right)(1-h_2^2) }\]The important point is that we reuse the gradient from the next time step.
4. Introducing delta
Define the derivative that has appeared repeatedly as
\[\delta_t=\frac{\partial L}{\partial a_t}\]This simplifies the expression we just derived:
\[\delta_2= \left( \frac{\partial L_2}{\partial h_2} +v\delta_3 \right)(1-h_2^2)\]Likewise, at $t=1$,
\[\delta_1= \left( \frac{\partial L_1}{\partial h_1} +v\delta_2 \right)(1-h_1^2)\]Generalizing to a vector RNN gives
\[\boxed{ \delta_t= \left( \frac{\partial L_t}{\partial h_t} + V^\top\delta_{t+1} \right) \odot(1-h_t^2) }\]Each term represents the following:
- $\partial L_t/\partial h_t$: the gradient from the loss at the current time step
- $V^\top\delta_{t+1}$: the gradient propagated from future time steps
- $1-h_t^2$: the derivative of tanh
$\delta_{t+1}$ already includes the effects of all subsequent losses.
We can therefore start backpropagation at the final time step and reuse computed gradients as we move backward.
5. The gradient of the shared weights
Return to our original goal: $\partial L/\partial V$.
At each time step, we compute
\[a_t=Ux_t+Vh_{t-1}+b\]The contribution to the gradient of $V$ at that time step is
\[\delta_t h_{t-1}^\top\]Summing these contributions gives the total gradient:
\[\boxed{ \frac{\partial L}{\partial V} = \sum_{t=1}^{T}\delta_t h_{t-1}^\top }\]This involves two calculations.
First, compute $\delta_t$ in reverse order to account for the effects of future losses.
Second, sum the gradient contributions from every time step that uses the same $V$.
This is the core of BPTT.
6. Does this also apply to LSTMs?
LSTMs pass the previous state forward to the next time step and reuse the same parameters across time steps.
BPTT therefore applies to them as well.
Vanilla RNNs compute using the hidden state $h_t$, while LSTMs use both a hidden state $h_t$ and a cell state $c_t$.
\[c_t=f_t\odot c_{t-1}+i_t\odot\tilde c_t\] \[h_t=o_t\odot\tanh(c_t)\]Backpropagation through an LSTM must account for the gradient paths through both states and compute the gradients of the gates as well.
The principle of BPTT remains the same, with a more complex computation graph at each time step.
7. Vanishing gradients
During BPTT, gradient propagation across time steps involves repeated multiplication by recurrent weights and activation derivatives.
If the gradients become very small, this leads to vanishing gradients. If they grow very large, it leads to exploding gradients.
LSTMs use the cell state and gates to improve gradient flow over longer time spans.
Summary
Three points are central to understanding BPTT.
- Each hidden state can affect both the current loss and future losses.
- Use the chain rule to backpropagate gradients, starting at the final time step.
- Sum the gradient contributions because the same parameters are used at multiple time steps.
Vanilla RNNs and LSTMs learn according to the same principle. They differ in the internal computations through which gradients must propagate over time.