LSTMs in NLP, explained in 5 lines

Module 04 · NLP · 20 concepts

In depth / 15 lines

LSTMs, in more detail

An LSTM is a recurrent network redesigned so that memory can survive long gaps.

It adds a second state alongside the hidden one, called the cell state.

The cell state runs the length of the sequence with only light interference.

Information can ride along it almost unchanged, which is what plain RNNs could not do.

Three gates decide what happens to it, and each gate is a small learned network.

The forget gate decides what fraction of the current memory to discard.

The input gate decides how much of the new information to write in.

The output gate decides how much of the memory to expose at this step.

Each gate outputs numbers between zero and one, which act as soft switches.

Because they are learned, the network works out for itself what is worth keeping.

The near-uncluttered path along the cell state lets gradients flow back much further.

That is what extends useful memory from a handful of steps to a few hundred.

The cost is four times the parameters and four times the computation of a plain RNN.

LSTMs dominated translation, speech and text generation from the late 1990s until 2017.

They still suit long streaming sequences where a transformer would be too expensive.

Want the full technical article? Wikipedia ↗

Deep dive

What is underneath LSTMs

A path the gradient can travel

\mathbf{c}_t = \mathbf{f}_t \odot \mathbf{c}_{t-1} + \mathbf{i}_t \odot \tilde{\mathbf{c}}_t

$\mathbf{c}_t$ : the cell state, the long-term memory.
$\odot$ : elementwise multiplication.

That line is the whole idea.

The cell state is carried forward by addition rather than by a matrix multiply, so the gradient flowing back through it is multiplied by $\mathbf{f}_t$ and nothing else.

\frac{\partial \mathbf{c}_t}{\partial \mathbf{c}_{t-1}} = \mathbf{f}_t

With a forget gate near 1 the gradient passes through unchanged across hundreds of steps, which is exactly what a plain recurrent network cannot do.

The three gates

\mathbf{f}_t = \sigma\big( W_f [\mathbf{h}_{t-1}, \mathbf{x}_t] + \mathbf{b}_f \big)

The forget gate decides, per element, how much of the old cell state survives. A sigmoid output of 0 erases it and 1 keeps it whole.

\mathbf{i}_t = \sigma\big( W_i [\mathbf{h}_{t-1}, \mathbf{x}_t] + \mathbf{b}_i \big)

\tilde{\mathbf{c}}_t = \tanh\big( W_c [\mathbf{h}_{t-1}, \mathbf{x}_t] + \mathbf{b}_c \big)

The input gate decides how much of the new candidate is written. The candidate itself is a tanh, so what gets written is bounded between $-1$ and 1.

\mathbf{o}_t = \sigma\big( W_o [\mathbf{h}_{t-1}, \mathbf{x}_t] + \mathbf{b}_o \big)

\mathbf{h}_t = \mathbf{o}_t \odot \tanh(\mathbf{c}_t)

The output gate decides how much of the cell is exposed as this step's hidden state, so the cell can hold something for later without acting on it now.

One timestep, in order

Concatenate the previous hidden state with the current input.
Compute all four gate pre-activations from that vector, as one matrix multiply.
Scale the old cell state by the forget gate.
Add the candidate, scaled by the input gate.
Squash the new cell state and mask it with the output gate to get $\mathbf{h}_t$ .

4 \big( d_h (d_h + d_x) + d_h \big)

The parameter count is four times a plain RNN's, since three gates and a candidate share the same shape.

The detail that decides whether it trains

\mathbf{b}_f = 1

The forget gate bias is initialised to 1 rather than 0, so the gate starts open and the cell state is preserved by default.

Initialised at zero the sigmoid sits at 0.5, which halves the memory at every step and destroys it in a few dozen, before the model has learned to keep anything.

That single line is the difference between an LSTM that learns long dependencies and one that quietly behaves like a plain RNN.

What it does not fix

\mathbf{h}_t \text{ still requires } \mathbf{h}_{t-1}

The computation is still sequential, so it cannot be parallelised across time and training is slow on hardware built for parallel work.

The state is still a fixed-size vector, so everything read so far must be compressed into the same few hundred numbers regardless of how much was said.

That compression is what attention removed by keeping every step's state available, and why encoder-decoder models with attention replaced the LSTM before transformers replaced both.

Test yourself

One question on this card. Every option explains itself once you pick, right or wrong.

Question 1

What problem do LSTMs solve compared to plain RNNs?

NLP

20 concepts in this module, then a graded test with every answer explained.

Study this module →