ML Grind

Sequence Models

Unroll a vanilla RNN with hand-written BPTT, then build single steps of PyTorch-compatible LSTM and GRU cells.

Watch the animated Sequence Models lecture (3 min) →

Sequence Models

Why It Matters

Text, audio, time series and user sessions are all sequences: order matters and length varies. Before transformers, recurrent networks (RNN, LSTM, GRU) were the standard way to model them, and they still power streaming speech recognition, on-device keyboards and many forecasting systems. State-space models (Mamba, RWKV) bring the same recurrent idea back to LLMs, and interviewers still ask you to write an LSTM step or explain vanishing gradients.

Core Ideas

The recurrence

An RNN reads one element at a time and updates a fixed-size hidden state that summarizes the past, reusing the same weights at every step:

h_t = tanh(x_t @ W_xh + h_{t-1} @ W_hh + b)        x_t: (N, D),  h_t: (N, H)

All problems here are batch first: a batch of sequences is X of shape (N, T, D) = (batch, seq_len, features) and the stacked hidden states are (N, T, H).

Backpropagation through time (BPTT)

Training unrolls the loop and backpropagates through every step. The gradient reaching h_1 from h_T is a product of T - 1 Jacobians:

dh_T/dh_1 = prod over t = T..2 of  diag(1 - h_t^2) @ W_hh.T

If the largest singular value of W_hh (times the tanh slope, at most 1) is below 1 the product vanishes; above 1 it can explode. Gradient clipping fixes explosions but not vanishing, so plain RNNs struggle beyond a few dozen steps.

Gated cells

Gates are sigmoids in [0, 1] that decide, per dimension, how much information flows.

CellStateGatesUpdate
RNNhnoneh = tanh(x W_xh + h W_hh + b)
LSTMh, cinput i, forget f, output o (+ candidate g)c = f*c_prev + i*g, h = o*tanh(c)
GRUhreset r, update z (+ candidate n)h = (1-z)*n + z*h_prev

The LSTM's cell state is updated additively: when f ≈ 1, dc_t/dc_{t-1} ≈ 1 and gradients travel hundreds of steps. That is why the forget-gate bias is often initialized to 1.

Conventions used in this section (memorize them)

Vanilla RNN problemLSTM (nn.LSTMCell)GRU (nn.GRUCell)
Weight layoutx @ W_xh, W_xh: (D, H), W_hh: (H, H)x @ W_ih.T, W_ih: (4H, D), W_hh: (4H, H)x @ W_ih.T, W_ih: (3H, D), W_hh: (3H, H)
Biasesone b: (H,)b_ih, b_hh: (4H,)b_ih, b_hh: (3H,)
Gate order (row blocks of H)nonei, f, g, or, z, n
ParametersDH + H^2 + H4(HD + H^2 + 2H)3(HD + H^2 + 2H)

Note the layout switch: the RNN problem uses (in, out) weights, the LSTM and GRU problems use PyTorch's (out, in) weights, hence the .T. TensorFlow/Keras orders LSTM gates i, f, c, o too, but stores kernels as (D, 4H).

Worked Micro-Example

RNN, scalar (D = H = 1), W_xh = 0.5, W_hh = 1, b = 0, h_0 = 0, input x_t = 1 for 3 steps:

h_1 = tanh(0.5)            = 0.4621
h_2 = tanh(0.5 + 0.4621)   = 0.7452
h_3 = tanh(0.5 + 0.7452)   = 0.8469

LSTM, one step with H = 1, c_prev = 1 and gate pre-activations (i, f, g, o) = (0, 2, 1, 0):

i = sigmoid(0) = 0.5,  f = sigmoid(2) = 0.8808,  g = tanh(1) = 0.7616,  o = 0.5
c = 0.8808 * 1 + 0.5 * 0.7616 = 1.2616
h = 0.5 * tanh(1.2616)       = 0.4258

Vanishing in numbers: a per-step gradient factor of 0.5 over 20 steps leaves 0.5^20 ≈ 9.5e-7 of the signal; a factor of 1.1 over 50 steps blows it up to about 117.

Common Pitfalls

  • Gate order: splitting an LSTM's 4H block as i, f, o, g (or a GRU's as z, r, n) runs without errors and gives wrong numbers. PyTorch is i, f, g, o and r, z, n.
  • Weight orientation: x @ W vs x @ W.T. Check the shape table before multiplying.
  • GRU reset gate: PyTorch applies r after the hidden matmul, to h_prev @ W_hn.T + b_hn; the original paper applies it to h_prev first. The numbers differ.
  • GRU update gate: in PyTorch z = 1 keeps h_prev; some papers swap z and 1 - z.
  • BPTT: weights are shared, so their gradients accumulate over time; the gradient into h_t has a direct part and a part from h_{t+1}.
  • Batch vs time axis: nn.RNN defaults to (T, N, D) unless batch_first=True. These problems are batch first.

Interview Angle

  • "Write an LSTM cell from scratch." Then: "count its parameters", "why does the cell state help gradients?", "why initialize the forget bias to 1?"
  • "Explain vanishing and exploding gradients in RNNs and the fixes": gating, clipping, orthogonal init, shorter truncated BPTT windows.
  • "LSTM vs GRU?" GRU has 3 weight blocks instead of 4 (25% fewer parameters) and no separate cell state; quality is usually similar.
  • "RNN vs Transformer?" RNN training is sequential in T but inference costs O(1) per token; attention trains in parallel but costs O(T) per token with a KV cache and O(T^2) for the full sequence.

Problems in This Section

  1. Vanilla RNN: Forward & BPTT: unroll h_t = tanh(x_t W_xh + h_{t-1} W_hh + b) over (N, T, D) inputs and backpropagate through time by hand.
  2. LSTM Cell: one step of nn.LSTMCell with gate order i, f, g, o.
  3. GRU Cell: one step of nn.GRUCell with gate order r, z, n and PyTorch's reset-gate placement.

Related: Gradient Clipping by Global Norm (Optimization & Training) for exploding gradients, and Scaled Dot-Product Attention (Transformer Architecture), which replaced recurrence in most language models.

Practice exercises

  1. Vanilla RNN: Forward & BPTT →
  2. LSTM Cell →
  3. GRU Cell →