Backpropagation through time for RNN: how to deal with recursively defined gradient updates?

Question

A simplified RNN architecture basically involves the following update

\begin{equation} \begin{cases} h_t & = \phi(w h_{t-1} + v x_t )\\ \hat y_t & = \theta(h_t ) \end{cases} \end{equation} for $t = 1 \ldots, T$, and $w,v$ are scalar parameters, $x_t$ is the input, $h_t$ is the state and $\hat y_t$ is the prediction, $\phi, \theta$ are two activation functions. For simplicity, assume everything is scalar.

I am a bit confused about the derivation of backpropagation for RNN.

---

Suppose we introduce the state $s_t = wh_{t-1} + v x_t$.

Then the RNN update equation reads

\begin{equation} \begin{cases} s_t & = w h_{t-1} + v x_t\\ h_t & = \phi(s_t)\\ \hat y_t & = \theta(h_t ) \end{cases} \end{equation}

Assume we have a loss function $L$ , then by the chain rule,

$$\dfrac{\partial L}{\partial s_t} = \dfrac{\partial L}{\partial h_t}\dfrac{ \partial h_t}{\partial s_t} = \dfrac{\partial L}{\partial h_t} \phi^\prime(s_t)$$

Now,

$$\dfrac{\partial L}{\partial h_t} = \dfrac{\partial L}{\partial {\hat y}_t}\dfrac{ \partial {\hat y}_t}{\partial h_t} + \dfrac{\partial L}{\partial s_{t+1}}\dfrac{ \partial s_{t+1}}{\partial h_t} = \dfrac{\partial L}{\partial {\hat y}_t} \phi^\prime(h_t) + \dfrac{\partial L}{\partial s_{t+1}}w $$

We see that if we were to combine these two equations together, we have,

$$\dfrac{\partial L}{\partial s_t} = \dfrac{\partial L}{\partial h_t}\dfrac{ \partial h_t}{\partial s_t} = \dfrac{\partial L}{\partial h_t} \phi^\prime(s_t) = (\dfrac{\partial L}{\partial {\hat y}_t} \phi^\prime(h_t) + \dfrac{\partial L}{\partial s_{t+1}}w) \phi^\prime(s_t)$$

which has $s_t$ appearing on the left hand side, and $s_{t+1}$ appear on the right hand side. Which means that this gradient update is recursively defined.

---

Question:

How do we find $\dfrac{\partial L}{\partial s_t}$ (unknown) when it is defined in terms of $\dfrac{\partial L}{\partial s_{t+1}}$ (unknown)?

I suspect that for $t = T$, $\dfrac{\partial L}{\partial s_{t+1}}$ vanish $(=0)$, then we have $\dfrac{\partial L}{\partial s_T}$ defined totally in terms of "knowns". Then each of the previous $\dfrac{\partial L}{\partial s_t}$ is solved backwards (dynamic programming). Is this correct?

Answer 1

To compute the gradients , first think that we unfold the RNN through time as below:

Though the notation is different, the essence of your problem can be understood very well with this figure.

TO compute gradients, we start from the last time step. $t = \textit{T}$:

$$ \begin{align} \dfrac{\partial L}{\partial h_{T}} &= \dfrac{\partial L}{\partial {\hat y}_T}\dfrac{ \partial {\hat y}_T}{\partial h_T} \\ &= \dfrac{\partial L}{\partial {\hat y}_T} \theta^{'}(h_{T}) \end{align} $$

$$ \frac{∂L}{∂s_{T}}=\frac{∂L}{∂h_{T}} \frac{∂h_{T}}{∂s_{T}}=\dfrac{\partial L}{\partial {\hat y}_T} \theta^{'}(h_{T}) ϕ′(s_{T}) $$ where $\frac{\partial L}{\partial {\hat y}_T}$ is the loss gradient with respect to prediction which can be computed easily.

Then for $t = \textit{T} -1 $, we use the relation that you've derived for $\frac{∂L}{∂s_{t}}$. Gradients for time intervals $t = 0. \dots , \textit{T} -1$ are computed like this. SO, if you compute gradients backward through time you can compute $\frac{∂L}{∂s_{t}}$'s as $\frac{∂L}{∂s_{t+1}}$ would be known to you(Your derivation is for $t = 0. \dots , \textit{T} -1$).

The gradients outside the time intervals are assumed to be zero for this procedure($t > \textit{T}$).

So, you should define training time steps carefully.