RNN & LSTM Primer

It reads left to right and keeps a single vector, mixing it at each step with the new token and squashing the result: h_t = tanh(W_h·h_{t−1} + W_x·x_t + b). We run it in one dimension, so every frame is checkable by hand.

The eight tokens are given, so they are grey from the start; the states already computed sit above them, and the state the loop is computing now is the one under your hand. Drag the slider to walk t forward:

Notice that nothing right of the slider exists. Those cells are dashed because h_4 is not a value the network is holding back, it is a value nobody has computed. That is the invariant, and it is worth writing as an assertion: at the top of every step, h_t is a function of x_0 … x_t and of nothing later.

The squash is not decoration. It folds any pre-activation into the open interval from −1 to 1, and its slope at the point you are standing on is what the next section turns out to be entirely about. Drag the point along the curve:

Watch the tangent. At z = 0 it is at its steepest and the readout gives tanh' = 1.000; by z = 2 the curve has flattened to 0.071. A saturated unit is one whose output has stopped answering its input — and one whose gradient has stopped travelling.

There is only one W_h. The same number multiplies every state on its way to the next, which is how one set of weights reads a sentence of any length. Move it and all eight move together — try W_h = 1.40:

Sharing the weight is what keeps the parameter count independent of the sequence: one d×d matrix for the recurrence, one for the input, one bias. At 1.40 the last state reads 0.841; at W_h = 0 it reads 0.310 — exactly tanh(0.8 × 0.4), its own token and nothing else.

An answer comes out of those states through one more matrix. Classification reads a single output off the last state; tagging reads one per position. Switch between them:

Either way, the read-out head sees nothing but the state beneath it. Everything the first token contributes to the last answer has to survive seven passes through W_h and seven squashes to get there. What that trip does to it is the next section.

To learn from the last token, the loss has to reach the first. It travels the same chain in reverse, and at each hop it is multiplied by ∂h_t/∂h_{t−1} = tanh'(z_t) · W_h. The grey curve is that slope on its own, and the bar is the whole hop — the slope times the one weight:

The slope is 1 only at z = 0 and falls away either side, so the bar is under W_h everywhere, and under it by a lot wherever the unit is doing anything. That is one hop. Send the gradient back through eight:

Seven hops take it from 1 to 0.057. Nothing went wrong — every factor here is a legitimate derivative of a working network. The trouble is only that there are seven of them, and that sentences are longer than eight words.

So look at the factors alone, over a chain of forty. The dashed rule is × 1, the line between shrinking and growing. Drag W_h, then switch the activation:

Because tanh' never exceeds 1, the factor never exceeds W_h, and because these states never sit at zero it stays strictly below it: at the shipped 0.90 the largest is 0.884. Push W_h to 1.80 and the bars still cannot reach the rule — saturation is the brake, and it holds the peak to 0.616. Switch to ReLU and the brake is gone: its derivative is exactly 1 wherever it is on.

Multiply those factors together, over a chain of two hundred this time, and you get the number that reaches h_0. The axis is logarithmic — each gridline is twenty-four orders of magnitude above the one below — so a constant factor draws a straight line:

At the shipped settings the product crosses the lower rule — the smallest normal float32 — at hop 161, and ends at 9.8e-48. Switch to ReLU and set W_h = 1.60 and it climbs instead, crossing the upper rule at hop 189. A float32 past that is infinity, and a loss that touches infinity is NaN.

In one dimension the brake always wins; across many, the largest singular value of W_h can act on a direction nothing has saturated in, and the product does run away. The standard guard is Pascanu, Mikolov and Bengio's 2013 rescale, and this step's raw norm came out at 84.0 against a threshold of 5 — drag it back down through the threshold:

Because the clip is a ceiling and not a floor, it does nothing for the other half. Vanishing fails silently: no exception, no NaN, the loss falls as usual, and the model quietly learns to predict from the last few tokens — the only evidence whose gradient survived the trip.

Hochreiter and Schmidhuber's 1997 answer was a second piece of memory, c, whose update is c_t = f · c_{t−1} + i · ĉ_t — a multiply by one number and an add. No weight matrix stands between c_{t−1} and c_t, and no squash.

The picture is §01's with exactly one element changed. The token row is the same, the cells are the same, and only the connector differs: the cell state travels a plain line, while what this step writes arrives from below. Walk it forward:

Compare that connector with §01's arrow. There the hop was tanh'(z)·W_h, two things the network cannot easily control; here it is the forget gate, one number the network computes on purpose at every step. That swap is the entire idea.

Both gates are fractions, so one update is two lengths laid end to end: the share of the old state that survives, then the share of the candidate that is written, and the remainder that is thrown away. Move the forget gate, then the input gate:

Set it to f = 1.00 and the discarded remainder vanishes: the cell state becomes a pure accumulator. Set it to 0 and the old state is gone in one step — which is what a network wants when one sentence ends and the next is unrelated.

The reason this fixes §02 is that the only path from c_0 to c_k is a product of forget gates — f^k and nothing else. Against it, in rose, is the plain RNN's product from the last section, on the same linear axis. Drag f:

At f = 0.95 a hundred steps still leave 0.006 — small, but a number a float can hold and an optimiser can use. The RNN's product is under 0.01 by step nine, which is why it hugs the axis for the rest of the plot. And the forget gate is no constant of the architecture: the network learns it, per dimension, per step.

Which is exactly where the classic bug lives. At initialisation the weights are near zero, so f is whatever its bias says, and a bias of 0 means the cell halves itself every step. The curve is the share of c_0 left after twenty steps, against that bias — on a log axis, so eight decades fit:

The default frame is the bug. bias = 0 gives f = 0.500 and 9.5e-07 after twenty steps — an LSTM born with the vanishing gradient it exists to fix. Drag to bias = 1 and the same twenty steps keep 0.002. Write b_f = 0 and it trains badly for no visible reason; write b_f = 1 — Jozefowicz, Zaremba and Sutskever, 2015 — and it does not.

The first is capacity. However well a cell state survives, it is one vector of fixed width, and no gate makes it larger.

Below, the state is 512 numbers wide and never changes; above it are the tokens it has to summarise, and the first token's share is the block at the left. Drag right to lengthen the sequence:

At four tokens each share is 128 numbers. At 256 tokens it is 2.00, and the ticks above have merged into a texture. That is not a bug to be tuned away — it is what “summarise the prefix in a fixed-size vector” means, and it is why an LSTM asked about something specific 300 tokens back answers plausibly rather than correctly.

The second is worse, and it is about the machine rather than the model. Rows are positions, columns are wall-clock steps, and a cell lights when that position's work happens. Drag the step forward and watch how much of the grid is ever busy:

Only the diagonal can light. Position 5 cannot start until position 4 is finished, so eight positions take eight sequential steps, and by the last one only 8 of the 64 cells have done any work. A GPU has tens of thousands of lanes; batching different sequences fills some, but the work inside one sequence is a chain, and a chain has no width.

The GRU (Cho et al., 2014) is the popular economy: two gates instead of three, and the input gate is not free — it is forced to 1 − f, so what is kept and what is written are one decision. Drag it:

Because the two lengths must sum to the whole track, a GRU cannot hold the old state and write hard in the same step; an LSTM can, and that is the one thing the extra gate buys. What it saves is a gate's worth of matrices — one bar against the other two, on a fixed budget:

Each cell is a whole number of 2d² + d blocks: one for the RNN, three for the GRU, four for the LSTM — 2.10 M, 6.29 M and 8.39 M at d = 1024. Neither page above changes with any of it: the state is still one fixed vector, and the loop is still a loop.

Self-attention computes each position from every position at once: score every pair, softmax the scores, take the weighted sum. There is no state to carry, so nothing has to be finished before anything else can start.

Here is §04's grid again, with a switch. The recurrent layer lights one cell per column; attention lights a whole column at once. Lengthen the sequence, then switch the layer:

Sequential depth is the number of columns that carry anything: n on the left and 1 on the right, at every n the slider reaches. That is the difference between a GPU running at a few per cent of its width and one running full — and it is why the 2017 base Transformer trained in 12 hours on eight P100s where the LSTM systems it beat took days on far more hardware.

The second measurement is distance. Each rung is one multiplication between two tokens, so the height of the path is how many stand between them. Drag the target token, then switch the layer:

Switch the layer and watch the height collapse. Recurrence puts k multiplications between token 0 and token k — the very product §02 is about — while attention puts one edge there, whatever k is. The maximum path length goes from O(n) to O(1), and with it goes the decay.

The third is what the decoder may look at. An RNN encoder hands over its last state and nothing else, so a whole source sentence passes through one vector; attention lets each output position read every input position. Switch the link:

At twelve source tokens the bottleneck hands the decoder 512 numbers however long the source was; the direct link hands it 6,144, and 24 tokens make that 12,288. Bahdanau, Cho and Bengio saw this in 2014 and bolted attention onto an RNN; 2017 kept the attention and deleted the RNN.

None of it is free. A recurrent layer does n·d² multiply-adds — n steps of a d×d matrix — while an attention layer does n²·d, one score per pair. Both axes below are logarithmic, so both are straight lines and attention's is the steeper. Drag d:

The steeper line crosses the shallower one at n = d, where the two costs are equal. At d = 512 attention is cheaper below 512 tokens and dearer above.

# keep · write · what · show
f = sigmoid(W_f[h,x] + b_f)
i = sigmoid(W_i[h,x] + b_i)
g = tanh(W_g[h,x] + b_g)
o = sigmoid(W_o[h,x] + b_o)

# the only path back to c_0
c = f * c + i * g
h = o * tanh(c)

That loop is never trained to its start. Backpropagation through time keeps every activation on the chain, so training cuts it to a window: only hops inside the window get a gradient, and everything past the cut gets none. Drag the window:

The cut costs no information — the forward pass still carries the state across it — it costs credit: nothing before the window can be blamed for the error at the end.

Serving is where that trade becomes money. A decoder holds the KV cache for every position it has seen, while a recurrent state is the same size at any length. Drag the context across three decades:

At 128 K tokens a 32-layer, 4096-wide model in fp16 holds 64 GB against 256 KB — a factor of 262,144, and most of an 80 GB accelerator before a weight is loaded. Grouped-query attention cuts it fourfold.

So the loop is returning where the crossing says it should: Mamba (Gu and Dao, 2023) and RWKV run a recurrent state whose update is arranged to compute in parallel while training. For the windows an LLM serves today, attention's three wins still decide it — but the thing worth carrying out of this page is older than any of them: a path with one multiplication on it does not decay.