LSTM and GRU

The big idea

A plain RNN ( Recurrent neural networks ) rewrites its whole memory at every step: the new state is a squashed mix of the old state and the new input. After 50 steps, whatever it saw at step 1 has been squashed 50 times, and so has the gradient that should teach it. The fix is to stop rewriting and start editing.

Picture a factory cell stateThe LSTM's separate memory vector, updated only by elementwise scaling (forget) and addition (input), so information and gradients can pass many steps unchanged. as a conveyor belt running past a row of workers. One worker can take items off the belt, one can put new items on, and one decides what to show the customer at the window. Items nobody touches ride along untouched for as long as the belt runs. A long short-term memoryLong short-term memory: a recurrent cell with a separate cell state and three learned gates (forget, input, output) that control what is erased, written and exposed at each step. (LSTM) cell works the same way, and the workers are learned gatesA vector of values between 0 and 1, produced by a sigmoid, that is multiplied elementwise with another vector to let a fraction of each entry through..

How it works

Every gate is a small layer that looks at the previous output $h_{t-1}$ and the new input $x_t$ and outputs a number between 0 and 1 for each entry of the memory, through a sigmoid $\sigma$. Multiplying by a gate value of 0 blocks an entry; 1 lets it through unchanged; 0.3 lets 30 percent through.

Step 1: The forget gate

The forget gateThe LSTM gate that scales the previous cell state entry by entry, deciding what to keep and what to erase. decides what to erase. Reading a review, a unit tracking "is the sentence negated?" should keep its value through "not very" and reset after the full stop. Interestingly, the original 1997 LSTM had no forget gate. Gers, Schmidhuber and Cummins added it in 2000 because, on continuous input streams with no marked ends, the cell state could grow without limit and break the network; the forget gate lets a cell learn to reset itself.

Step 2: The input gate and the candidate

$g_t$ is the candidate content, like the new state of a plain RNN. The input gateThe LSTM gate that scales the candidate content before it is added to the cell state, deciding what new information to write. $i_t$ decides how much of it to write. A word like "however" might open the input gate wide; filler words might barely move it.

Step 3: The conveyor belt update

($\odot$ means entry-by-entry multiplication.) This one line is the reason LSTMs work. The memory is never passed through a weight matrix and a squashing function on its way from step to step, only scaled and added to. Look at the gradient along the belt: $\partial c_t/\partial c_{t-1}$ along this path is just $f_t$. If the forget gate stays near 1, the gradient passes back through 100 steps almost unchanged, instead of being multiplied by a matrix and a tanh slope 100 times. Hochreiter and Schmidhuber's original paper calls this enforcing "constant error flow through constant error carousels" and reports bridging time lags of more than 1,000 steps.

Step 4: The output gate

The cell can hold information that is not useful right now. The output gateThe LSTM gate that scales tanh of the cell state to produce the hidden output h(t), deciding what part of the memory to expose. decides what part of the memory to expose as $h_t$, the vector the next step's gates read and the next layer receives. These are exactly the equations PyTorch's nn.LSTM documents, where the candidate $g_t$ is called the "cell gate".

The GRU: two gates instead of three

The gated recurrent unitGated recurrent unit: a simpler gated cell with an update gate and a reset gate and no separate cell state. (GRU), proposed by Cho et al. in 2014 and described by them as "motivated by the LSTM unit but much simpler to compute and implement", merges the cell state and the output into one vector $h$ and uses two gates:

(biases omitted). The update gateThe GRU gate that mixes the old state and the candidate: near 1 keeps the old state, near 0 replaces it. $z_t$ does the jobs of forget and input together: a value near 1 keeps the old state, near 0 replaces it with the candidate. The same "keep a fraction, add the rest" form gives it the LSTM's gradient highway. The reset gateThe GRU gate that scales how much of the previous state is used when computing the candidate; near 0 lets the unit start fresh from the current input. $r_t$ lets the candidate ignore the past; in Cho et al.'s words, near 0 it "allows the hidden state to drop any information that is found to be irrelevant later in the future". The sign convention for $z$ differs between papers and libraries; this is PyTorch's.

Going deeper

Code: does the start still matter after 100 steps?

import numpy as np
rng = np.random.default_rng(1)
D, H, T = 3, 8, 100
sig = lambda z: 1 / (1 + np.exp(-z))
xs = rng.normal(0, 1, (T, D))
Wx, Wh = rng.normal(0, 0.3, (4 * H, D)), rng.normal(0, 0.3, (4 * H, H))   # rows: f, i, g, o
U, W = rng.normal(0, 0.3, (H, H)), rng.normal(0, 0.3, (H, D))

# Run each cell for 100 steps while tracking J = d(state_t)/d(state_0), the memory of the start
def lstm(forget_bias):
    c, h, Jc, Jh, fs = np.zeros(H), np.zeros(H), np.eye(H), np.zeros((H, H)), []
    for x in xs:
        z, dz = Wx @ x + Wh @ h, Wh @ Jh
        f, i, g, o = sig(z[:H] + forget_bias), sig(z[H:2*H]), np.tanh(z[2*H:3*H]), sig(z[3*H:])
        df, di = (f * (1 - f))[:, None] * dz[:H], (i * (1 - i))[:, None] * dz[H:2*H]
        dg, do = (1 - g**2)[:, None] * dz[2*H:3*H], (o * (1 - o))[:, None] * dz[3*H:]
        Jc = df * c[:, None] + f[:, None] * Jc + di * g[:, None] + i[:, None] * dg
        c = f * c + i * g                                 # the conveyor belt: scale, then add
        h = o * np.tanh(c)
        Jh = do * np.tanh(c)[:, None] + (o * (1 - np.tanh(c)**2))[:, None] * Jc
        fs.append(f.mean())
    return np.linalg.norm(Jc), np.mean(fs)

def rnn():
    h, J = np.zeros(H), np.eye(H)
    for x in xs:
        h = np.tanh(W @ x + U @ h)
        J = (1 - h**2)[:, None] * (U @ J)
    return np.linalg.norm(J)

print(f"vanilla RNN          |dh_100/dh_0| = {rnn():.1e}")
for fb in [0.0, 1.0, 3.0, 6.0]:
    s, f = lstm(fb)
    print(f"LSTM forget bias {fb:.0f}:  |dc_100/dc_0| = {s:.1e}  (average forget gate {f:.2f})")

# Parameters per layer for input size 300 and hidden size 128 (one bias vector per block)
d, n = 300, 128
block = n * d + n * n + n
print(f"parameters  RNN {block:,}  GRU {3 * block:,}  LSTM {4 * block:,}")

vanilla RNN          |dh_100/dh_0| = 7.0e-11
LSTM forget bias 0:  |dc_100/dc_0| = 1.4e-18  (average forget gate 0.50)
LSTM forget bias 1:  |dc_100/dc_0| = 3.9e-06  (average forget gate 0.73)
LSTM forget bias 3:  |dc_100/dc_0| = 1.8e+01  (average forget gate 0.95)
LSTM forget bias 6:  |dc_100/dc_0| = 1.6e+01  (average forget gate 1.00)
parameters  RNN 54,912  GRU 164,736  LSTM 219,648

The code runs a vanilla RNN and an LSTM over the same 100 random inputs and tracks, exactly, how much the final state depends on the starting state. For the plain RNN the dependence is about $10^{-10}$: the start is forgotten, and no gradient can reach it.

For the LSTM it depends on the forget gate. With a forget-gate bias of 0 the gates average about 0.5, so the belt halves its contents every step and the dependence falls to about $10^{-18}$, worse than the RNN. With bias 3 the gates average 0.95 and the dependence stays around 18: the memory and the gradient survive 100 steps. Gates are not magic; they give the network a path that keeps information when the forget gate learns to stay open. That is why a positive forget-gate bias at initialization is common advice. Jozefowicz, Zaremba and Sutskever (2015) found that adding a bias of 1 to the LSTM's forget gate closed the gap between the LSTM and the GRU in their experiments.

The last line counts parameters. Each gate is a full layer, so an LSTM layer has 4 times the parameters of a plain RNN layer of the same size and a GRU has 3 times: 219,648 and 164,736 against 54,912 for input size 300 and hidden size 128. (PyTorch keeps two bias vectors per gate, so its counts are slightly higher.)

LSTM or GRU?

Chung et al. (2014) compared both on polyphonic music and speech signal modelling and found gated units clearly better than plain tanh units, and the GRU "comparable to LSTM". Jozefowicz et al. evaluated over ten thousand architectures and found no single design that won on every task. The practical rule: the GRU is smaller and a bit faster, the LSTM has a separate memory and output and is the more common default; try both if it matters, and tune the forget-gate bias for the LSTM.

What gates do not fix

• Still sequential. Step 50 waits for step 49. Gates change what is remembered, not how fast it is computed.
• Still a fixed-size memory. Everything the network knows about a 500-word paragraph must fit in one vector. For translation, that bottleneck is what attention removes ( Seq2seq and attention ).
• Exploding gradients can still happen through the gate paths, so gradient clipping ( Recurrent neural networks ) stays in the training loop.

Common extensions: bidirectionalTwo RNNs read the sequence in opposite directions and their states are concatenated, so each position sees both left and right context. layers run one LSTM left to right and another right to left and concatenate their states, so each position knows its full context (useful for tagging, not for predicting the next word). Stacking several LSTM layers, each reading the previous layer's outputs, adds depth.

Trade-offs

Common mistakes

• Saying the LSTM "solves" vanishing gradients unconditionally. It gives a path that can carry them; a forget gate stuck near 0.5 still forgets fast.
• Confusing the cell state $c_t$ with the hidden output $h_t$. Only $h_t$ is passed to the next layer and the output head.
• Using a bidirectional LSTM for next-word prediction or real-time streaming, where the future is not available.
• Comparing LSTM and GRU at the same hidden size and calling the GRU worse, when it has a quarter fewer parameters.
• Dropping gradient clipping because "LSTMs fixed the gradient problem".

In the wild

From the mid-2010s until transformers, LSTMs were the default model for sequences. Speech recognition and machine translation systems were built on stacked LSTM layers; Google's 2016 neural machine translation paper, for example, describes a deep LSTM network with 8 encoder and 8 decoder layers using attention and residual connections. Deep learning libraries ship LSTM and GRU layers as standard building blocks, and they are still a good choice for streaming data on small devices, where processing one step at a time with a fixed-size state is exactly what you want.

Check yourself

Recap