Lecture 6: Building a WaveNet

What this lecture covers

The lecture answers how to make the makemore character-level model deeper without simply widening a single hidden layer. It replaces the flat multi-layer perceptron from earlier lectures with a hierarchical, tree-like architecture that progressively fuses pairs of characters, arriving at something structurally similar to DeepMind's WaveNet. It follows directly from the MLP and BatchNorm work in the earlier makemore lectures (the manual backprop lecture is treated as a self-contained aside).

Along the way, the existing custom layers (Linear, BatchNorm1d, Tanh) are reorganized into reusable, PyTorch-like building blocks, including new Embedding and FlattenConsecutive modules and a Sequential container. After watching, you can explain why context is fused gradually instead of all at once, restructure a flat network into a hierarchical one, and recognize a subtle BatchNorm bug that appears when moving from two-dimensional to three-dimensional tensors.

Key ideas

• Progressive fusion: instead of flattening all context characters into one big vector before the first layer, WaveNet-style architectures fuse two characters at a time, then pairs of pairs, and so on, preserving structure longer.
• PyTorch-like containers: a custom Sequential module holds a list of layers and applies them in order, mirroring torch.nn.Sequential.
• Matrix multiply over extra batch dimensions: PyTorch's matmul operates on the last dimension and treats all leading dimensions as batch dimensions, which is what makes grouped, parallel processing of character pairs possible.
• FlattenConsecutive: a custom layer that reshapes (B, T, C) into (B, T/n, C*n), fusing n consecutive elements into the last dimension instead of flattening everything at once.
• BatchNorm dimensionality bug: applying the two-dimensional BatchNorm implementation to three-dimensional input silently computes per-position statistics instead of per-channel ones; fixing it means reducing over dimensions (0, 1) instead of just 0.
• Model state and bugs: forgetting to set a layer's training flag correctly, or leaving a module in a stale mode, can silently corrupt results, as shown when a single-example forward pass hits BatchNorm in training mode.
• Convolutions as an efficiency trick: sliding this same tree structure over an input sequence using convolutions computes the same result as looping in Python, but faster, because it reuses shared intermediate nodes and moves the loop into fast kernels.

Walkthrough

Intro and motivation (0:00)

The lecture explains the goal: take more than three characters of context and process them through a deeper network that fuses information gradually, rather than squashing everything into one hidden layer at once. This is connected to the 2016 WaveNet paper's hierarchical, tree-like prediction structure.

Starter code walkthrough (0:01:40)

The code picks up from the earlier MLP lecture, reusing the custom Linear, BatchNorm1d, and Tanh layer classes built to mimic torch.nn's API. The dataset and initial training loop are unchanged, reaching a validation loss around 2.10.

Fixing the loss plot and simplifying the forward pass (0:06:56)

The noisy per-step loss plot is smoothed by reshaping the loss list into rows and averaging. The embedding lookup and the flattening/concatenation step are then extracted into dedicated Embedding and Flatten modules so the entire forward pass becomes a single call through a Sequential model.

Pytorchifying the code: layers, containers, and a bug (0:09:16)

After introducing the Sequential container, a bug appears: sampling from the model with a single example produces gibberish because BatchNorm was left in training mode and computed a variance over one example, which is undefined. The fix highlights how easy it is to introduce bugs by not tracking a layer's train/eval state.

Implementing the WaveNet-style architecture (0:21:36)

The context length is increased from 3 to 8 characters. A naive flat model with the larger context already improves validation loss to about 2.02. The lecture then builds the FlattenConsecutive layer, which reshapes (B, T, C) into (B, T/n, C*n) by fusing n consecutive character embeddings, and stacks multiple Linear, BatchNorm1d, Tanh blocks with this flattening between them to build the tree-like structure, matching parameter count to the earlier flat model for a fair comparison.

Fixing the BatchNorm1d bug for 3D inputs (0:38:50)

Once the network processes three-dimensional tensors, the existing BatchNorm1d computes per-position statistics (1 x 4 x 68) instead of per-channel statistics, because it only reduces over dimension 0. The fix reduces over dimensions (0, 1) when the input is three-dimensional, correctly treating both the batch and the group position as batch dimensions. Retraining with the fix gives a small improvement, from 2.029 to 2.022.

Scaling up and final results (0:46:07)

Increasing embedding size and hidden units within the same hierarchical architecture pushes validation loss down to about 1.993, though the lecture notes there is no systematic experimental harness yet, only manual guessing and checking.

Convolutions, torch.nn, and the development process (0:47:44)

The lecture connects the hand-built tree structure to convolutions: a convolution slides the same small network across an input sequence and reuses shared intermediate computations, which is what makes it efficient compared to independently forwarding every position. It closes with reflections on the deep learning development process (prototyping in notebooks, wrestling with tensor shapes, PyTorch's inconsistent documentation) and previews future topics: dilated causal convolutions, residual and skip connections, and recurrent architectures.

Before you watch

• Complete the earlier makemore MLP lecture that builds the Linear, BatchNorm1d, and Tanh layers and the initial two-layer network, since this lecture's starter code builds directly on it.
• Be comfortable with BatchNorm's training/eval mode distinction and running statistics from that lecture.
• Familiarity with torch.view and broadcasting rules will help with the shape reasoning throughout.

Check your understanding

1. Why does fusing two characters at a time, rather than flattening all context characters immediately, allow the network to be made deeper productively?
2. What change does FlattenConsecutive make to the tensor shape, and how does it differ from the original Flatten layer?
3. Why did the two-dimensional BatchNorm1d implementation silently produce incorrect per-position statistics once the input became three-dimensional, and what dimensions does the fix reduce over?
4. What caused the single-example sampling bug, and why does it relate to BatchNorm specifically?
5. In what sense does a convolution over a sequence compute the same result as looping over each position with the tree-structured model, only faster?