LSTM, GRU and Backpropagation Through Time Explained With Numbers

Read this sentence: the cat, which already ate three bowls of food and ignored every toy in the house, was full. To pick the word was, you must remember the cat, twenty words back. A simple recurrent network usually can't. And the reason is not mysterious. It is a product of numbers, and today we will compute it.

• Short memory is a math problem
• Gradients shrink across time

This is lesson 53. In lesson 28 we met the recurrent network, and in lesson 52 we finished convolutional architectures. Today we derive backpropagation through time, see exactly why gradients vanish or explode, fix it with the LSTM and the GRU, and then write an LSTM cell in NumPy and check it against PyTorch.

• BPTT and the Jacobian product
• LSTM and GRU gates
• Training tricks and code

Recall the recurrent network. At each step it takes the input x t and the previous hidden state h t minus one, and computes h t equals tanh of W times h t minus one, plus U times x t, plus a bias. The key point is that the same matrices W and U are reused at every time step. That sharing is what makes the network recurrent.

• h_t = tanh(W h_{t-1} + U x_t + b)
• Same weights every step

To train it, we unroll the loop. Imagine a very deep feedforward network with one layer per time step, where every layer has the same weights. Inputs x one through x T go in, hidden states h one through h T come out, and each hidden state can produce a loss. Once unrolled, it is just an ordinary computational graph, as in lesson 36.

• One copy of the cell per step
• Shared weights W, U

Backpropagation through time, or BPTT, is simply backpropagation on that unrolled graph. Because W appears at every step, its gradient is the sum of contributions from each step. The total gradient of the loss with respect to W equals the sum over t of the partial of the loss with respect to h t, times the local effect of W on h t. Nothing new, except the sum.

• BPTT = ordinary backprop on unrolled graph
• Sum gradients over all steps

Now the crucial part. How does the loss at the final step T depend on an early hidden state h k? By the chain rule, we multiply the local derivatives of every step in between. The gradient of L T with respect to h k equals the gradient with respect to h T, times the product, from j equals k plus one up to T, of the Jacobian of h j with respect to h j minus one.

• dL_T/dh_k = dL_T/dh_T times product
• Each factor: dh_j / dh_{j-1}

Each Jacobian is easy to write. Since h j is tanh of a pre-activation, its derivative is one minus h j squared, placed on a diagonal, multiplied by W. So J j equals diag of one minus h j squared, times W. The tanh slope is at most one, and the matrix W is the same every time. We are multiplying nearly the same matrix again and again.

• J_j = diag(1 - h_j^2) W
• tanh derivative times W

Strip it down to one number. Let the hidden state be a single scalar and ignore tanh, so h t equals w times h t minus one. Then the derivative of h T with respect to h zero is just w multiplied by itself T times, which is w to the power T. Everything about vanishing and exploding gradients is hiding in that power.

• h_t = w h_{t-1}
• dh_T/dh_0 = w^T

Take w equal to zero point five. After ten steps the gradient is zero point five to the tenth, about one thousandth. After twenty steps it is about nine point five times ten to the minus seven, under one in a million. The error signal from step twenty reaches step zero as essentially nothing, so the network cannot learn that the early word mattered.

• 0.5^10 ≈ 0.001
• 0.5^20 ≈ 0.00000095

Now take w equal to one point five. After ten steps the factor is about fifty-eight. After twenty steps it is about three thousand three hundred and twenty-five. A tiny error becomes a huge update, and the weights get thrown across the loss surface. That is an exploding gradient, and training often ends with not a number values in the loss.

• 1.5^10 ≈ 58
• 1.5^20 ≈ 3,325

For a real network, W is a matrix, and the behavior is set by its singular values, from lesson 5. If the largest singular value of the Jacobian is below one, the product shrinks toward zero. If it is above one, the product can blow up. Only a razor-thin balance near one keeps the signal alive over long distances.

• Largest singular value decides
• Below 1 vanishes, above 1 explodes

Tanh makes vanishing even easier. Suppose w is zero point nine, which looks safe, but the hidden unit sits where the tanh slope is only zero point five. Each step multiplies the gradient by zero point four five. Over ten steps that is about three point four times ten to the minus four. A weight below one is not even required to get a vanishing gradient.

• Per step: 0.9 times 0.5 = 0.45
• 0.45^10 ≈ 0.00034

A common mistake is to see a flat loss and blame the data or the learning rate. If the model never uses long-range context, suspect the gradients first. Log the gradient norm with respect to each hidden state, from the last step back to the first. If it decays exponentially as you move backward, you have found the culprit.

• Loss plateaus are not always data issues
• Check gradient norms per time step

The LSTM, long short-term memory, fixes this with one design principle: give the network a memory lane that is updated by addition, not by repeated matrix multiplication and squashing. This lane is the cell state, called c t. Information can ride along it, step after step, while small gates decide what to erase, what to write, and what to reveal.

• Add, do not overwrite
• A protected memory lane

A gate is a sigmoid layer whose output lies between zero and one, applied by element-wise multiplication. A value near zero blocks a signal, a value near one lets it through, and values in between let a fraction pass. Each gate looks at the current input x t and the previous hidden state h t minus one, concatenated together.

• Sigmoid outputs between 0 and 1
• Multiply to open or close

Here are the equations. The forget gate f, the input gate i, and the output gate o are each a sigmoid of a linear function of x t and h t minus one. The candidate g is a tanh of another linear function. Then the cell is c t equals f times c t minus one, plus i times g. And the hidden state is h t equals o times tanh of c t.

• f, i, o gates and candidate g
• c_t = f ⊙ c_{t-1} + i ⊙ g
• h_t = o ⊙ tanh(c_t)

Think of it like a notebook. The forget gate decides which old notes to erase. The input gate decides how much of the new candidate to write down. The output gate decides how much of the notebook to read out as the hidden state, which the rest of the network sees. The notebook itself, the cell state, is private.

• Forget: what to erase
• Input: what to write
• Output: what to reveal

Let's compute one scalar step. Suppose c t minus one is two. After the linear layers, the pre-activations are: forget one, input zero, candidate one, output two. Sigmoid of one is zero point seven three one, so f is point seven three one. Sigmoid of zero is zero point five, so i is point five. Tanh of one is zero point seven six two, so g is point seven six two. Sigmoid of two is point eight eight one, so o is point eight eight one.

• x=1, h_prev=0.5, c_prev=2
• f=0.731, i=0.5, g=0.762, o=0.881

Now the update. The new cell state is point seven three one times two, which is one point four six two, plus point five times point seven six two, which is point three eight one. That gives c t equals one point eight four three. Tanh of that is about point nine five one. Multiply by the output gate, point eight eight one, and the hidden state is about point eight three eight.

• c = 0.731·2 + 0.5·0.762 = 1.843
• h = 0.881 · tanh(1.843) ≈ 0.838

Now look at the gradient along the cell state. Because c t is f times c t minus one plus something, the direct derivative of c t with respect to c t minus one is just f, here point seven three one. No weight matrix and no tanh sit on this path. If the network learns to keep f close to one, the product of these factors stays near one, and the gradient survives.

• dc_t/dc_{t-1} = f (direct path)
• No matrix, no squashing

A practical detail. If the forget gate starts near zero point five, then after twenty steps the cell path scales by roughly one millionth, and you are back to vanishing gradients. Many practitioners initialize the forget gate bias to one, so sigmoid starts near zero point seven three and memory begins open. Remember this when training LSTMs on long sequences.

• Init forget bias near 1
• Otherwise memory starts leaky

The gated recurrent unit, or GRU, keeps the same idea with fewer parts. It has no separate cell state. It has an update gate z, which decides how much of the old state to keep, and a reset gate r, which decides how much of the old state to use when forming the new candidate. Fewer gates means fewer parameters.

• Two gates: update and reset
• No separate cell state

The final GRU update is h t equals one minus z, times h t minus one, plus z times the candidate state. It is a blend of old and new, with weights adding to one. When z is near zero, the state is copied forward unchanged, which gives the same additive-style gradient path as the LSTM, with a derivative of one minus z along the direct route.

• h_t = (1-z) ⊙ h_{t-1} + z ⊙ h̃
• Convex blend of old and new

Let's count. A block has an input matrix of size h by d, a recurrent matrix of size h by h, and biases. With input size one hundred twenty-eight and hidden size two fifty-six, one block has thirty-two thousand seven hundred sixty-eight plus sixty-five thousand five hundred thirty-six weights, which is ninety-eight thousand three hundred four, plus five hundred twelve biases in PyTorch, giving ninety-eight thousand eight hundred sixteen.

• LSTM: 4 blocks, GRU: 3 blocks
• d=128, h=256 example

A vanilla RNN has one block, so ninety-eight thousand eight hundred sixteen parameters. A GRU has three blocks, giving two hundred ninety-six thousand four hundred forty-eight. An LSTM has four blocks, giving three hundred ninety-five thousand two hundred sixty-four. So the GRU is exactly twenty-five percent smaller than the LSTM. In practice their accuracy is often close, so try both.

• RNN 98,816
• GRU 296,448
• LSTM 395,264

Now practical training. When the model generates text, its input at each step is the previous token. In teacher forcing, during training we feed the true previous token from the data rather than the model's own guess. This makes training stable and parallel across targets. The drawback is exposure bias: at test time the model sees its own, possibly wrong, outputs.

• Feed the true previous token
• Fast, but train-test mismatch

Batches need equal lengths. Suppose three sequences have lengths three, five and two. We pad them with a special token up to length five. If we compute the loss on pad positions, we teach the model to predict padding, which is wrong. So we build a mask of ones and zeros, multiply the per-step loss by it, and divide by the number of real tokens, ten, not fifteen.

• Pad to equal length
• Mask the loss on pad positions

Masking the loss is not enough, because for a forward LSTM the hidden state still passes through pad steps, and for the last state this is wrong. PyTorch offers packed sequences. You sort by length, pack, and the LSTM processes only real tokens, so the final hidden state is the one at each sequence's true last step. Then you unpack the output if needed.

• Skip padded steps entirely
• pack_padded_sequence in PyTorch

For very long sequences, unrolling everything is expensive. Truncated BPTT splits the sequence into windows, say length one hundred. You carry the hidden state forward between windows, but detach it so gradients stop at the window boundary. Memory stays bounded, though the model cannot learn dependencies longer than the window through gradients alone.

• Backprop over windows of k steps
• Carry state, cut the graph

LSTMs reduce vanishing, but exploding gradients can still happen. The standard fix is clipping by global norm. If the gradient norm exceeds a threshold, scale the whole vector down. For example, a gradient of six and eight has norm ten. With a threshold of one, we multiply by one tenth and get zero point six and zero point eight. The direction is preserved, only the length changes.

• If norm > threshold, rescale
• [6, 8] norm 10 -> [0.6, 0.8]

Time to code. In NumPy, we concatenate x and h previous, multiply by one big weight matrix of shape four h by d plus h, add a bias, and slice the result into four chunks. Apply sigmoid to three of them and tanh to one. Then c equals f times c previous plus i times g, and h equals o times tanh of c. About ten lines.

• Stack four gates into one matrix
• Slice into f, i, g, o

How do we know it's right? Create a PyTorch nn.LSTM, then copy its weights into our cell. Careful: PyTorch orders its gates as input, forget, cell candidate, output, and it keeps two bias vectors, which we add together. Run the same input through both and call allclose. If the difference is around ten to the minus six, our implementation matches.

• Copy weights from PyTorch
• Compare outputs with allclose

Most bugs are small. Wrong gate order gives outputs that look plausible but differ. Forgetting that PyTorch has two biases shifts every result. Mixing batch-first and sequence-first shapes silently breaks things. Always test on one tiny sequence, compare against the reference at each step, and only then scale up to a real dataset.

• Wrong gate order
• Forgetting both biases
• Shape mix-ups

Now a character-level language model. Given a text such as hello, the inputs are h, e, l, l and the targets are e, l, l, o. Each character is a one-hot vector or an embedding, the LSTM updates its state, and a linear layer with softmax predicts a distribution over the vocabulary. We use cross-entropy from lesson 10 as the loss.

• Predict the next character
• Targets are inputs shifted by one

What should the loss look like? Before training, the model guesses uniformly, so the loss is the natural log of the vocabulary size. With sixty-five characters, that is about four point one seven. If your first loss is far from that, something is wrong. Perplexity is e to the loss, so a loss of one point five means the model is as uncertain as choosing among about four and a half characters.

• Start near ln(V)
• Vocabulary 65: about 4.17

The training loop is familiar from lesson 50: forward, loss, backward, clip the gradient norm, optimizer step. To generate text, feed a seed character, sample the next one from the softmax, and feed it back in. A temperature divides the logits: low temperature gives safe, repetitive text, high temperature gives creative but messier text. Early samples look like random letters, later ones form words.

• Clip gradients, truncated BPTT
• Sample with temperature

Let's recap. BPTT is backpropagation on the unrolled network, and the gradient to an early step is a product of Jacobians. Powers like zero point five to the twenty vanish and one point five to the twenty explode. The LSTM adds an additive cell path controlled by forget, input and output gates. The GRU simplifies it and is twenty-five percent smaller. Masking, packing and gradient clipping make training practical.

• BPTT: Jacobian product over time
• LSTM: additive cell path
• GRU: 25% fewer parameters

Recurrent models still read text one step at a time, and that sequential bottleneck is what we will relieve later in the course. But first, in lesson 54, we ask a more basic question: how do we turn raw text into numbers? We will study tokenization and embeddings.

• Lesson 54 next
• Subscribe to continue