LSTM and GRU

The plain recurrent cell of the previous note overwrites its whole state at every step through a multiplication by $W_h$, and the gradient across $k$ steps is a product of $k$ such factors. The fix proposed by Hochreiter and Schmidhuber (1997) and simplified by Cho et al. (2014) is architectural: give the cell a path along which the state is added to rather than transformed, and let the network learn, per coordinate and per step, how much to keep, write and read.

Notation. As before, $x_t \in \mathbb R^d$ is the input, $h_t \in \mathbb R^H$ the hidden state. The LSTM adds a cell state $c_t \in \mathbb R^H$. $\sigma$ is the logistic function with range $(0, 1)$; $\tanh$ has range $(-1, 1)$; $\odot$ is element-wise multiplication. Each gate has its own weights $W_{\cdot}$ (on $x_t$), $U_{\cdot}$ (on $h_{t-1}$) and bias $b_{\cdot}$.

Baseline: one state, overwritten every step#

$h_t = \tanh(W_x x_t + W_h h_{t-1} + b)$. The state is recomputed from scratch at each step; nothing in the cell can say “leave coordinate 7 alone”. Whether old information survives depends on $W_h$, which is shared across all steps and all coordinates, and the gradient factor per step is $\mathrm{diag}(1 - h_t^2) W_h$.

LSTM#

Four quantities are computed from $(x_t, h_{t-1})$, three of them gates in $(0, 1)$ and one a candidate value in $(-1, 1)$:

$$\begin{aligned} f_t &= \sigma(W_f x_t + U_f h_{t-1} + b_f) && \text{forget gate: how much of } c_{t-1} \text{ to keep},\\ i_t &= \sigma(W_i x_t + U_i h_{t-1} + b_i) && \text{input gate: how much of the candidate to write},\\ \tilde c_t &= \tanh(W_c x_t + U_c h_{t-1} + b_c) && \text{candidate: what to write},\\ o_t &= \sigma(W_o x_t + U_o h_{t-1} + b_o) && \text{output gate: how much of the memory to expose}. \end{aligned}$$

Then the memory is updated and the working state read from it:

$$c_t = f_t \odot c_{t-1} + i_t \odot \tilde c_t, \qquad h_t = o_t \odot \tanh(c_t).$$

Every operation is per coordinate. A coordinate of $c$ with $f \approx 1$ and $i \approx 0$ is copied unchanged for as many steps as the gates stay that way; one with $f \approx 0$ and $i \approx 1$ is overwritten. The sigmoid range makes the gates fractions and the $\tanh$ range makes the candidate and the exposed value bounded, so $c_t$ can grow only by repeated addition of bounded values, not geometrically.

Worked example#

One coordinate, one step. Previous memory $c_{t-1} = 2$, previous state $h_{t-1} = 1$, input $x_t = 1$. Take scalar weights (toy values) so that each pre-activation is a small sum:

Gate	Pre-activation	Value
forget $f$	$1.63 \cdot 1 + 2.70 \cdot 1 + 1.62 = 5.95$	$\sigma(5.95) = 0.997$
input $i$	$1.65 \cdot 1 + 2.00 \cdot 1 + 0.62 = 4.27$	$\sigma(4.27) = 0.986$
candidate $\tilde c$	$0.94 \cdot 1 + 1.41 \cdot 1 - 0.32 = 2.03$	$\tanh(2.03) = 0.966$
output $o$	$-0.19 \cdot 1 + 4.38 \cdot 1 + 0.59 = 4.78$	$\sigma(4.78) = 0.992$

Memory update: $c_t = 0.997 \times 2 + 0.986 \times 0.966 = 1.995 + 0.953 = 2.948$. Exposed state: $h_t = 0.992 \times \tanh(2.948) = 0.992 \times 0.9945 = 0.986$.

The old memory survived almost entirely (forget gate near 1) and a new value of about $0.95$ was added. Had the forget gate come out at $0.1$, the memory would have dropped to $0.2 + 0.95$: same input, different decision, chosen by weights that see $x_t$ and $h_{t-1}$.

Step by step#

Figure

One LSTM step: what goes where

Following the numbers in the figure:

1. Forget gate. $f = \sigma(1.63 \cdot 1 + 2.70 \cdot 1 + 1.62) = 0.997$: how much of the old memory to keep, coordinate by coordinate.
2. Input gate and candidate. $i = 0.986$ says how much to write, $\tilde c = 0.966$ says what to write. The candidate is the plain RNN update of the previous note, with its own weights.
3. Keep. The old memory is multiplied by the forget gate: $0.997 \times 2 = 1.995$.
4. Write. The gated candidate $0.986 \times 0.966 = 0.953$ is added: $c_t = 2.948$. Steps 3 and 4 are the only operations on the memory line, a multiplication by a number in $(0, 1)$ and an addition.
5. Output gate. $o = 0.992$: how much of the memory to expose.
6. Read. $h_t = o \cdot \tanh(c_t) = 0.992 \times 0.9945 = 0.986$. Both $c_t$ and $h_t$ go on to step $t + 1$: the gates of the next step read $h_t$, and the memory line carries $c_t$.

Against the RNN figure of the previous note, the inputs and the working state $h_t$ are drawn the same way; what is new is the memory line at the top, in the amber outline, and the three gates that act on it.

Why the gradient survives#

Differentiate the memory update with respect to the previous memory:

$$\frac{\partial c_t}{\partial c_{t-1}} = \mathrm{diag}(f_t) + (\text{terms through } f_t, i_t, \tilde c_t \text{ that depend on } h_{t-1}).$$

The first term is the forget gate itself. Along the memory line the gradient across $k$ steps contains the product $\prod f_s$, and the network can keep that near 1 by keeping the forget gates near 1 wherever memory is needed, whereas the plain RNN’s factor $\mathrm{diag}(1 - h^2) W_h$ is a fixed matrix power. This does not prove that LSTMs learn arbitrary long dependencies (the other terms still vanish or explode, and gates saturate), but it removes the geometric decay that made the plain cell fail at tens of steps; in practice LSTMs learn dependencies over hundreds of tokens. The bias $b_f$ is often initialized to 1 so that the forget gates start open (Jozefowicz, Zaremba and Sutskever, 2015).

Cost#

Each of the four blocks has an $H \times d$ matrix, an $H \times H$ matrix and a bias, so an LSTM cell has $4(Hd + H^2 + H)$ parameters, four times the plain cell (PyTorch stores two bias vectors per block, giving $4(Hd + H^2 + 2H)$). The translation model of the seq2seq note uses $d = H = 128$: $4(128 \cdot 128 + 128 \cdot 128 + 2 \cdot 128) = 132{,}096$ parameters per direction. The computation per step is also four matrix products instead of one, and still strictly sequential.

GRU#

The gated recurrent unit (Cho et al., 2014) merges the memory and the working state and uses two gates:

$$\begin{aligned} z_t &= \sigma(W_z x_t + U_z h_{t-1} + b_z) && \text{update gate},\\ r_t &= \sigma(W_r x_t + U_r h_{t-1} + b_r) && \text{reset gate},\\ \tilde h_t &= \tanh\big(W_h x_t + U_h (r_t \odot h_{t-1}) + b_h\big) && \text{candidate},\\ h_t &= (1 - z_t) \odot h_{t-1} + z_t \odot \tilde h_t . \end{aligned}$$

The update gate plays the roles of both $f$ and $i$, tied as $f = 1 - z$: what is not kept is overwritten. The reset gate controls how much of the previous state the candidate may read. There is no output gate; $h_t$ is the memory. The parameter count is $3(Hd + H^2 + H)$, three quarters of an LSTM at the same size, and one fewer matrix product per step. The AG News classifier of the previous note is a GRU with $d = 100$ and $H = 128$.

Conventions. The equations above write $z_t$ as the fraction of the state that is overwritten. The original paper (Cho et al., 2014, eqs. 5–8) and PyTorch’s nn.GRU interpolate the other way round, $h_t = z_t \odot h_{t-1} + (1 - z_t) \odot \tilde h_t$, so $z_t$ there is $1 - z_t$ here. PyTorch also applies the reset after the recurrent linear map, $\tilde h_t = \tanh\big(W_h x_t + b_h + r_t \odot (U_h h_{t-1} + b'_h)\big)$, while the equations above, like Cho et al., reset $h_{t-1}$ before multiplying by $U_h$; the two candidates differ in general. With its second bias vector per gate, PyTorch’s count is $3(Hd + H^2 + 2H)$. Both are standard GRUs, but a formula and an implementation must use the same one.

	Plain RNN	LSTM	GRU
States	$h$	$h, c$	$h$
Gates	0	3 ($f, i, o$)	2 ($z, r$)
Memory update	overwrite	$f \odot c + i \odot \tilde c$	$(1 - z) \odot h + z \odot \tilde h$
Parameters	$Hd + H^2 + H$	$4\times$	$3\times$
Gradient factor along memory	$\mathrm{diag}(1-h^2)W_h$	$\approx \mathrm{diag}(f)$	$\approx \mathrm{diag}(1 - z)$

Which cell is better is empirical and task-dependent. Direct comparisons find the two close (Chung et al., 2014; Jozefowicz, Zaremba and Sutskever, 2015), and a large study of LSTM variants found that coupling the input and forget gates, as the GRU does, costs little (Greff et al., 2017).

Reading in both directions#

A recurrent state at position $t$ summarizes $x_{\le t}$ and knows nothing of $x_{> t}$. For a classifier or an encoder that sees the whole input at once, run a second cell from right to left and concatenate the two states at each position: the bidirectional RNN gives $[\overrightarrow{h}_t; \overleftarrow{h}_t] \in \mathbb R^{2H}$, a representation of each token informed by its full context, at twice the cost. The encoder in the next note is a bidirectional LSTM. A language model cannot be bidirectional, since it must not see the token it predicts; the BERT note finds a different way to use both directions.

What changed#

	Plain RNN	Gated RNN
Can keep a coordinate unchanged	Only if $W_h$ happens to	Yes, by gating
Long dependencies	Tens of tokens	Hundreds, in practice
Parameters per cell	$1\times$	$3\times$ (GRU), $4\times$ (LSTM)
Sequential	Yes	Yes
Fixed-size state	Yes	Yes

Limitations#

• Still sequential. Step $t$ depends on step $t - 1$. Training and inference time grow with $T$, and the sequence dimension cannot be parallelized.
• Still one vector. However well it remembers, the state at the end is $H$ (or $2H$) numbers that must hold everything the next stage needs. The seq2seq note makes this failure visible in translation.
• Gates saturate. Sigmoids near 0 or 1 have small derivatives; a gate that has learned to close is slow to reopen.
• Capacity costs. Four matrices per step; a 2-layer bidirectional LSTM with $H = 512$ over 512-dimensional inputs has 10.5 M parameters before any embedding.

What you should remember#

• LSTM: $c_t = f_t \odot c_{t-1} + i_t \odot \tilde c_t$, $h_t = o_t \odot \tanh(c_t)$; gates in $(0, 1)$ from $\sigma$, candidate in $(-1, 1)$ from $\tanh$.
• The additive memory update makes the gradient factor along $c$ approximately the forget gate, which the network can hold near 1.
• GRU: two gates, one state, $h_t = (1 - z_t)\odot h_{t-1} + z_t \odot \tilde h_t$; three quarters of the LSTM’s parameters.
• Bidirectional cells give each position its full context; they cannot be used for next-token prediction.
• Gating fixes the gradient, not the sequential computation or the single-vector summary.

Notebook#

week04_lstm_attention_transformer.ipynb. Cell numbers count every cell, markdown included, from 1. Quoted numbers come from the run saved with the course; cells changed since then for current library versions have no saved output.

• Cell 39 defines Seq2SeqLSTM with nn.LSTM(128, 128, bidirectional=True) as the encoder and a unidirectional nn.LSTM(128, 128) as the decoder; the encoder’s two final hidden and cell states (h_enc, c_enc of shape [2, B, 128]) are concatenated to [1, B, 256] and projected back to 128 by h_proj and c_proj to initialize the decoder.
• Cell 44 prints 8,380,474 parameters in total and 492,602 outside the embedding and output matrices; the latter is $2 \times 132{,}096$ (bidirectional encoder) $+ 132{,}096$ (decoder) $+ 65{,}792$ (the two projections) $+ 30{,}522$ (the output bias, which the notebook’s subtraction leaves in).

week02_NN_text_classification.ipynb, cell 40, is the GRU classifier referred to above. It uses nn.GRU, so its gates follow the PyTorch convention of the GRU section: its update gate is the fraction kept, and the reset acts after the recurrent matrix.

Next#

With a cell that can remember, the course takes on its first generation task with two sequences: translate a German sentence into English. One recurrent network reads the source, another writes the target. The next note builds that encoder–decoder, measures it with BLEU, and discovers that the single vector handed from encoder to decoder is now the limiting factor.

References#

• S. Hochreiter, J. Schmidhuber. Long Short-Term Memory. Neural Computation, 9(8), 1997.
• K. Cho, B. van Merriënboer, C. Gulcehre, D. Bahdanau, F. Bougares, H. Schwenk, Y. Bengio. Learning Phrase Representations using RNN Encoder–Decoder for Statistical Machine Translation. EMNLP, 2014. The GRU.
• J. Chung, C. Gulcehre, K. Cho, Y. Bengio. Empirical Evaluation of Gated Recurrent Neural Networks on Sequence Modeling. NIPS Deep Learning Workshop, 2014. LSTM and GRU compared on music and speech modelling.
• R. Jozefowicz, W. Zaremba, I. Sutskever. An Empirical Exploration of Recurrent Network Architectures. ICML, 2015. Forget-gate bias initialization; LSTM and GRU compared.
• K. Greff, R. K. Srivastava, J. Koutník, B. R. Steunebrink, J. Schmidhuber. LSTM: A Search Space Odyssey. IEEE TNNLS, 28(10), 2017.
• C. Olah. Understanding LSTM Networks. 2015. The gate-by-gate walkthrough that the worked example follows.
• M. Schuster, K. K. Paliwal. Bidirectional recurrent neural networks. IEEE Transactions on Signal Processing, 45(11), 1997.