Next Word Prediction Using LSTM and GRU
Gated memory and learning to forget
LSTM Overview
Previously we ended on the note that gradient vanishing poses a real issue when training an RNN-based model. Hochreiter and Schmidhuber (1997) proposed the Long Short-Term Memory (LSTM) architecture to mitigate this issue. It was later refined by Gers, Schmidhuber, and Cummins (2000), whose version is now considered the modern LSTM architecture.
At a high level, LSTM preserves sequential nature by propagating the cell state that updates additively, which prevents 0.xxx from being multiplied by another 0.xxx during backpropagation. This keeps the loss gradient from turning into a value so small that the weights stop updating. It also uses gates, values in the range \((0, 1)\) that control what portion of information flows through.
The Building Blocks
To recap, in an RNN, the partial derivative for the hidden weight at the last time step is computed as a sum over all paths, and each path is the sole carrier of the lesson from its own time step. The path from \(t=1\) is the only route through which the loss can teach the network about a 30-step dependency. Notice what happens if we have 30 sequential steps:
\[ \frac{\partial L_{30}}{\partial W_{hh}} = \frac{\partial L_{30}}{\partial h_{30}} \left[ \underbrace{ \frac{\partial^+ h_{30}}{\partial W_{hh}} }_{\text{used at } t=30} + \underbrace{ \frac{\partial h_{30}}{\partial h_{29}} \frac{\partial^+ h_{29}}{\partial W_{hh}} }_{\text{used at } t=29} + \dotsm + \underbrace{ \frac{\partial h_{30}}{\partial h_{29}} \frac{\partial h_{29}}{\partial h_{28}} \dotsm \frac{\partial h_{2}}{\partial h_{1}} \frac{\partial^+ h_1}{\partial W_{hh}} }_{\text{used at } t=1}\right] \]
By \(t=1\), the product of time step factors can easily shrink to around \(10^{-9}\). The total gradient stays healthy, but it is dominated by the most recent terms. \(t=1\) contributes a billionth of what \(t=30\) contributes. Now the network only optimizes for the latest time steps.
The RNN also has trouble holding onto information in the forward pass. At every step, the hidden state gets multiplied by \(W_{hh}\), mixed with new input, and bounded by \(\tanh\).
LSTM uses a second vector \(c_t\) to carry information across time steps alongside \(h_t\). This lets the cell state update without passing through a weight matrix or a nonlinearity.
\[ c_t = f_t \odot c_{t-1} + i_t \odot \tilde{c}_t, \quad h_t = o_t \odot \tanh(c_t) \]
The terms can be broken down into:
The forget gate: \(f_t = \sigma(x_t W_{xf} + h_{t-1} W_{hf})\) – As the name suggests, it controls what fraction of the accumulated cell state flows into the current time step.
The input gate: \(i_t = \sigma(x_t W_{xi} + h_{t-1} W_{hi})\) – It controls what fraction of the candidate value is written into the cell state.
The candidate value: \(\tilde{c}_t = \tanh(x_t W_{xc} + h_{t-1} W_{hc})\) – It proposes the new content to be written, computed from the current input and previous hidden state.
The output gate: \(o_t = \sigma(x_t W_{xo} + h_{t-1} W_{ho})\) – It determines what fraction of the cell state is exposed as the hidden state.
The rest of the forward pass from this point on is identical to RNN (\(h_tW_{hy} \rightarrow \text{softmax} \rightarrow \hat{y}_t\)).
LSTM Forward Pass
At this point, it’s clear that backpropagation is the spine of training a deep learning model. The forward pass is no more than a linear combination with some nonlinearity applied, so we’ll run through it quickly and spend our time where this new architecture pays off. The training sample will be identical to the one used in previous RNN example:
the rise of the machine
Create Arbitrary Weights
\(W_{xc}\) and \(W_{hc}\) are initialized to same values as the RNN’s \(W_{xh}\) and \(W_{hh}\), and \(W_{hy}\) is unchanged. This is deliberate, so the difference is easier to visualize. The candidate \(\tilde{c}_t = \tanh(x_t W_{xc} + h_{t-1} W_{hc})\) has the same form as the RNN’s hidden state update. All the gate matrices are new, since the RNN has nothing corresponding to them.
Embedding
\[ E = \begin{bmatrix} 0.6 & -0.2 \\ 0.4 & 0.7 \\ -0.5 & 0.3 \\ 0.1 & -0.8 \end{bmatrix} \]
Input Weights and Recurrent Weights of Candidate
\[ W_{xc} = \begin{bmatrix} 0.5 & -0.3 \\ 0.2 & 0.4 \end{bmatrix}, \quad W_{hc} = \begin{bmatrix} 0.1 & 0.4 \\ -0.2 & 0.3 \end{bmatrix} \]
Input Weights and Recurrent Weights of the Forget Gate
\[ W_{xf} = \begin{bmatrix} 0.3 & -0.2 \\ 0.4 & 0.1 \end{bmatrix}, \quad W_{hf} = \begin{bmatrix} 0.2 & 0.3 \\ -0.1 & 0.2 \end{bmatrix} \]
Input Weights and Recurrent Weights of the Input Gate
\[ W_{xi} = \begin{bmatrix} 0.5 & 0.2 \\ -0.3 & 0.4 \end{bmatrix}, \quad W_{hi} = \begin{bmatrix} 0.1 & -0.2 \\ 0.3 & 0.2 \end{bmatrix} \]
Input Weights and Recurrent Weights of the Output Gate
\[ W_{xo} = \begin{bmatrix} -0.4 & 0.2 \\ 0.3 & 0.5 \end{bmatrix}, \quad W_{ho} = \begin{bmatrix} 0.2 & -0.3 \\ 0.1 & 0.4 \end{bmatrix} \]
Output Weights
\[ W_{hy} = \begin{bmatrix} 0.7 & -0.1 & 0.2 & -0.5 \\ -0.3 & 0.6 & -0.4 & 0.2 \end{bmatrix} \]
\(t=1\) - input “the”, target “rise”
At index one, the model does not have any sequential memory, so we can say that: \[ x_1 = \begin{bmatrix} 0.6 & -0.2 \end{bmatrix}, \quad h_0 = \begin{bmatrix} 0 & 0 \end{bmatrix}, \quad c_0 = \begin{bmatrix} 0 & 0 \end{bmatrix} \]
Forget Gate:
\[ \begin{aligned} a^f_1 &= x_1 W_{xf} + h_0 W_{hf} = \begin{bmatrix} 0.1 & -0.14 \end{bmatrix} \\ f_1 &= \sigma(a^f_1) = \begin{bmatrix} 0.52498 & 0.46506 \end{bmatrix} \end{aligned} \]
Input Gate:
\[ \begin{aligned} a^i_1 &= x_1 W_{xi} + h_0 W_{hi} = \begin{bmatrix} 0.36 & 0.04 \end{bmatrix} \\ i_1 &= \sigma(a^i_1) = \begin{bmatrix} 0.58904 & 0.51000 \end{bmatrix} \end{aligned} \] Candidate:
\[ \begin{aligned} a^c_1 &= x_1 W_{xc} + h_0 W_{hc} = \begin{bmatrix} 0.26 & -0.26 \end{bmatrix} \\ \tilde{c}_1 &= \tanh(a^c_1) = \begin{bmatrix} 0.25430 & -0.25430 \end{bmatrix} \end{aligned} \]
Output Gate:
\[ \begin{aligned} a^o_1 &= x_1 W_{xo} + h_0 W_{ho} = \begin{bmatrix} -0.3 & 0.02 \end{bmatrix} \\ o_1 &= \sigma(a^o_1) = \begin{bmatrix} 0.42556 & 0.50500 \end{bmatrix} \end{aligned} \]
\[ c_1 = \underbrace{f_1 \odot c_0}_{[0,\;0]} + i_1 \odot \tilde{c}_1 = \begin{bmatrix} 0.14979 & -0.12969 \end{bmatrix} \]
The half-open input gate admitted about half of what the candidate proposed. Now it needs to go through the output gate to get \(h_1\):
\[ h_1 = o_1 \odot \tanh(c_1) = \begin{bmatrix} 0.06327 & -0.06513 \end{bmatrix} \]
The rest is identical to the previous RNN computation, \(z_1 = h_1 W_{hy}\):
\[ \begin{bmatrix} 0.06327 & -0.06513 \end{bmatrix} \begin{bmatrix} 0.7 & -0.1 & 0.2 & -0.5 \\ -0.3 & 0.6 & -0.4 & 0.2 \end{bmatrix} = \begin{bmatrix} 0.063828 & -0.045405 & 0.038706 & -0.044661 \end{bmatrix} \]
\[ z_1 = h_1 W_{hy} = \begin{bmatrix} 0.06383 & -0.04540 & 0.03871 & -0.04466 \end{bmatrix} \]
\[ e^{z_1} = \begin{bmatrix} 1.06591 & 0.95561 & 1.03946 & 0.95632 \end{bmatrix}, \quad \sum_{j=1}^{4} e^{z_{1,j}} = 4.01731 \]
Here \(j\) indexes the position within the logit vector (the raw output before softmax), running from 1 to 4 over the vocabulary.
\[ \hat{y}_1 = \begin{bmatrix} 0.26533 & \mathbf{0.23787} & 0.25875 & 0.23805 \end{bmatrix} \]
\(t=2\) — input “rise”, target “of”
Now the recurrent terms are live:
\[ x_2 = \begin{bmatrix} 0.4 & 0.7 \end{bmatrix}, \quad h_1 = \begin{bmatrix} 0.06327 & -0.06513 \end{bmatrix}, \quad c_1 = \begin{bmatrix} 0.14979 & -0.12969 \end{bmatrix} \]
Forget Gate:
\[ \begin{aligned} a^f_2 &= x_2 W_{xf} + h_1 W_{hf} = \begin{bmatrix} 0.41917 & -0.00404 \end{bmatrix} \\ f_2 &= \sigma(a^f_2) = \begin{bmatrix} 0.60328 & 0.49899 \end{bmatrix} \end{aligned} \]
Input Gate:
\[ \begin{aligned} a^i_2 &= x_2 W_{xi} + h_1 W_{hi} = \begin{bmatrix} -0.02321 & 0.33432 \end{bmatrix} \\ i_2 &= \sigma(a^i_2) = \begin{bmatrix} 0.49420 & 0.58281 \end{bmatrix} \end{aligned} \]
Candidate:
\[ \begin{aligned} a^c_2 &= x_2 W_{xc} + h_1 W_{hc} = \begin{bmatrix} 0.35935 & 0.16577 \end{bmatrix} \\ \tilde{c}_2 &= \tanh(a^c_2) = \begin{bmatrix} 0.34464 & 0.16427 \end{bmatrix} \end{aligned} \]
Output Gate:
\[ \begin{aligned} a^o_2 &= x_2 W_{xo} + h_1 W_{ho} = \begin{bmatrix} 0.05614 & 0.38497 \end{bmatrix} \\ o_2 &= \sigma(a^o_2) = \begin{bmatrix} 0.51403 & 0.59507 \end{bmatrix} \end{aligned} \]
The forget gate keeps roughly 60% of the first cell coordinate and half of the second. The input gate attenuates the candidate to 49% of its value in the first coordinate and 58% in the second:
\[ c_2 = \underbrace{f_2 \odot c_1}_{[0.09037,\;-0.06471]} + \underbrace{i_2 \odot \tilde{c}_2}_{[0.17032,\;0.09574]} = \begin{bmatrix} 0.26069 & 0.03102 \end{bmatrix} \]
\[ h_2 = o_2 \odot \tanh(c_2) = \begin{bmatrix} 0.13105 & 0.01845 \end{bmatrix} \]
\[ z_2 = h_2 W_{hy} = \begin{bmatrix} 0.08620 & -0.00203 & 0.01883 & -0.06183 \end{bmatrix} \]
\[ e^{z_2} = \begin{bmatrix} 1.09002 & 0.99797 & 1.01901 & 0.94004 \end{bmatrix}, \quad \sum_{j=1}^{4} e^{z_{2,j}} = 4.04704 \]
\[ \hat{y}_2 = \begin{bmatrix} 0.26934 & 0.24659 & \mathbf{0.25179} & 0.23228 \end{bmatrix} \]
\(t=3\) — input “of”, target “the”
\[ x_3 = \begin{bmatrix} -0.5 & 0.3 \end{bmatrix}, \quad h_2 = \begin{bmatrix} 0.13105 & 0.01845 \end{bmatrix}, \quad c_2 = \begin{bmatrix} 0.26069 & 0.03102 \end{bmatrix} \]
Forget Gate:
\[ \begin{aligned} a^f_3 &= x_3 W_{xf} + h_2 W_{hf} = \begin{bmatrix} -0.00564 & 0.17301 \end{bmatrix} \\ f_3 &= \sigma(a^f_3) = \begin{bmatrix} 0.49859 & 0.54314 \end{bmatrix} \end{aligned} \]
Input Gate:
\[ \begin{aligned} a^i_3 &= x_3 W_{xi} + h_2 W_{hi} = \begin{bmatrix} -0.32136 & -0.00252 \end{bmatrix} \\ i_3 &= \sigma(a^i_3) = \begin{bmatrix} 0.42034 & 0.49937 \end{bmatrix} \end{aligned} \]
Candidate:
\[ \begin{aligned} a^c_3 &= x_3 W_{xc} + h_2 W_{hc} = \begin{bmatrix} -0.18059 & 0.32796 \end{bmatrix} \\ \tilde{c}_3 &= \tanh(a^c_3) = \begin{bmatrix} -0.17865 & 0.31668 \end{bmatrix} \end{aligned} \]
Output Gate:
\[ \begin{aligned} a^o_3 &= x_3 W_{xo} + h_2 W_{ho} = \begin{bmatrix} 0.31805 & 0.01807 \end{bmatrix} \\ o_3 &= \sigma(a^o_3) = \begin{bmatrix} 0.57885 & 0.50452 \end{bmatrix} \end{aligned} \]
\[ c_3 = \underbrace{f_3 \odot c_2}_{[0.12998,\;0.01685]} + \underbrace{i_3 \odot \tilde{c}_3}_{[-0.07509,\;0.15814]} = \begin{bmatrix} 0.05488 & 0.17499 \end{bmatrix} \]
\[ h_3 = o_3 \odot \tanh(c_3) = \begin{bmatrix} 0.03174 & 0.08740 \end{bmatrix} \]
\[ z_3 = h_3 W_{hy} = \begin{bmatrix} -0.00400 & 0.04926 & -0.02861 & 0.00161 \end{bmatrix} \]
\[ e^{z_3} = \begin{bmatrix} 0.99601 & 1.05050 & 0.97179 & 1.00161 \end{bmatrix}, \quad \sum_{j=1}^{4} e^{z_{3,j}} = 4.01991 \]
\[ \hat{y}_3 = \begin{bmatrix} \mathbf{0.24777} & 0.26132 & 0.24175 & 0.24916 \end{bmatrix} \]
\(t=4\) — input “the”, target “machine”
“the” appears for the second time, so \(x_4\) is the same embedding row as \(x_1\). It is distinguished by values written to \(h_3\) and \(c_3\):
\[ x_4 = \begin{bmatrix} 0.6 & -0.2 \end{bmatrix}, \quad h_3 = \begin{bmatrix} 0.03174 & 0.08740 \end{bmatrix}, \quad c_3 = \begin{bmatrix} 0.05488 & 0.17499 \end{bmatrix} \]
Forget Gate:
\[ \begin{aligned} a^f_4 &= x_4 W_{xf} + h_3 W_{hf} = \begin{bmatrix} 0.09761 & -0.11300 \end{bmatrix} \\ f_4 &= \sigma(a^f_4) = \begin{bmatrix} 0.52438 & 0.47178 \end{bmatrix} \end{aligned} \]
Input Gate:
\[ \begin{aligned} a^i_4 &= x_4 W_{xi} + h_3 W_{hi} = \begin{bmatrix} 0.38939 & 0.05113 \end{bmatrix} \\ i_4 &= \sigma(a^i_4) = \begin{bmatrix} 0.59614 & 0.51278 \end{bmatrix} \end{aligned} \]
Candidate:
\[ \begin{aligned} a^c_4 &= x_4 W_{xc} + h_3 W_{hc} = \begin{bmatrix} 0.24569 & -0.22109 \end{bmatrix} \\ \tilde{c}_4 &= \tanh(a^c_4) = \begin{bmatrix} 0.24087 & -0.21755 \end{bmatrix} \end{aligned} \]
Output Gate:
\[ \begin{aligned} a^o_4 &= x_4 W_{xo} + h_3 W_{ho} = \begin{bmatrix} -0.28491 & 0.04544 \end{bmatrix} \\ o_4 &= \sigma(a^o_4) = \begin{bmatrix} 0.42925 & 0.51136 \end{bmatrix} \end{aligned} \]
\[ c_4 = \underbrace{f_4 \odot c_3}_{[0.02878,\;0.08256]} + \underbrace{i_4 \odot \tilde{c}_4}_{[0.14359,\;-0.11156]} = \begin{bmatrix} 0.17237 & -0.02900 \end{bmatrix} \]
\[ h_4 = o_4 \odot \tanh(c_4) = \begin{bmatrix} 0.07327 & -0.01482 \end{bmatrix} \]
\[ z_4 = h_4 W_{hy} = \begin{bmatrix} 0.05573 & -0.01622 & 0.02058 & -0.03960 \end{bmatrix} \]
\[ e^{z_4} = \begin{bmatrix} 1.05732 & 0.98391 & 1.02080 & 0.96118 \end{bmatrix}, \quad \sum_{j=1}^{4} e^{z_{4,j}} = 4.02320 \]
\[ \hat{y}_4 = \begin{bmatrix} 0.26280 & 0.24456 & 0.25373 & \mathbf{0.23891} \end{bmatrix} \]
Calculate Cross-Entropy Loss
Previously, we saw that the cross-entropy loss can be expressed as:
\[ L_t = -\ln \hat{y}_t[\text{target}] \]
We compute loss at each timestep:
\[ \begin{aligned} L_1 = -\ln(0.23787) = 1.43602 \\ L_2 = -\ln(0.25179) = 1.37916 \\ L_3 = -\ln(0.24777) = 1.39526 \\ L_4 = -\ln(0.23891) = 1.43167 \\ \end{aligned} \]
Then the mean loss is:
\[ \begin{aligned} L &= \frac{1}{4}(L_1 + L_2 + L_3 + L_4)\\ & = \frac{1}{4}(1.43602 + 1.37916 + 1.39526 + 1.43167) = 1.41053 \end{aligned} \]
As a sanity check, an untrained model spreading its mass evenly over a four-word vocabulary would score \(-\ln(1/4) = 1.38629\). We’re sitting just above that, which is what we should expect before any learning has happened.
Backpropagation
Again, we’ve omitted the bias terms, which real-world models typically include, but here they would only add more terms to differentiate without adding much clarity. There are still ten matrices to differentiate, so grab some coffee, settle in, and let’s get started.
Calculate Output Gradient
Similar to RNN output gradient, \(W_{hy}\) stays outside of the hidden state boundary. So the gradient is relatively easy to derive:
\[ \frac{\partial L_t}{\partial W_{hy}} = \frac{\partial L_t}{\partial \hat{y}_t} \cdot \frac{\partial \hat{y}_t}{\partial z_t} \cdot \frac{\partial z_t}{\partial W_{hy}}, \quad \frac{\partial L}{\partial W_{hy}} = \frac{1}{4}\sum_{t=1}^{4} \frac{\partial L_t}{\partial W_{hy}} \]
The loss is the mean over all four time steps:
\[ \delta^z_t = \frac{\partial L}{\partial z_t} = \frac{1}{4}\left(\hat{y}_t - y_t\right) \]
\[ \begin{aligned} \delta^z_4 &= \begin{bmatrix} 0.06570 & 0.06114 & 0.06343 & \mathbf{-0.19027} \end{bmatrix} \\ \delta^z_3 &= \begin{bmatrix} \mathbf{-0.18806} & 0.06533 & 0.06044 & 0.06229 \end{bmatrix} \\ \delta^z_2 &= \begin{bmatrix} 0.06733 & 0.06165 & \mathbf{-0.18705} & 0.05807 \end{bmatrix} \\ \delta^z_1 &= \begin{bmatrix} 0.06633 & \mathbf{-0.19053} & 0.06469 & 0.05951 \end{bmatrix} \end{aligned} \]
Each row sums to zero, which is the softmax-with-cross-entropy signature.
Since \(z_t = h_t W_{hy}\) and \(W_{hy}\) is used at all four steps, its gradient is the sum of four outer products:
\[ h_4^\top \delta^z_4 = \begin{bmatrix} 0.00481 & 0.00448 & 0.00465 & -0.01394 \\ -0.00097 & -0.00091 & -0.00094 & 0.00282 \end{bmatrix} \]
\[ h_3^\top \delta^z_3 = \begin{bmatrix} -0.00597 & 0.00207 & 0.00192 & 0.00198 \\ -0.01644 & 0.00571 & 0.00528 & 0.00544 \end{bmatrix} \]
\[ h_2^\top \delta^z_2 = \begin{bmatrix} 0.00882 & 0.00808 & -0.02451 & 0.00761 \\ 0.00124 & 0.00114 & -0.00345 & 0.00107 \end{bmatrix} \]
\[ h_1^\top \delta^z_1 = \begin{bmatrix} 0.00420 & -0.01206 & 0.00409 & 0.00377 \\ -0.00432 & 0.01241 & -0.00421 & -0.00388 \end{bmatrix} \]
\[ \frac{\partial L}{\partial W_{hy}} = \sum_{t=1}^{4} h_t^\top \delta^z_t = \begin{bmatrix} 0.01187 & 0.00258 & -0.01385 & -0.00059 \\ -0.02049 & 0.01835 & -0.00332 & 0.00546 \end{bmatrix} \]
Candidate Gradients
We’ll examine the candidate gradient first, because it has a similar responsibility to the hidden-state gradient in the RNN. Through the derivation, we can see the key difference that reduces gradient vanishing.
\({\partial L_4}/{\partial W_{xc}}\) Chain Preview
Written out in full, the term runs off the page. It’s easier to group the computational paths and express the chain rule across multiple lines.
Since \(h_4 = o_4 \odot \tanh(c_4)\), the gradient can flow back to both \(c_4\) and \(o_4\).
So we can write the overall structure as:
\[ \frac{\partial L_4}{\partial W_{xc}} = \frac{\partial L_4}{\partial h_4} \left[P_4 + P_3 + P_2 + P_1\right] \]
\(P_t\) denotes the sum of all paths to the use of \(W_{xc}\) at time step \(t\). Similar to RNN workflow, we’ll establish a common factor, but the number of computational paths grows rapidly, making the full derivative impractical to write manually. However, we can still get a glimpse of why LSTM is resilient towards gradient vanishing.
Output to hidden state
\({\partial L_4}/{\partial h_4}\) is shared by every path with \(h_4\) as origin, so we’ll establish it as a reusable common factor:
\[ \delta^h_4 = \frac{\partial L_4}{\partial h_4} = \frac{\partial L_4}{\partial \hat{y}_4} \frac{\partial \hat{y}_4}{\partial z_4} \frac{\partial z_4}{\partial h_4} = \delta^z_4 W_{hy}^{\top}. \]
Path to \(W_{xc}\) at timestep 4
At \(t=4\), there’s only one path back to \(\partial W_{xc}\):
\[ P_4 = \frac{\partial h_4}{\partial c_4} \frac{\partial^{+} c_4}{\partial \tilde{c}_4} \frac{\partial \tilde{c}_4}{\partial W_{xc}} \]
Path to \(W_{xc}\) at timestep 3
At \(t=3\), there are five paths back to \(\partial W_{xc}\):
- \(h_4 \rightarrow c_4 \rightarrow c_3 \rightarrow \tilde{c}_3 \rightarrow W_{xc}\)
- \(h_4 \rightarrow o_4 \rightarrow h_3 \rightarrow c_3 \rightarrow \tilde{c}_3 \rightarrow W_{xc}\)
- \(h_4 \rightarrow c_4 \rightarrow f_4 \rightarrow h_3 \rightarrow c_3 \rightarrow \tilde{c}_3 \rightarrow W_{xc}\)
- \(h_4 \rightarrow c_4 \rightarrow i_4 \rightarrow h_3 \rightarrow c_3 \rightarrow \tilde{c}_3 \rightarrow W_{xc}\)
- \(h_4 \rightarrow c_4 \rightarrow \tilde{c}_4 \rightarrow h_3 \rightarrow c_3 \rightarrow \tilde{c}_3 \rightarrow W_{xc}\)
The partial derivative written out line by line:
\[ \begin{aligned} P_3 &= \frac{\partial h_4}{\partial c_4} \frac{\partial^{+} c_4}{\partial c_3} \frac{\partial^{+} c_3}{\partial \tilde{c}_3} \frac{\partial \tilde{c}_3}{\partial W_{xc}} \\[6pt] &+ \frac{\partial h_4}{\partial o_4} \frac{\partial o_4}{\partial h_3} \frac{\partial h_3}{\partial c_3} \frac{\partial^+ c_3}{\partial \tilde{c}_3} \frac{\partial \tilde{c}_3}{\partial W_{xc}} \\[6pt] &+ \frac{\partial h_4}{\partial c_4} \frac{\partial^{+} c_4}{\partial f_4} \frac{\partial f_4}{\partial h_3} \frac{\partial h_3}{\partial c_3} \frac{\partial^+ c_3}{\partial \tilde{c}_3} \frac{\partial \tilde{c}_3}{\partial W_{xc}} \\[6pt] &+ \frac{\partial h_4}{\partial c_4} \frac{\partial^{+} c_4}{\partial i_4} \frac{\partial i_4}{\partial h_3} \frac{\partial h_3}{\partial c_3} \frac{\partial^+ c_3}{\partial \tilde{c}_3} \frac{\partial \tilde{c}_3}{\partial W_{xc}} \\[6pt] &+ \frac{\partial h_4}{\partial c_4} \frac{\partial^{+} c_4}{\partial \tilde{c}_4} \frac{\partial \tilde{c}_4}{\partial h_3} \frac{\partial h_3}{\partial c_3} \frac{\partial^+ c_3}{\partial \tilde{c}_3} \frac{\partial \tilde{c}_3}{\partial W_{xc}} \end{aligned} \]
As we can see, at \(t=3\), we have to visually stack terms to prevent the equation from overflowing into abyss. To get an intuition of the improvements from RNN, we differentiate a few of them.
First, we look at cell state, which carries sequential data. Comparing it with RNN’s method, we get:
\[ \text{LSTM} = \frac{\partial^{+} c_4}{\partial c_3} = f_4 \quad \text{vs} \quad \text{RNN} = \frac{\partial^{+} h_4}{\partial h_3} = (1-h^2_4)W_{hh} \]
You might notice that even though the term at each timestep is simpler for LSTM, \(f_t\) is still sigmoid activated, which means it’ll be between 0 and 1. Multiply enough of them together and the gradient still has a chance to vanish. However, the forget gate is learned from trainable weights. With each \(f_t\) close to 1, the gradient at \(t=30\) can still survive. For example, \(0.98^{30} = 0.54548\).
If we apply that concept to the section of the path that leads to \(W_{xc}\) through the forget gate, we get:
\[ \begin{aligned} \frac{\partial h_4}{\partial c_4} \frac{\partial^{+} c_4}{\partial f_4} \frac{\partial f_4}{\partial h_3} &= \left[o_4 \odot (1-\tanh^2(c_4))\right] \odot c_3 \odot \left[f_4 \odot \left(1-f_4\right)\right] \odot W_{hf} \end{aligned} \]
The critical section is: \(\left[f_4 \odot \left(1-f_4\right)\right] \odot W_{hf}\)
If we plug in the example number \(0.98\), or anything close to 1, indicating that the gate should retain nearly all of the previous cell-state content, we get a vanishing gradient, because now we have a term that could be extremely small: \(1-0.98=0.02\). Imagine a timestep of 30, this could be reduced to almost nothing.
I hope now we can understand and truly appreciate the nuance here. When the gradient from one path vanishes, we still have the gradient from another path that can hop back through timesteps with little attenuation.
The remaining timesteps
The number of paths grows rapidly for the remaining timesteps, because each earlier cell and hidden state introduces another set of branches. For example, once this path reaches \(c_3\):
\[ \begin{aligned} h_4 \rightarrow c_4 \rightarrow f_4 \rightarrow h_3 \rightarrow c_3 &\rightarrow c_2 \rightarrow \tilde{c}_2 \rightarrow W_{xc} \\ &\rightarrow f_3 \rightarrow h_2 \rightarrow c_2 \rightarrow \tilde{c}_2 \rightarrow W_{xc} \\ &\rightarrow i_3 \rightarrow h_2 \rightarrow c_2 \rightarrow \tilde{c}_2 \rightarrow W_{xc} \\ &\rightarrow \tilde{c}_3 \rightarrow h_2 \rightarrow c_2 \rightarrow \tilde{c}_2 \rightarrow W_{xc} \\ &\dotsm \end{aligned} \]
For \(P_2\), there are 24 paths, and for \(P_1\), there are 115. Writing them all out would be quite a waste of eyesight.
Startover using PyTorch
The stubborn part of me thinks it’s crucial to work through backpropagation by hand first. However, with close to 150 paths, computing each one manually feels like returning to the Stone Age. So far, we’ve shown examples in both PyTorch and TensorFlow, but from this point on, we’ll stick with PyTorch, since it has become the dominant framework in deep learning research and the default choice for new projects.
Step 1 - Import PyTorch and Set Configuration
import torch
import torch.nn.functional as F
# Print up to 5 decimals
torch.set_printoptions(precision=5, sci_mode=False)
dtype = torch.float64
# Helper function to create tensor
def create_parameter(x):
return torch.tensor(x, dtype=dtype, requires_grad=True)Step 2 - Initialize Parameters
We initialize the parameters identically to the previous forward pass section:
# Create Parameters
E = create_parameter([
[ 0.6, -0.2],
[ 0.4, 0.7],
[-0.5, 0.3],
[ 0.1, -0.8],
])
W_xc = create_parameter([
[0.5, -0.3],
[0.2, 0.4],
])
W_hc = create_parameter([
[ 0.1, 0.4],
[-0.2, 0.3],
])
W_xf = create_parameter([
[0.3, -0.2],
[0.4, 0.1],
])
W_hf = create_parameter([
[ 0.2, 0.3],
[-0.1, 0.2],
])
W_xi = create_parameter([
[ 0.5, 0.2],
[-0.3, 0.4],
])
W_hi = create_parameter([
[0.1, -0.2],
[0.3, 0.2],
])
W_xo = create_parameter([
[-0.4, 0.2],
[ 0.3, 0.5],
])
W_ho = create_parameter([
[0.2, -0.3],
[0.1, 0.4],
])
W_hy = create_parameter([
[ 0.7, -0.1, 0.2, -0.5],
[-0.3, 0.6, -0.4, 0.2],
])
# the=0, rise=1, of=2, machine=3
inputs = torch.tensor([0, 1, 2, 0])
targets = torch.tensor([1, 2, 0, 3])
h = torch.zeros(2, dtype=dtype)
c = torch.zeros(2, dtype=dtype)Step 3 - Apply Forward Pass and Calculate Loss
The forward pass is the literal translation from the manual step to pytorch:
# Run forward pass, store pre softmax output in logits array
logits = []
for token in inputs:
x = E[token]
f = torch.sigmoid(x @ W_xf + h @ W_hf)
i = torch.sigmoid(x @ W_xi + h @ W_hi)
c_tilde = torch.tanh(x @ W_xc + h @ W_hc)
o = torch.sigmoid(x @ W_xo + h @ W_ho)
c = f * c + i * c_tilde
h = o * torch.tanh(c)
z = h @ W_hy
logits.append(z)
# Take all the outputs in array and create a new dimension
logits = torch.stack(logits)
# Calculate cross-entropy loss from the logits
loss_t = F.cross_entropy(
logits,
targets,
reduction="none"
)Step 4 - Calculate Gradient
After looking at all those “flipped 6” (\(\partial\)) signs, it’s incredible that all the tedious work can be done with two lines of code:
loss = loss_t.mean()
# Full backpropagation through time
loss.backward()To verify, you can print W_hy.grad and compare it with our manual gradient computation. On my machine, it prints:
∂L/∂W_hy:
tensor([[ 0.01187, 0.00258, -0.01385, -0.00059],
[-0.02049, 0.01835, -0.00332, 0.00546]],
dtype=torch.float64)Step 5 - Training the Model
The beauty of the machine is that it takes a fraction of the manual effort to run the forward pass, backward pass, and parameter update 500 times:
params = [
E,
W_xc, W_hc,
W_xf, W_hf,
W_xi, W_hi,
W_xo, W_ho,
W_hy,
]
optimizer = torch.optim.Adam(params, lr=0.05)
epochs = 500
for epoch in range(1, epochs + 1):
# Hidden state([0 0]) and cell state([0 0])
h = torch.zeros(2, dtype=dtype)
c = torch.zeros(2, dtype=dtype)
logits = []
# Forward pass
for token in inputs:
x = E[token]
f = torch.sigmoid(x @ W_xf + h @ W_hf)
i = torch.sigmoid(x @ W_xi + h @ W_hi)
c_tilde = torch.tanh(x @ W_xc + h @ W_hc)
o = torch.sigmoid(x @ W_xo + h @ W_ho)
c = f * c + i * c_tilde
h = o * torch.tanh(c)
z = h @ W_hy
logits.append(z)
logits = torch.stack(logits)
loss = F.cross_entropy(logits, targets)
# Clear gradients left over from the previous epoch
optimizer.zero_grad()
# Calculate gradients
loss.backward()
# Update all trainable parameters
optimizer.step()
# Print every 50 epochs
if epoch % 50 == 0:
predictions = logits.argmax(dim=1)
print(
f"epoch {epoch:3d} | "
f"loss = {loss.item():.5f} | "
f"predictions = {predictions.tolist()}"
)This will print loss every 50 epochs. Starting from 1.41053, the loss falls by roughly three orders of magnitude, and the predicted indices settle onto the targets [1, 2, 0, 3] — that is, rise, of, the, machine:
epoch 50 | loss = 0.19650 | predictions = [1, 2, 0, 3]
epoch 100 | loss = 0.01642 | predictions = [1, 2, 0, 3]
epoch 150 | loss = 0.00943 | predictions = [1, 2, 0, 3]
epoch 200 | loss = 0.00646 | predictions = [1, 2, 0, 3]
epoch 250 | loss = 0.00478 | predictions = [1, 2, 0, 3]
epoch 300 | loss = 0.00372 | predictions = [1, 2, 0, 3]
epoch 350 | loss = 0.00300 | predictions = [1, 2, 0, 3]
epoch 400 | loss = 0.00248 | predictions = [1, 2, 0, 3]
epoch 450 | loss = 0.00209 | predictions = [1, 2, 0, 3]
epoch 500 | loss = 0.00179 | predictions = [1, 2, 0, 3]It’s definitely worth opening the champagne and feeling a little bit of joy after coming this far, but it’s important to recognize that, in real life, we need to feed many training sequences through the same loop. What we have here is essentially memorization. The model is intentionally overfitting a single training sequence.
Gated Recurrent Unit (GRU)
The Gated Recurrent Unit (GRU) was designed with the same general motivation as the LSTM.
GRU has a simplified architecture with three components:
\[ \text{recurrent state} = h_t, \quad \text{reset gate} = r_t, \quad \text{update gate} = u_t \]
It only uses a single vector \(h_t\) to carry information across time steps:
\[ h_t = u_t \odot h_{t-1} + (1-u_t)\odot\tilde{h}_t \]
The terms can be broken down into:
- The reset gate: \(r_t = \sigma(x_tW_{xr} + h_{t-1}W_{hr})\) - controls how much of the previous hidden state contributes to the new candidate.
- The update gate: \(u_t = \sigma(x_tW_{xu} + h_{t-1}W_{hu})\) - decides how much of the previous hidden state to keep.
- The hidden candidate: \(\tilde{h}_t = \tanh \left( x_tW_{xh} + r_t \odot (h_{t-1}W_{hh}) \right)\)
The flow to output is identical to RNN and LSTM:
\[ z_t = h_tW_{hy}, \quad \hat{y} = softmax(z_t) \]
Understand GRU through backpropagation
Instead of running the forward pass again, we can examine backpropagation from \(t=4\) to \(t=2\) to demonstrate the cleverness of this architecture.
Through the direct path:
This is similar to how an LSTM uses the cell state to create a direct path back to the weight.
\[ \frac{\partial^+ h_4}{\partial h_3} \frac{\partial^+ h_3}{\partial h_2} \frac{\partial^+ h_2}{\partial \tilde h_2} \frac{\partial \tilde h_2}{\partial W_{xh}} = x_2^\top \left[ u_4 \odot u_3 \odot (1-u_2) \odot (1-\tilde h_2^2) \right] \]
If the update gate has a large value, say 0.98 again, the gradient can survive through many timesteps.
Through the update gate \(u_t\):
\[ \begin{aligned} &\frac{\partial h_4}{\partial u_4} \frac{\partial u_4}{\partial h_3} \frac{\partial h_3}{\partial u_3} \frac{\partial u_3}{\partial h_2} \frac{\partial h_2}{\partial \tilde h_2} \frac{\partial \tilde h_2}{\partial W_{xh}} \\ &= x_2^\top \bigg[ \big( h_3 - \tilde{h}_4 \big) \odot u_4 \odot \big( 1 - u_4 \big) W_{hu}^{\top} \\ & \qquad \odot \big( h_2 - \tilde{h}_3 \big) \odot u_3 \odot \big( 1 - u_3 \big) W_{hu}^{\top} \\ & \qquad \odot \big( 1 - u_2 \big) \odot \big( 1 - \tilde{h}_2^2 \big) \bigg] \end{aligned} \]
In contrast to the direct path, 0.98 here can cause the gradient to become small because the sigmoid derivative contains \(1 - u_t\).
Through the reset gate \(r_t\):
\[ \begin{aligned} &\frac{\partial h_4}{\partial \tilde h_4} \frac{\partial \tilde h_4}{\partial r_4} \frac{\partial r_4}{\partial h_3} \frac{\partial h_3}{\partial \tilde h_3} \frac{\partial \tilde h_3}{\partial r_3} \frac{\partial r_3}{\partial h_2} \frac{\partial h_2}{\partial \tilde h_2} \frac{\partial \tilde h_2}{\partial W_{xh}} \\ &= x_2^\top \Bigg[ \Bigg( \Big[ \big( 1 - u_4 \big) \odot \big( 1 - \tilde{h}_4^2 \big) \odot \left( h_3 W_{hh} \right) \odot r_4 \odot \big( 1 - r_4 \big) \Big] W_{hr}^{\top} \\ &\quad\quad \odot \Big[ \left( 1 - u_3 \right) \odot \left( 1 - \tilde{h}_3^2 \right) \odot \left( h_2 W_{hh} \right) \odot r_3 \odot \left( 1 - r_3 \right) \Big] W_{hr}^{\top} \Bigg) \\ &\qquad \odot \left( 1 - u_2 \right) \odot \left( 1 - \tilde{h}_2^2 \right) \Bigg] \end{aligned} \]
If the path through the update gate creates a vanishing gradient, the reset gate path can still carry a stronger gradient.
Through the candidate \(\tilde h_t\):
\[ \begin{aligned} &\frac{\partial h_4}{\partial \tilde h_4} \frac{\partial \tilde h_4}{\partial h_3} \frac{\partial h_3}{\partial \tilde h_3} \frac{\partial \tilde h_3}{\partial h_2} \frac{\partial h_2}{\partial \tilde h_2} \frac{\partial \tilde h_2}{\partial W_{xh}} \\ &= x_2^\top \Bigg[ \Bigg( \Big[ \Big( 1 - u_4 \Big) \odot \Big( 1 - \tilde{h}_4^2 \Big) \odot r_4 \Big] W_{hh}^{\top} \odot \Big( 1 - u_3 \Big) \odot \Big( 1 - \tilde{h}_3^2 \Big) \odot r_3 \Bigg) W_{hh}^{\top} \\ &\qquad \odot \big( 1 - u_2 \big) \odot \big( 1 - \tilde{h}_2^2 \big) \Bigg] \end{aligned} \]
This path resembles the recurrent path of a vanilla RNN. With an update gate value of 0.98, it makes the candidate path small, but at the same time 0.98 will appear in the direct path unchanged.
When the GRU decides to preserve its previous hidden state, it suppresses the candidate path while opening a nearly unchanged direct path backward through time.
Wrap up
LSTM and GRU are obviously quite useful for creating sequential models. GRU can offer greater simplicity and better hardware efficiency, while LSTM provides more explicit control over memory. Since neither consistently outperforms the other across all tasks, in my opinion, the final choice is often preference. It’s also worth noting that transformers have recently become the dominant choice for language models, which we’ll explore in the next article.
Additional Resources
- Understanding LSTM Networks — Christopher Olah’s illustrated walkthrough, the classic companion to these equations
- Long Short-Term Memory — Hochreiter & Schmidhuber (1997), the original paper and its constant error carousel
- Learning to Forget: Continual Prediction with LSTM — Gers, Schmidhuber & Cummins (2000), where the forget gate was born