Lecture 19: RNN Implementation

What this lecture covers

This lecture answers a practical question: what does an LSTM actually compute, line by line? Building on the earlier lecture that introduced recurrent networks conceptually, it walks through a from-scratch NumPy implementation of an LSTM cell, checks it against PyTorch's built-in LSTMCell and LSTM modules for numerical agreement, then extends the cell to a full sequence and to batched inputs.

After watching, you can implement an LSTM cell and a full sequence LSTM in plain array code, explain why LSTM inputs are laid out time-first rather than batch-first, and describe how truncated backpropagation through time and hidden-state repackaging make training on long sequences tractable. This directly prepares you for the homework assignment that asks you to reimplement the same logic in the needle automatic differentiation framework.

Key ideas

• Gate vector: PyTorch's LSTM packs the four gate computations (input, forget, cell/candidate, output) into two big weight matrices, weight_hh and weight_ih, each four times the hidden size, rather than four separate small matrices.
• Single matrix multiply, then split: an LSTM cell computes W_hh @ h + W_ih @ x + b once, then splits the resulting vector into four equal chunks with np.split, applying sigmoid to three of them and tanh to one.
• Cell state update: the new cell state is c_out = f * c + i * g (element-wise), and the new hidden state is h_out = o * tanh(c_out).
• Time-first batching: LSTM inputs are stored as (time, batch, hidden) rather than (batch, time, hidden) so that slicing out all examples at a single time step returns a contiguous block of memory, which matrix multiplication needs.
• Sequence vs. cell: a full LSTM is just the cell function called in a loop over time steps, returning the hidden state at every step but only the final cell state.
• Truncated backpropagation through time: training on very long sequences by running the full compute graph end to end runs out of memory, so training instead splits the sequence into fixed-size blocks and backpropagates only within each block.
• Hidden unit repackaging: when moving from one truncated block to the next, the final hidden and cell states are detached from the compute graph and copied in as the next block's initial state, so information carries forward without keeping the whole graph in memory.

Walkthrough

Inspecting PyTorch's LSTM weights (1:01)

The lecture creates a PyTorch LSTMCell with input size 20 and hidden size 100, then inspects weight_hh and weight_ih. It shows these are 400x100 and 400x20 matrices: four stacked 100-row blocks, one for each gate, avoiding four separate matrix multiplications. It also notes PyTorch keeps two separate bias terms (bias_hh and bias_ih) that are always added together, which the lecture treats as a quirk and simplifies by summing them.

Implementing a single LSTM cell in NumPy (6:04)

A sigmoid function is defined manually since NumPy has no built-in. The LSTM cell function takes x, h, c, the two weight matrices, and a bias, computes one combined matrix-vector product, splits it into i, f, g, o, applies the appropriate nonlinearities, and computes the new cell and hidden states. The result is checked against PyTorch's LSTMCell output and matches to numerical precision.

Extending to a full sequence (14:14)

Since a single cell only advances one time step, the lecture defines an LSTM function that loops the cell over all time steps of an input sequence X, collecting every hidden state into an output array but keeping only the final cell state, matching PyTorch's LSTM module convention. This is validated against PyTorch's LSTM class, again matching to numerical precision.

Why LSTM inputs are time-first (22:45)

The lecture explains that batching requires slicing out a matrix of all examples at a given time step, and this slice is only contiguous in memory if the array is laid out as (time, batch, hidden) rather than the more familiar (batch, time, hidden) convention used elsewhere. It ties this back to earlier lectures on why compacting memory before matrix multiplication matters for performance.

Batched LSTM cell and sequence (29:57)

The cell function is rewritten to operate on batched matrices instead of vectors: h @ W_hh + x @ W_ih + b, splitting along the second axis instead of the first. A batch of 128 sequences of length 50 is tested against PyTorch's batched LSTM output, again matching to numerical precision, in roughly a dozen lines of code.

Training an LSTM and truncated backpropagation (36:04)

The lecture sketches, conceptually, how training would work if this were implemented in needle: call the LSTM forward, compute a loss against a target, then call loss.backward() and step an optimizer. It extends this to multi-layer (deep) LSTMs, computing one layer across all time steps before moving to the next layer. It then explains why running backpropagation across an entire long sequence (for example 10,000 steps) is impractical, motivating truncated backpropagation through time: the sequence is chopped into blocks (such as 100 steps), and gradients flow only within a block. Hidden unit repackaging is introduced as a refinement, where the final hidden and cell states of one block are detached and reused as the initial state of the next block, preserving information across blocks without carrying the compute graph forward.

Before you watch

• Be comfortable with the LSTM equations and gate structure from the preceding lecture on recurrent architectures.
• Review why matrix-matrix multiplication is preferred over repeated matrix-vector products, covered in the earlier lectures on efficient linear algebra.
• Know the basics of automatic differentiation and the needle framework's backward() and optimizer step pattern, since this lecture assumes it without re-deriving it.

Check your understanding

1. Why does PyTorch store weight_hh and weight_ih as 4H x H and 4H x I matrices rather than four separate H x H and H x I matrices?
2. Why must LSTM inputs be laid out as (time, batch, hidden) instead of (batch, time, hidden) for efficient batched computation?
3. What problem does truncated backpropagation through time solve, and what is the tradeoff of chopping a sequence into independent blocks?
4. How does hidden unit repackaging differ from simply resetting the hidden state to zero at the start of each block?