LSTMs & GRUs

This is the hands-on post, and it is intentionally code-heavy. The previous posts built understanding from concept through mechanics; here you assemble a complete, runnable character-level LSTM in PyTorch so the ideas become a working artifact rather than abstractions.

The model learns to predict the next character in a sequence — the simplest task that exercises everything a gated cell does. Each snippet is real code that runs in order, so you can copy the whole thing, train it, and then modify it to understand each piece by breaking it. The closing slides show that switching to a GRU is a single line.

Setup defines a tiny dataset and the encoding that turns characters into the integer ids a network can consume. We build two lookup tables — stoi (string to index) and itos (index to string) — and convert the text into a tensor of ids.

This encode/decode boundary is something every sequence model needs: networks operate on numbers, not raw text, so you always need a clean mapping in both directions. Keeping the vocabulary tiny here means the model trains in seconds, which is exactly what you want when experimenting and breaking things to learn.

The model class shows the canonical gated-RNN stack: an embedding layer mapping each id to a dense vector, an nn.LSTM that processes the sequence while maintaining its (h, c) state, and a linear head mapping each step's hidden output to a score over the vocabulary.

The forward method threads an optional state through the LSTM and returns it alongside the outputs. Returning the state matters for generation, where you feed the model one character at a time and must carry both the hidden and cell memory forward between calls. This embed -> LSTM -> linear pattern is the backbone of nearly every recurrent sequence model you'll write.

The training loop is where backpropagation through time happens, and notice how little code it takes. You form inputs and next-character targets by shifting the sequence by one, run a forward pass to get logits, compute cross-entropy loss, and call loss.backward(), which triggers autograd to run BPTT across the entire unrolled sequence.

The one recurrent-specific addition is clip_grad_norm_ before opt.step(). This matters even for LSTMs: gating fixes the vanishing gradient but does nothing for the exploding one, so clipping is still mandatory. Everything you traced by hand in the mechanics post is compressed into these few lines; the framework handles the bookkeeping while you know exactly what it's doing.

The pipeline diagram names the five stages every training step moves through: encode characters into ids, run the LSTM forward over all time steps, compute the next-character cross-entropy loss, run BPTT backward through time, and clip-then-step with the optimizer.

Seeing the stages laid out helps when debugging. A shape error lives in encode or forward; a NaN loss points to the backward/clip stage and means your clipping is missing or too loose; a flat loss curve often means the learning rate or step is wrong. The diagram doubles as a mental checklist for the loop above, and the explicit clip stage is a reminder that recurrent training without it is fragile.

Generation demonstrates the model's learned behavior and a subtle mechanic: stateful, one-step-at-a-time inference. You seed with a character, then repeatedly feed the current character and the carried state back in, take the argmax of the logits as the next character, and append it.

The crucial part is threading the state through every call — for an LSTM that state is the (h, c) tuple, and carrying both is how the model remembers what it has generated so far. Using argmax gives deterministic greedy generation; sampling from the softmax instead would produce more varied output. This loop is the inference-time mirror of the training loop, and it's where forgetting to pass the state back produces gibberish.

These notes on the (h, c) state tuple capture the single biggest difference between using an LSTM and using an RNN or GRU in PyTorch. nn.LSTM returns its state as a tuple (h_n, c_n): h_n is the hidden output and c_n is the cell memory carrying long-term information. nn.GRU and nn.RNN return a single tensor instead.

For stateful inference you must pass the whole tuple back in each call, not just h_n, or you silently drop the cell memory. And between truncated chunks during training you detach the state to cut the autograd graph while keeping the values. Getting this tuple handling right is the difference between an LSTM that remembers and one that quietly resets, which is why it gets its own slide and reappears in the mistakes post.

The GRU swap shows off the modularity of PyTorch's RNN family: changing nn.LSTM to nn.GRU is a one-line edit and the training loop keeps working — with one caveat. The GRU returns a single state tensor, not a (h, c) tuple, so you must adjust how you unpack the state wherever you use it.

The slide also previews two production levers: num_layers stacks multiple recurrent layers for more capacity, and bidirectional=True runs the sequence forward and backward and concatenates the results, which helps when the whole sequence is available at once (not for streaming). Both are single keyword arguments, which is why people reach for them so readily — and why understanding their cost matters.

This slide catalogs the runtime errors that actually bite, several of which preview the dedicated mistakes post. Unpacking the LSTM state as a single tensor crashes because it's a (h, c) tuple, while the GRU returns one tensor — code written for one breaks on the other. Forgetting batch_first produces a silent dimension mismatch. Passing softmax outputs into CrossEntropyLoss double-applies the softmax and weakens gradients. Carrying state across batches without .detach() grows the autograd graph until you run out of memory.

The pattern is that the worst recurrent-net bugs are quiet — they don't crash, they just give you a model that won't learn. Knowing them in advance is the cheapest debugging you'll ever do.

The data-flow diagram traces tensor shapes through the model, the single most useful thing to hold in your head when wiring a sequence model. Integer ids of shape (B, T) become embeddings (B, T, H), pass through the LSTM preserving (B, T, H), and the linear head maps them to logits of shape (B, T, V) — one score per vocabulary item per time step.

Most wiring bugs are shape bugs, so being able to recite this chain lets you catch a mismatch by reasoning rather than by trial and error. When something breaks, mentally walk these shapes and the offending layer usually reveals itself. Note that the LSTM's internal cell state doesn't appear in this output shape — it's carried separately in the state tuple.

This recap consolidates the build: the embed -> LSTM -> linear head stack is the standard skeleton; loss.backward() performs BPTT for you; gradient clipping is mandatory even for gated cells; the LSTM state is an (h, c) tuple while the GRU state is a single tensor; and swapping between nn.LSTM and nn.GRU takes almost no changes.

With a working model in hand, the only thing left is to learn the traps — the quiet mistakes that break LSTM and GRU training even when the code looks right. That is the focus of the final post.

The teaser points to the common-mistakes post. You now have a model that runs; the next step is hardening it by learning the quiet failure modes — state-tuple mishandling, skipping gradient clipping, defaulting to LSTM over GRU, ignoring the forget-gate bias, the wrong dropout, and padding without packing — and the one-line fix for each.