Recurrent Neural Networks

Unroll an RNN through time and watch the hidden state and vanishing gradients

IntermediateHidden StateUnrollingVanishing GradientsLSTM

Time step t = 1 of 5

Unrolled through time

The same cell (same weights) is applied at every step. Each hidden state is 4 numbers: orange positive, blue negative.

h₀ = 0tanhh₁hx₁tanhh₂ex₂tanhh₃lx₃tanhh₄lx₄tanhh₅ox₅

Vanilla RNN step

h(t) = tanh(W_h · h(t−1) + W_x · x(t) + b)

h(t−1) 0.00 0.00 0.00 0.00
W_h · h(t−1) 0.00 0.00 0.00 0.00
W_x · x(t) ['h']-0.36 0.95 0.20 0.29
b 0.10-0.10 0.00 0.05
sum a(t)-0.26 0.85 0.20 0.34
h(t) = tanh(a)-0.25 0.69 0.19 0.33

x(t) for 'h' = [-0.12, 0.99, 0.50]. Weights are hand-picked for illustration, not trained.

LSTM cell at the same step

Same input sequence, separate seeded weights. Gates decide what to forget, write and output.

c(t−1)×+c(t)forget σinput σcandidate tanhoutput σ×tanh×h(t)[h(t−1), x(t)] feeds every gategate colours: white = 0 (closed), green = 1 (open)
f forget 0.66 0.72 0.71 0.84
i input 0.40 0.45 0.66 0.39
g candidate-0.06-0.29 0.07-0.21
o output 0.56 0.54 0.37 0.34
c(t)-0.02-0.13 0.04-0.08
h(t)-0.01-0.07 0.02-0.03

Vanishing and exploding gradients

How much a loss at step T still depends on a hidden state k steps earlier. Each step back multiplies the gradient by the same kind of factor (log scale).

1e-121e-81e-411e41e81e1208152330steps back in time k
Linear RNN |w|k

at T: 4.2e-2

tanh RNN ∏ w·(1 − h²)

at T: 1.6e-3

LSTM cell path fk

at T: 4.0e-1

A one-unit model keeps the maths visible: with a matrix W_h the role of |w| is played by its largest singular value. The tanh curve uses a small fixed input sequence; because tanh′ ≤ 1 it can only shrink the gradient further, so it vanishes even for some w above 1.

How a recurrent network reads a sequence

A recurrent neural network processes a sequence one element at a time. At every step it combines the new input x(t) with its previous hidden state h(t−1) and squashes the result with tanh. Drawing each step side by side, called unrolling, shows that an RNN is really a very deep network whose layers all share the same weights.

That depth is also the problem. Training multiplies one factor per step on the way back, so the gradient from a distant step either fades towards zero or grows without bound. LSTMs fix most of this with a separate cell state that is updated by addition and controlled by gates.

Everything here runs in your browser on a toy model: each letter is mapped to three fixed numbers, the RNN weights are hand-picked and the LSTM weights come from a seeded random generator. The values are illustrative, but every number is computed exactly as the formulas below describe.

Key formulas

Vanilla RNN step

h(t) = tanh(Wh h(t−1) + Wx x(t) + b)

h(t)
hidden state after step t (h(0) = 0)
x(t)
input at step t
W_h, W_x, b
recurrent weights, input weights and bias, shared across steps

Gradient through k steps

∂h(T) / ∂h(T−k) = ∏t=T−k+1…T diag(1 − h(t)²) · Wh

For a single unit this is ∏ w·(1 − h²), which behaves like |w|k: it vanishes when |w| < 1 and explodes when |w| > 1.

LSTM cell

f = σ(Wf[h(t−1), x(t)] + bf), and likewise i, o; g = tanh(Wg[h(t−1), x(t)] + bg)

c(t) = f ⊙ c(t−1) + i ⊙ g, h(t) = o ⊙ tanh(c(t))

Along the cell state the gradient is multiplied by f at each step, which the network can keep close to 1.

How to use this visualizer

  1. 1
    Type a short word. Each letter becomes one time step.
  2. 2
    Press Step to process one letter at a time, or Play to run through the whole word. The highlighted cell is the current step.
  3. 3
    Read the vanilla RNN table to see W_h·h(t−1), W_x·x(t) and b added together before tanh produces the new hidden state.
  4. 4
    Compare with the LSTM cell for the same step: green gate cells near 1 are open, white cells near 0 are closed.
  5. 5
    In the gradient panel, move the recurrent weight w above and below 1 and lengthen the sequence. Then raise the forget gate towards 1 and see how the LSTM path keeps the gradient alive.

Key concepts

Hidden state

A vector that summarises everything the network has seen so far. It is the only thing passed from one step to the next.

Weight sharing

W_h, W_x and b are the same at every time step, so one small cell can read sequences of any length.

Backpropagation through time

Training unrolls the loop into a deep chain and backpropagates through it, multiplying one Jacobian per step.

Vanishing gradients

If those per-step factors are below 1 the product shrinks exponentially, and early inputs stop influencing learning.

Exploding gradients

If they are above 1 the product blows up. Gradient clipping caps the norm so a single update cannot wreck the weights.

LSTM gates

Sigmoid gates between 0 and 1 control what to forget, what to write and what to output, keeping an additive path for the cell state.

Code

NumPy
import numpy as np

def rnn_forward(xs, W_h, W_x, b):
    h = np.zeros(W_h.shape[0])
    hs = []
    for x in xs:                                  # one step per input
        h = np.tanh(W_h @ h + W_x @ x + b)        # same weights every step
        hs.append(h)
    return hs

def lstm_step(x, h, c, W, b):
    z = np.concatenate([h, x])
    f = sigmoid(W["f"] @ z + b["f"])   # forget gate
    i = sigmoid(W["i"] @ z + b["i"])   # input gate
    g = np.tanh(W["g"] @ z + b["g"])   # candidate values
    o = sigmoid(W["o"] @ z + b["o"])   # output gate
    c = f * c + i * g                  # additive cell-state update
    h = o * np.tanh(c)
    return h, c

Practice problems

Problem 1: Two steps by hand

A one-unit RNN has h(t) = tanh(0.5·h(t−1) + 1·x(t)), with no bias and h(0) = 0. The input sequence is x = [1, 0]. Compute h(1) and h(2).
Show solution

h(1) = tanh(0.5·0 + 1·1) = tanh(1) ≈ 0.762

h(2) = tanh(0.5·0.762 + 1·0) = tanh(0.381) ≈ 0.363

h(1) ≈ 0.762, h(2) ≈ 0.363: with no new input the memory decays

Problem 2: How fast do gradients vanish?

Ignoring the tanh derivative, a gradient travelling back 50 steps is multiplied by w each step. What factor survives for w = 0.9 and w = 1.1? What about an LSTM cell path whose forget gate stays at 0.99?
Show solution

w = 0.9: 0.9^50 ≈ 0.0052 (the signal almost disappears)

w = 1.1: 1.1^50 ≈ 117 (the signal explodes)

LSTM: 0.99^50 ≈ 0.61 (most of the signal survives)

A small change around 1 makes a huge difference over long sequences

Discussion

Questions, ideas or a different approach? Join the conversation.