AI
Andrej KarpathyBuilding makemore Part 5: Building a WaveNet
synopsis
Starting from the part-3 character-level MLP (three characters of context feeding one hidden layer, validation loss 2.10), he restructures it into a WaveNet-style hierarchical network — progressively fusing pairs of characters across multiple layers instead of squashing all context into one layer at once — and lands, after fixing a BatchNorm bug and scaling up, at a validation loss of 1.993 on the same names dataset, while also converting the codebase to a small PyTorch-like module library (Embedding, Flatten, Sequential) reusable going forward.
PyTorch-ifying the codebase. Before touching the architecture, he refactors the ad hoc layer list into proper module classes — Embedding and Flatten alongside the existing Linear, BatchNorm1d, and Tanh — plus a Sequential container that chains them, mirroring torch.nn's own API. This isn't cosmetic: the forward pass collapses to a single call through the model, exposing a real bug where forgetting to leave BatchNorm in evaluation mode (rather than training mode) caused a single-example forward pass to compute a variance over one number, silently corrupting output — a reminder that BatchNorm's train/eval state is an easy source of hidden bugs.
Baseline: just widening the flat network. Simply increasing context from 3 to 8 characters in the existing one-hidden-layer architecture drops validation loss from 2.10 to 2.02, with no structural change — establishing a baseline against which the hierarchical redesign has to earn its keep.
Building the hierarchical (WaveNet-style) network. Rather than flattening all 8 characters into one long vector immediately, he fuses them progressively: a FlattenConsecutive layer groups every 2 adjacent character embeddings into a single unit, which a linear/BatchNorm/Tanh block processes; this repeats, so 8 characters become 4 pairs, then 2 groups of 4, before a final layer produces the prediction — matching WaveNet's tree-like, dilated causal convolution structure, though implemented purely by reshaping tensors rather than by using real convolutions. He shows PyTorch's matrix multiply already supports this: it operates only on a tensor's last dimension, treating any number of leading dimensions as batch dimensions, so grouping character-pairs into a new "batch-like" dimension is just a view/reshape, not a new operation.
The 3D BatchNorm bug. Feeding the now-3D intermediate tensors (batch × groups × channels) into a BatchNorm1D layer written for 2D input caused it to silently compute separate running statistics per group instead of per channel — reducing only over dimension 0 instead of over both the batch and group dimensions. The fix (reducing jointly over dimensions 0 and 1 when the input is 3D) recovers a proper per-channel normalization and gives a small but real improvement, from 2.029 to 2.022 — and he notes this convention deliberately departs from PyTorch's own BatchNorm1d, which places the channel dimension in the middle rather than last.
Results and what convolutions actually buy you. With the corrected hierarchical architecture and its parameter budget matched to the flat baseline, results are initially a wash (~2.02 either way), but scaling embedding and hidden sizes up (to about 76,000 parameters) pushes validation loss to 1.993. He closes by clarifying that real WaveNet-style convolutions wouldn't change what's computed here at all — they're purely an efficiency mechanism, sliding the same shared filters across a sequence inside fast CUDA kernels and reusing intermediate values shared between overlapping positions, rather than recomputing the whole hierarchy independently at every character position as his implementation does.
- Fusing context characters gradually in a tree, rather than all at once in one layer, improves prediction without adding parameters. Mirroring WaveNet's dilated causal convolutions, he restructures the MLP to fuse two characters per layer, then pairs of those, doubling context each step; holding parameter count fixed at roughly the same budget as a flat network, this hierarchical structure still needed a bug fix before it beat the naive approach, but scaling it up (block size 8, then a bigger hierarchical version) pushed validation loss from 2.10 down to 1.993, versus 2.02 for simply widening a flat one-hidden-layer network to the same context length.
- A subtle BatchNorm bug appeared when the same 2D normalization code was fed 3D tensors, and PyTorch never raised an error. Reducing only over dimension zero left the middle "group" dimension untouched, so the layer silently maintained a separate running mean and variance for each of the 4 groups × 68 channels instead of just 68 channels; averaging over both the batch and group dimensions together fixed it, improving validation loss slightly (2.029 to 2.022) — a reminder that broadcasting can make a wrong shape run without complaint.
- Convolutions are purely an efficiency trick over this same computation, not a different model. Applying the identical stack of layers independently at every position of a name like "DeAndre" (one call per position) computes the same values a convolution would, but redundantly — a convolution instead slides shared weights across the sequence inside fast CUDA kernels and reuses intermediate values that are shared between adjacent output positions, avoiding the recomputation his naive per-position loop would require.