Learn LLM

LSTM: Long Short-Term Memory

Long Short-Term Memory networks, or LSTMs, address the vanishing gradient problem by introducing a separate cell state c_t in addition to the hidden state h_t. The cell state acts like a protected conveyor belt that can carry information across many timesteps with only minimal changes.

LSTMs use three gates to regulate information. Each gate is a small neural network whose output is squashed between zero and one by a sigmoid. A value near one means "let everything through," while a value near zero means "block everything." These gates decide what to forget, what to store, and what to output.

f_t = \sigma\left(W_f \begin{bmatrix} h_{t-1} \ x_t \end{bmatrix} + b_f\right)

The forget gate decides what to throw away from the cell state. It looks at the previous hidden state and the current input, and it produces a vector of values between zero and one. It multiplies the old cell state element-wise, scaling each dimension down if it is no longer relevant.

i_t = \sigma\left(W_i \begin{bmatrix} h_{t-1} \ x_t \end{bmatrix} + b_i\right), \quad \tilde{c}_t = \tanh\left(W_c \begin{bmatrix} h_{t-1} \ x_t \end{bmatrix} + b_c\right)

The input gate decides what new information to store. The sigmoid layer determines which dimensions of the cell state should be updated, while the tanh layer creates candidate values that could be added. The product of the input gate and the candidate is the new information written into the cell state.

o_t = \sigma\left(W_o \begin{bmatrix} h_{t-1} \ x_t \end{bmatrix} + b_o\right)

The output gate decides what to reveal as the next hidden state. It filters the cell state through a tanh non-linearity and then uses the output gate values to decide which parts of the cell state should be exposed at this timestep.

c_t = f_t \odot c_{t-1} + i_t \odot \tilde{c}_t, \quad h_t = o_t \odot \tanh(c_t)

These two equations are the heart of the LSTM. The cell state c_t is updated by adding new information and removing old information, while the hidden state h_t is a filtered view of the cell state. The circle-dot is element-wise multiplication, so each gate controls one dimension of the cell independently.

Why gates work

The key insight is that the cell state update is mostly additive. If the forget gate is close to one and the input gate is close to zero, the cell state changes very little. It simply passes the previous cell state forward, almost unchanged, for as many timesteps as necessary.

During backpropagation, the gradient can flow straight through the unchanged cell state without being repeatedly multiplied by small weights. The forget gate acts as a throttle: values near one mean the gradient passes through freely, while values near zero mean the gradient is stopped at that timestep.

A gradient highway

LSTM does not eliminate vanishing gradients entirely. Instead, it gives them a fast lane. When the forget gate stays open, the gradient can flow through the cell state c_t like a skip connection, bypassing the nonlinear tanh that would otherwise shrink it at every step.

How LSTM reads a sentence

Let us walk through the sentence "The cat, which was black, sat." with an LSTM. We will track the cell state as the network processes the clause between the commas. The goal is to keep the subject "cat" available so that the model can connect it to the verb "sat".

Token   | Forget | Input | Output | Cell state contains | Notes
--------|--------|-------|--------|---------------------|----------------------------------
the     |   0.9  |  0.1  |  0.1   | the                 | article; low impact
cat     |   1.0  |  0.9  |  0.9   | the, cat            | stores the subject strongly
,       |   1.0  |  0.0  |  0.1   | the, cat            | punctuation; mostly passes state
which   |   1.0  |  0.2  |  0.2   | the, cat            | relative pronoun; preserves subject
was     |   1.0  |  0.1  |  0.1   | the, cat            | auxiliary verb; little new info
black   |   0.9  |  0.8  |  0.7   | the, cat, black     | adds a descriptive modifier
,       |   1.0  |  0.0  |  0.1   | the, cat, black     | end of clause; preserve state
sat     |   0.9  |  0.4  |  0.9   | the, cat, black, sat| verb connects to preserved subject

The forget gate stays close to one through the intervening clause, so the cell state still carries "cat" when the network reaches "sat". The output gate opens at "sat", allowing the model to use the preserved subject to interpret the verb. This is exactly the kind of long-range dependency a plain RNN struggles to learn.

Of course, these numbers are illustrative. A real LSTM learns gate values from data, and they will vary from sentence to sentence. The important point is that the architecture makes this kind of preservation possible.

Truncated BPTT and teacher forcing

Unrolling a long sequence creates a very deep computational graph, and backpropagating through every step is expensive and unstable. In practice we truncate the graph after a fixed number of timesteps, then continue from the last hidden state as a new starting point. This is called truncated backpropagation through time, and it is the standard way to train RNNs on long documents.

Seq2seq models also rely on teacher forcing during training: the decoder is fed the true previous token as input instead of its own prediction. This keeps gradients stable early in training. At inference the model must feed itself, which is why autoregressive decoding can drift and produce repetitive or nonsensical output.

Seq2Seq with Attention

LSTMs became the backbone of early neural machine translation. In the sequence-to-sequence, or seq2seq, framework introduced by Sutskever et al. and extended with attention by Bahdanau et al. in 2014, an encoder LSTM reads the source sentence and a decoder LSTM generates the translation.

The encoder processes the source sentence token by token and produces a final hidden state. In the simplest version, this single fixed-size vector, often called the context vector, is supposed to encode the meaning of the entire source sentence. The decoder then uses it as the starting point for generating the target sentence.

The fixed-size context vector is a bottleneck. A long sentence must be squeezed into a vector of the same size as a short sentence, and fine-grained source information can be lost. Attention solves this by letting the decoder look back at all encoder states, a topic we explore in the Attention chapter.

Why we moved on

LSTMs and GRUs dominated sequence modeling for years, but they have a fundamental speed problem. Because each hidden state depends on the previous one, tokens must be processed in order. You cannot compute the hidden state for the tenth word until you have computed the hidden states for the first nine words.

This sequential dependency makes it impossible to parallelize training over the sequence dimension. Modern GPUs and TPUs are most efficient when they can process many independent operations at once. RNNs and LSTMs leave most of that hardware unused.

Practical LSTMs also have a limited effective context window. Although they can in principle remember information indefinitely, in practice they begin to forget important details after roughly one hundred tokens. Modern Transformer models, which we cover in the Transformer chapter, can attend to 128,000 tokens or more and process them all in parallel.

GRU: a simpler gate design

GRU, or Gated Recurrent Unit, is a simplified LSTM with two gates instead of three. It often performs similarly while being faster to train. We cover LSTM because its explicit cell state makes the gradient highway more intuitive, but GRUs are a perfectly valid alternative.

RNN and LSTM cell explorer

Step through a sentence and watch how the hidden state and cell state update at each token. Switch between plain RNN and LSTM modes to see the effect of gating.

Sequence Cell Explorer

thecat,whichwasblack,sat

Step 1 of 8

c_t = f_t \odot c_{t-1} + i_t \odot \tilde{c}_t, \quad h_t = o_t \odot \tanh(c_t)

Forget gate f_t (avg 0.50)

0.50.50.50.40.50.6

Input gate i_t (avg 0.52)

0.50.50.50.60.50.5

Output gate o_t (avg 0.52)

0.60.60.50.50.50.5

Cell state c_t

-0.10.10.20.10.0-0.1

Hidden state h_t

-0.10.00.10.10.0-0.0

This is a synthetic seeded model intended to show the mechanics of recurrence. The exact gate values are not learned from data; they are generated from small random matrices so you can step through the update equations.

Key takeaway

LSTMs protect long-term information with a cell state and three gating mechanisms, creating a gradient highway that makes long-range dependencies learnable. Even so, their sequential nature limits parallelism and context length, setting the stage for attention and Transformers.

What comes next

The next chapter moves from recurrent cells to dense word vectors. Embeddings turn discrete tokens into continuous vectors where similar words live close together, preparing us for the attention and Transformer chapters that follow.