Recurrent Networks Lab

Step a recurrent network through a sequence and watch its memory fade or hold.

Intermediate interactive lab, about 15 minutes. Techniques: RNN, LSTM, Vanishing gradients.

About

A recurrent neural network reads a sequence one element at a time and carries a hidden state from each step to the next. That state is the network's only memory: a vector of numbers that has to summarise everything it has read so far that might matter later.

This lab trains three recurrent networks on the same text, in your browser, while you watch: a plain (vanilla) RNN, a GRU and an LSTM. Each is a character-level language model. Given the characters so far, it outputs a probability for every possible next character. Nothing is pre-computed or faked. The weights start random, the forward pass, backpropagation through time and the Adam updates are all written out in about 400 lines of TypeScript, and the pictures are read straight from the live networks.

The interesting part is what the three disagree on. On the nursery rhymes they all learn spelling and common words. On the memory test, where a letter must be recalled after eight dashes, the vanilla RNN does no better than guessing, the GRU gets it right almost every time, and the LSTM is usually partway there when training stops. Gates make long memory learnable; they do not all learn it equally fast. The gradient chart shows why the vanilla RNN never gets there.

What is running

Hidden units
32 per model
Training windows
24 characters, random positions
Training steps
3,000 per model, batch size 1
Optimiser
Adam, learning rate 0.01, gradient norm clipped at 5
Initial state
All zeros, for training and for your text
Vanilla RNN parameters
2,876
GRU parameters
6,780
LSTM parameters
8,732

Things to try

  1. Watch it learnPress R to retrain and keep Generate pressed every few seconds. Early samples are random letters, then word shapes, then real words.
  2. Memory testSwitch to the memory test (3). When training finishes, step to the letter after the dashes and compare the three prediction panels.
  3. Break the gradientIn the gradient chart, drag σ down to 0.5, then up to 2.5. Then set input strength to 0 and try again.
  4. TemperatureGenerate at 0.2 and at 1.3. Low temperature repeats the likeliest phrase; high temperature invents words.

How it works

Notation: xt is the character read at step t (one-hot), ht the hidden state, σ the logistic sigmoid, ⊙ element-wise multiplication. Every model ends in the same output layer: pt = softmax(Wyht + by), a probability for each next character.

Vanilla RNN (Elman, 1990)

ht = tanh(Wxxt + Whht−1 + b)

One matrix mixes the old state with the new input, then tanh squashes the result into (−1, 1). The same weights are reused at every step.

LSTM (Hochreiter and Schmidhuber, 1997; forget gate from Gers et al., 2000)

i, f, o = σ(W·[xt, ht−1] + b)g = tanh(Wg·[xt, ht−1] + bg)ct = f ⊙ ct−1 + i ⊙ ght = o ⊙ tanh(ct)

A separate cell state c runs alongside h. The forget gate f decides how much of it survives each step, the input gate i how much new content g is written, the output gate o how much is exposed as h.

GRU (Cho et al., 2014)

r, z = σ(W·[xt, ht−1] + b)n = tanh(Wnxt + Un(r ⊙ ht−1) + bn)ht = (1 − z) ⊙ n + z ⊙ ht−1

A lighter design with two gates and no separate cell. The update gate z interpolates between keeping the old state and taking the new candidate n.

Training: backpropagation through time

To train, the network is unrolled over a window of 24 characters, so it becomes a 24-layer feed-forward network in which every layer shares the same weights. The loss is the average cross-entropy of the true next characters. Backpropagation runs from the last step to the first, and each shared weight collects gradient from every step. That is backpropagation through time (BPTT).

The additive update ct = f ⊙ ct−1 + … is the key design choice. Its derivative with respect to ct−1 is just f, with no weight matrix and no squashing. When the forget gate sits near 1, gradient passes back through many steps almost unchanged. Hochreiter and Schmidhuber called this the constant error carousel. This lab initialises the forget-gate bias to 1 so the carousel starts switched on, following Jozefowicz et al. (2015).

Want the same ideas in a longer read? The Natural Language Processing lesson covers sequence models, and Transformers and Attention explains what replaced them.

Vanishing gradients

When the loss at step T is sent back to step T − k, it is multiplied by the Jacobian of every step in between:

∂hT/∂hT−k = JT JT−1 ⋯ JT−k+1

Each factor is diag(1 − h²) W for a tanh RNN. Multiply k of them together and the size of the product behaves roughly like (typical gain)ᵏ.

If the typical gain per step is 0.8, after 20 steps the gradient is 0.8²⁰ ≈ 0.012 of its size. After 50 steps it is 0.8⁵⁰ ≈ 0.000014. The network still receives a gradient, but the part that says “this character 50 steps ago mattered” is drowned out by nearby steps. That is the vanishing gradient problem, analysed by Hochreiter (1991) and by Bengio, Simard and Frasconi (1994). If the gain is above 1, the product explodes instead, and a single update can wreck the weights. Pascanu, Mikolov and Bengio (2013) proposed gradient clipping for that case, which this lab uses.

The random-RNN chart computes the product exactly: it builds a 32 × 32 recurrent matrix with entries drawn from N(0, σ²/32), runs 60 steps with random inputs, and multiplies the Jacobians together. The number shown is the Frobenius norm divided by √32, so the identity matrix scores exactly 1.

What the chart shows

  • σ below 1Straight lines on a log axis: exponential decay. The network is effectively blind to anything more than a few steps back.
  • σ near 1Still decays, because tanh′ is below 1 wherever units are active. Saturated units (near ±1) pass almost no gradient at all.
  • σ well above 1The product grows. With strong inputs the units saturate and tame it; set input strength to 0 to see how violent it can be.
  • Trained modelsMeasured by nudging each model's state k steps back along random directions and re-running to the end. The LSTM's cell path keeps the line high.

Reading the heat strip

Each strip is a picture of one network's memory. A column is the hidden state right after reading the character underneath it. Each of the 32 rows is one hidden unit, coloured with the viridis scale: dark purple at −1, teal around 0, yellow at +1.

Look for rows that change colour at the same kind of character, such as every space or every vowel. Those units have learned a feature. Karpathy, Johnson and Fei-Fei (2015) found individual LSTM cells that track quote marks, line length and indentation in exactly this way. In a 32-unit model trained for a few seconds the features are cruder, but word boundaries are often visible.

Compare how quickly each strip changes. The vanilla RNN has to rewrite its whole state at every step, so its columns tend to look unrelated to each other. The GRU and LSTM can hold units steady across many steps, which shows up as horizontal bands. On the memory test, look for units that switch on at the first letter and stay put until the dashes end.

The gate readout

Under the LSTM and GRU predictions, the lab prints the average gate values at the selected step. An LSTM forget gate averaging 0.9 means most of the cell state is being carried forward. A GRU update gate near 1 means the same: keep the old state, ignore the new candidate. Step through a word and watch these numbers dip at spaces, where the model can afford to forget.

RNNs today

LSTMs were the workhorse of sequence modelling for most of the 2010s: speech recognition, handwriting recognition and, from 2014, neural machine translation with encoder-decoder models. In 2017 the transformer replaced recurrence with attention, which looks at every position at once. That removed the long gradient path and, just as importantly, let training run in parallel across the sequence instead of step by step.

Recurrence has come back in a new form. A transformer's cost per generated token grows with the length of the context, while a recurrent model carries a fixed-size state. State space models such as Mamba (Gu and Dao, 2023), RWKV (Peng et al., 2023) and xLSTM (Beck et al., 2024) design their recurrences so that training can still be parallelised, and several production models now mix these layers with attention. Feng et al. (2024) showed that stripping LSTMs and GRUs down to “minimal” versions makes them trainable in parallel too.

The trade-off you can see in this lab is the one these designs still negotiate: a fixed-size memory is cheap, but it has to decide what to forget.

Sources and further reading

  1. 01Long short-term memory. Hochreiter and Schmidhuber, Neural Computation, 1997
  2. 02Learning to forget: continual prediction with LSTM. Gers, Schmidhuber and Cummins, Neural Computation, 2000
  3. 03Learning phrase representations using RNN encoder-decoder. Cho et al., 2014 (introduces the GRU)
  4. 04Learning long-term dependencies with gradient descent is difficult. Bengio, Simard and Frasconi, IEEE Trans. Neural Networks, 1994
  5. 05On the difficulty of training recurrent neural networks. Pascanu, Mikolov and Bengio, ICML 2013
  6. 06An empirical exploration of recurrent network architectures. Jozefowicz, Zaremba and Sutskever, ICML 2015
  7. 07Visualizing and understanding recurrent networks. Karpathy, Johnson and Fei-Fei, 2015
  8. 08Mamba: linear-time sequence modeling with selective state spaces. Gu and Dao, 2023
  9. 09RWKV: reinventing RNNs for the transformer era. Peng et al., 2023
  10. 10xLSTM: extended long short-term memory. Beck et al., 2024
  11. 11Were RNNs all we needed?. Feng et al., 2024

Related