ai_ml / deep learning / 05_rnn_lstm.md

RNNs and LSTMs

6 interview angles 4 min read source

RNNs and LSTMs

Largely replaced by transformers, and still worth knowing because the reasons they lost are exactly the reasons transformers won. That comparison is the interview value here.

The recurrent idea

Process a sequence one step at a time, carrying a hidden state forward:

h_t = tanh(W_hh @ h_{t-1} + W_xh @ x_t + b)
y_t = W_hy @ h_t

The same weights apply at every timestep — weight sharing across time, as convolution shares across space. The hidden state is a fixed-size summary of everything seen so far.

That fixed size is the first structural problem: an entire 500-token history must compress into one vector.

Backpropagation through time, and why it fails

Unroll the recurrence and backprop. The gradient across k steps involves W_hh multiplied by itself k times.

  • Largest eigenvalue < 1 → gradient vanishes exponentially → the network cannot learn long-range dependencies.
  • Largest eigenvalue > 1 → gradient explodes → NaN.

Exploding gradients are easy to fix with clipping. Vanishing gradients are the fundamental limitation — a plain RNN reliably learns dependencies of maybe 10-20 steps.

LSTM

Adds a cell state with additive updates and three learned gates:

f_t = sigmoid(W_f @ [h_{t-1}, x_t])     # forget: what to drop from the cell
i_t = sigmoid(W_i @ [h_{t-1}, x_t])     # input: what to write
o_t = sigmoid(W_o @ [h_{t-1}, x_t])     # output: what to expose
g_t = tanh(W_g @ [h_{t-1}, x_t])        # candidate values

c_t = f_t * c_{t-1} + i_t * g_t         # THE key line
h_t = o_t * tanh(c_t)

The cell state update is additive, not multiplicative. With the forget gate near 1, c_t ≈ c_{t-1} + something, so the gradient flows back along that path without repeated matrix multiplication. It’s the same trick as a residual connection in a CNN or transformer — give the gradient a highway.

That single line is the answer to “how does an LSTM solve vanishing gradients”, and it’s what makes 100+ step dependencies learnable.

GRU merges the forget and input gates into one update gate and drops the separate cell state. Fewer parameters, comparable performance, faster. When choosing between them, try GRU first.

Why transformers replaced them

Three reasons, in order of importance:

1. Sequential computation. An RNN must compute step t before step t+1. That’s O(n) sequential operations, and it cannot be parallelised across the sequence during training. A transformer processes all positions simultaneously, so it uses modern accelerators fully. This is the dominant reason — it’s not that transformers are smarter, it’s that they train on far more data in the same wall-clock time.

2. Path length. In an RNN, information from position 1 reaching position 500 traverses 500 steps, degrading along the way. In self-attention, every position is one operation from every other — constant path length regardless of distance.

3. The bottleneck. A fixed-size hidden state must compress the whole history. Attention keeps every position available and learns what to look at.

RNN/LSTM Transformer
Training parallelism none across time full
Path between distant tokens O(n) O(1)
Compute per layer O(n * d²) O(n² * d)
Memory O(n) O(n²) for attention scores
Long sequences cheap but forgetful expensive but accurate

Note the last rows: transformers are quadratic in sequence length, which RNNs are not. For very long sequences that’s a real cost, and it’s why efficient-attention and state-space models (Mamba and relatives) remain an active area — they aim for RNN-like linear scaling with transformer-like quality.

Where RNNs still appear

  • Very long or streaming sequences where quadratic attention is impractical.
  • Tiny models on constrained hardware — an LSTM can be a few hundred KB.
  • Classical time-series forecasting, though gradient boosting on lag features frequently beats both.
  • Legacy systems you’ll be asked to maintain or migrate.

Practical notes

nn.LSTM(input_size=128, hidden_size=256, num_layers=2,
        batch_first=True, bidirectional=True, dropout=0.2)
  • batch_first=True — otherwise PyTorch expects (seq, batch, feature), which is a routine source of silent shape bugs.
  • Bidirectional doubles the output size and requires the whole sequence up front, so it’s unavailable for streaming or autoregressive generation.
  • Pack padded sequences (pack_padded_sequence) so the RNN doesn’t process padding tokens and pollute the final hidden state.
  • Gradient clipping is effectively mandatory.

Interview angle

  • “Why do plain RNNs struggle with long sequences?” — backpropagation through time multiplies the recurrent weight matrix repeatedly, so gradients vanish or explode exponentially with distance. Clipping handles explosion; vanishing is the structural limit, capping useful memory at roughly 10-20 steps.
  • “How does an LSTM fix that?” — the cell state updates additively, gated by a forget gate. With the gate near 1 the gradient flows back along a near-identity path instead of through repeated matrix multiplication. Same principle as a residual connection.
  • “LSTM or GRU?” — GRU has fewer parameters and trains faster with comparable accuracy on most tasks. Try it first; use LSTM if you need the extra capacity.
  • “Why did transformers replace RNNs?” — mainly training parallelism. An RNN is inherently sequential across time; a transformer processes all positions at once and saturates modern hardware, so it consumes far more data per unit time. Also constant path length between any two positions, versus linear for an RNN.
  • “Is anything worse about transformers?” — attention is quadratic in sequence length in time and memory, where an RNN is linear. That’s why very long contexts are expensive and why state-space models pursuing linear scaling are an active research direction.
  • “When would you still choose an LSTM?” — streaming or very long sequences where quadratic cost is prohibitive, or severely constrained hardware where a small recurrent model fits and a transformer doesn’t.