Generative Deep Learning, Chapters 1 to 3

From probability models to variational autoencoders

---

This notebook follows the first three chapters of David Foster, Generative Deep Learning (O'Reilly, first edition). The book's code is written in Keras (GDL_code). Here the same models are rewritten in PyTorch, with the same layer sizes.

Part	Book chapter	The question it answers
1	1 · Generative Modeling	What does it mean for a model to generate data, and why is it hard for images?
2	2 · Deep Learning	How do deep networks learn features from raw pixels?
3	3 · Variational Autoencoders	How do we learn a small space of codes that we can sample new images from?
4		Summary and exercises

Dependencies: numpy, matplotlib and torch. The notebook downloads MNIST (11 MB) and CIFAR-10 (163 MB) into a data/ folder the first time it runs.

Running on Colab. Open the notebook with the badge above. It asks Colab for a GPU runtime. If the first code cell prints that no GPU was found, choose Runtime, then Change runtime type, then a GPU, and run the cells again. Without a GPU the notebook still runs, but slowly: about 12 minutes on a 24-core desktop CPU, and much longer on Colab's CPU runtime. The number of training epochs is set in the first code cell. Lower it for a faster run with somewhat worse models.

Part 0 · Setup

import gzip, pickle, tarfile, time, urllib.request
from pathlib import Path
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import TensorDataset, DataLoader

%matplotlib inline
np.set_printoptions(precision=3, suppress=True, linewidth=100)
torch.manual_seed(0)
device = "cuda" if torch.cuda.is_available() else "cpu"
DATA = Path("data")
EPOCHS = {"mlp": 10, "cnn": 5, "ae": 5, "vae": 10}

C = {"blue": "#2a78d6", "red": "#e34948", "ink": "#0b0b0b", "ink2": "#52514e",
     "muted": "#898781", "grid": "#e1e0d9", "axis": "#c3c2b7", "surface": "#fcfcfb"}
mpl.rcParams.update({
    "figure.facecolor": C["surface"], "axes.facecolor": C["surface"],
    "font.size": 11, "axes.titlesize": 11.5, "axes.titleweight": "bold",
    "axes.labelcolor": C["ink2"], "axes.edgecolor": C["axis"],
    "xtick.color": C["muted"], "ytick.color": C["muted"],
    "legend.frameon": False, "figure.dpi": 110,
})
print("device:", device)
if device == "cpu":
    print("No GPU found. Training will be slow. On Colab: Runtime, Change runtime type, GPU.")

device: cpu
No GPU found. Training will be slow. On Colab: Runtime, Change runtime type, GPU.

The helpers below load the two data sets, train a model and draw rows of images. Every model in the notebook is trained with the same train function.

def download(url, name):
    path = DATA / name
    if not path.exists():
        path.parent.mkdir(parents=True, exist_ok=True)
        print("downloading", url)
        urllib.request.urlretrieve(url, path)
    return path

def load_mnist():
    """Handwritten digits: images of shape (n, 1, 28, 28) with values in [0, 1], and labels."""
    base = "https://ossci-datasets.s3.amazonaws.com/mnist/"
    def read(name, offset):
        with gzip.open(download(base + name, "mnist/" + name)) as f:
            return np.frombuffer(f.read(), np.uint8, offset=offset)
    x_train = read("train-images-idx3-ubyte.gz", 16).reshape(-1, 1, 28, 28)
    x_test = read("t10k-images-idx3-ubyte.gz", 16).reshape(-1, 1, 28, 28)
    y_train = read("train-labels-idx1-ubyte.gz", 8)
    y_test = read("t10k-labels-idx1-ubyte.gz", 8)
    return ((x_train / 255).astype(np.float32), y_train.astype(np.int64),
            (x_test / 255).astype(np.float32), y_test.astype(np.int64))

def load_cifar10():
    """Small color photos: images of shape (n, 3, 32, 32) with values in [0, 1], and labels."""
    archive = download("https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz", "cifar-10-python.tar.gz")
    folder = DATA / "cifar-10-batches-py"
    if not folder.exists():
        with tarfile.open(archive) as t:
            t.extractall(DATA, filter="data")
    def read(name):
        with open(folder / name, "rb") as f:
            batch = pickle.load(f, encoding="bytes")
        return batch[b"data"].reshape(-1, 3, 32, 32), np.array(batch[b"labels"])
    train_parts = [read(f"data_batch_{i}") for i in range(1, 6)]
    x_train = np.concatenate([x for x, _ in train_parts])
    y_train = np.concatenate([y for _, y in train_parts])
    x_test, y_test = read("test_batch")
    return ((x_train / 255).astype(np.float32), y_train.astype(np.int64),
            (x_test / 255).astype(np.float32), y_test.astype(np.int64))

def train(model, x, y, loss_fn, epochs, lr, batch_size=32):
    """Adam on loss_fn(model, x_batch, y_batch). Prints and returns the mean loss of each epoch."""
    loader = DataLoader(TensorDataset(torch.from_numpy(x), torch.from_numpy(y)),
                        batch_size=batch_size, shuffle=True)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    history = []
    for epoch in range(epochs):
        model.train()
        start, total = time.time(), 0.0
        for xb, yb in loader:
            xb, yb = xb.to(device), yb.to(device)
            optimizer.zero_grad()
            loss = loss_fn(model, xb, yb)
            loss.backward()
            optimizer.step()
            total += loss.item() * len(xb)
        history.append(total / len(x))
        print(f"epoch {epoch + 1:2d}/{epochs}   loss {history[-1]:.4f}   ({time.time() - start:.0f} s)")
    return history

@torch.no_grad()
def apply(fn, model, x, batch_size=1000):
    """Run fn on a NumPy array in batches, in evaluation mode, and return a NumPy array."""
    model.eval()
    out = [fn(torch.from_numpy(x[i:i + batch_size]).to(device)).cpu() for i in range(0, len(x), batch_size)]
    return torch.cat(out).numpy()

def show_images(images, titles=None, rows=1, size=1.25, cmap="gray_r"):
    """Draw images of shape (n, 1, h, w) or (n, 3, h, w) in a grid."""
    n = len(images)
    cols = int(np.ceil(n / rows))
    fig, axes = plt.subplots(rows, cols, figsize=(size * cols, size * rows * (1.3 if titles else 1)))
    for i, ax in enumerate(np.ravel(axes)):
        ax.axis("off")
        if i < n:
            img = np.asarray(images[i])
            if img.shape[0] == 3:
                ax.imshow(img.transpose(1, 2, 0))
            else:
                ax.imshow(img[0], cmap=cmap, vmin=0, vmax=1)
            if titles:
                ax.set_title(titles[i], fontsize=8, fontweight="normal", color=C["ink2"])
    plt.tight_layout(); plt.show()

def count_parameters(model):
    return sum(p.numel() for p in model.parameters())

---

Part 1 · Generative modeling

---

1.1 What a generative model is

The book's definition:

A generative model describes how a dataset is generated, in terms of a probabilistic model. By sampling from this model, we are able to generate new data.

The model must be probabilistic. Taking the average of every pixel over a data set of horses gives one blurry image, the same one every time. That is a fixed calculation, not a generative model. A generative model contains randomness, so each sample is different.

A discriminative model does something else. Given paintings labeled "Van Gogh" or "not Van Gogh", it learns $p(y \mid x)$, the probability of a label given an observation. A generative model learns $p(x)$, the probability of the observation itself, and usually needs no labels. The book's point: even a perfect Van Gogh detector has no idea how to paint a Van Gogh.

	Discriminative	Generative
Estimates	$p(y \mid x)$	$p(x)$, or $p(x \mid y)$ if labels are available
Training data	labeled	usually unlabeled
Can produce new observations	no	yes, by sampling

1.2 The framework

The book's framework has four steps:

1. We have a data set of observations $X$.
2. We assume the observations were drawn from an unknown distribution $p_{\text{data}}$.
3. A generative model $p_{\text{model}}$ tries to imitate $p_{\text{data}}$.
4. We sample from $p_{\text{model}}$ to make new observations.

The model is good if it follows two rules.

• Rule 1. Its samples look as if they were drawn from $p_{\text{data}}$.
• Rule 2. Its samples are suitably different from the observations in $X$. A model that only reproduces its training data is not impressive.

The book illustrates this with points on a world map. The data are spread uniformly over land and never fall in the sea. A simple $p_{\text{model}}$ is a rectangle drawn around the points. A sample in the sea breaks Rule 1. A sample right next to a training point breaks Rule 2. A sample on land and away from the data satisfies both.

To fit a model, the book uses four terms.

Term	Meaning
sample space	all values an observation can take
density function $p(x)$	how likely each point of the sample space is; it sums (or integrates) to 1
parametric model $p_\theta(x)$	a family of densities described by a few parameters $\theta$; the rectangle has four, its corners
maximum likelihood	choose $\hat\theta = \arg\max_\theta \sum_{x \in X} \log p_\theta(x)$, the parameters that make the data most probable

Below, the same idea with an arc instead of a continent and a Gaussian instead of a rectangle. For a Gaussian, the maximum likelihood parameters are the sample mean and the sample covariance.

rng = np.random.default_rng(0)
angle = rng.uniform(0, np.pi, 400)
points = np.column_stack([np.cos(angle), np.sin(angle)]) + rng.normal(0, 0.08, (400, 2))

mean = points.mean(axis=0)                    # maximum likelihood estimates
cov = np.cov(points, rowvar=False)
samples = rng.multivariate_normal(mean, cov, size=400)

fig, axes = plt.subplots(1, 2, figsize=(9, 3.6), sharex=True, sharey=True)
for ax, xy, title, color in [(axes[0], points, "Training data", C["blue"]),
                             (axes[1], samples, "Samples from the fitted Gaussian", C["red"])]:
    ax.scatter(xy[:, 0], xy[:, 1], s=8, color=color, alpha=0.7)
    ax.set_title(title, pad=8)
    ax.set_aspect("equal")
plt.tight_layout(); plt.show()

This is the best Gaussian for these data, and many of its samples still land inside the arc, where the data never go. They break Rule 1. Fitting cannot fix this, because no Gaussian has the shape of an arc. The choice of model family limits what the model can learn.

1.3 Naive Bayes: Planet Wrodl and Planet Pixel

The book's first full example is a story called "Hello Wrodl!", a play on "Hello, World!", the first program people write in a new language. You are the new Chief Fashion Officer of Planet Wrodl and must invent new outfits from 50 examples. Each outfit has five features: hairstyle (7 options), hair color (6), glasses (3), clothing type (4) and clothing color (8).

• One probability per combination. There are $7 \cdot 6 \cdot 3 \cdot 4 \cdot 8 = 4{,}032$ combinations. The maximum likelihood estimate of each is its count divided by 50, so every combination that is not in the data gets probability 0. The model can only repeat what it has seen, which breaks Rule 2. Adding 1 to every count fixes that, but then all unseen combinations are equally likely, and the model has learned no structure.
• Naive Bayes assumes the features are independent, so $p(x) = \prod_{k} p(x_k)$. Each feature gets its own small table of probabilities, $23$ parameters in total instead of $4{,}031$. Popular values such as white clothing come up often, and combinations never seen before can still be sampled. On Wrodl this works well, because the five features really are close to independent.

Then the story moves to Planet Pixel, where each outfit is given only as the pixels of a small image. Naive Bayes is applied to each pixel separately, and the samples are noise. We can repeat this on MNIST handwritten digits. Each pixel is treated as an independent coin that is black with probability $p_j$, the fraction of training images in which pixel $j$ is black. The second row uses only images of the digit 7, which is naive Bayes with the label given.

x_mnist, y_mnist, x_mnist_test, y_mnist_test = load_mnist()
black = (x_mnist > 0.5).astype(np.float32)

p_all = black.mean(axis=0)                    # one probability per pixel, shape (1, 28, 28)
p_seven = black[y_mnist == 7].mean(axis=0)

rng = np.random.default_rng(1)
row_all = (rng.random((7, 1, 28, 28)) < p_all).astype(np.float32)
row_seven = (rng.random((7, 1, 28, 28)) < p_seven).astype(np.float32)

show_images(np.concatenate([[p_all], row_all, [p_seven], row_seven]),
            titles=["p, all digits"] + ["sample"] * 7 + ["p, sevens"] + ["sample"] * 7, rows=2)

The model gets right what every image shares: the border is always white. The strokes, which differ from image to image, become scattered dots. Each pixel is drawn without looking at its neighbors, so nothing makes black pixels form a line. Even the model that knows the digit is a 7 cannot draw one.

1.4 The two challenges, and representation learning

The book draws two challenges from Planet Pixel.

1. Features depend on each other. Neighboring pixels are usually similar, and pixels far apart combine into strokes and shapes. Naive Bayes assumes this away.
2. Good observations are a tiny part of a huge sample space. There are $2^{784} \approx 10^{236}$ black-and-white $28 \times 28$ images. Almost none look like digits, so sampling pixels one at a time almost never finds one.

The book's answer to both is deep learning, used for representation learning. Instead of modeling pixels directly, describe each observation by a point $z$ in a low-dimensional latent space and learn a function $f$ that maps $z$ to an image. The book's example is a set of photos of biscuit tins. Two numbers, height and width, describe each tin. With $f$ learned, you can draw a tin that is not in the data, and making a tin taller means adding 1 to the height coordinate, not changing hundreds of pixels by hand.

A network has to discover the latent dimensions by itself. Part 2 covers the networks, and Part 3 learns a latent space for digits.

---

Part 2 · Deep learning

---

2.1 Structured and unstructured data

The book's definition:

Deep learning is a class of machine learning algorithm that uses multiple stacked layers of processing units to learn high-level representations from unstructured data.

Structured data come in columns of features, such as age, income and number of visits. Each column carries information by itself, and models such as logistic regression or random forests work well on it.

Unstructured data, such as images, sound and text, do not. Knowing that pixel 234 is a muddy brown says nothing about whether the photo shows a house or a dog. If the chimney of the house moved to the other side, the same information would sit in completely different pixels. A deep network learns to build informative features, such as "chimney", from the raw pixels itself. This is why deep learning drives generative modeling: the things we want to generate, images, music and text, are unstructured.

The book's data set is CIFAR-10: 60,000 color photos of size $32 \times 32$ in ten classes.

CLASSES = ["airplane", "automobile", "bird", "cat", "deer", "dog", "frog", "horse", "ship", "truck"]
x_cifar, y_cifar, x_cifar_test, y_cifar_test = load_cifar10()
print("train:", x_cifar.shape, "  test:", x_cifar_test.shape)
show_images(x_cifar[:10], titles=[CLASSES[k] for k in y_cifar[:10]])

train: (50000, 3, 32, 32)   test: (10000, 3, 32, 32)

2.2 A multilayer perceptron

The book's first network flattens each image into $3 \cdot 32 \cdot 32 = 3072$ numbers and passes them through dense layers of 200 and 150 units with ReLU activations, then a layer with one output per class. It trains with Adam at a learning rate of 0.0005 and batches of 32.

The book's Keras model ends in a softmax and uses one-hot labels. In PyTorch the model returns raw scores, and F.cross_entropy applies the softmax and takes integer labels. The two set-ups compute the same loss.

def classification_loss(model, x, y):
    return F.cross_entropy(model(x), y)

def test_accuracy(model):
    predicted = apply(model, model, x_cifar_test).argmax(axis=1)
    return (predicted == y_cifar_test).mean()

torch.manual_seed(0)
mlp = nn.Sequential(
    nn.Flatten(),
    nn.Linear(3 * 32 * 32, 200), nn.ReLU(),
    nn.Linear(200, 150), nn.ReLU(),
    nn.Linear(150, 10),
).to(device)
print("parameters:", count_parameters(mlp))
train(mlp, x_cifar, y_cifar, classification_loss, epochs=EPOCHS["mlp"], lr=5e-4)
print(f"MLP test accuracy: {test_accuracy(mlp):.1%}   (book: 49.0% after 10 epochs)")

parameters: 646260

epoch  1/10   loss 1.8474   (4 s)

epoch  2/10   loss 1.6687   (6 s)

epoch  3/10   loss 1.5799   (6 s)

epoch  4/10   loss 1.5217   (6 s)

epoch  5/10   loss 1.4809   (6 s)

epoch  6/10   loss 1.4485   (6 s)

epoch  7/10   loss 1.4209   (6 s)

epoch  8/10   loss 1.3949   (6 s)

epoch  9/10   loss 1.3760   (10 s)

epoch 10/10   loss 1.3509   (10 s)
MLP test accuracy: 49.1%   (book: 49.0% after 10 epochs)

Guessing gives 10%, so the network has learned something. But its first step throws away the layout of the image. To a dense layer, pixel 1 and pixel 2 are no more related than pixel 1 and pixel 3000.

2.3 Convolutions

A convolutional layer keeps the layout. It slides a small grid of weights, a filter, over the image, multiplies it with the patch of pixels under it and adds up the result. The output is large and positive where the patch looks like the filter, and large and negative where it looks like the inverse.

The book shows two hand-made filters on a grayscale image. The first responds to horizontal edges and the second to vertical edges.

gray = torch.from_numpy(x_cifar_test[3].mean(axis=0))[None, None]      # shape (1, 1, 32, 32)
horizontal = torch.tensor([[ 1.0,  1.0,  1.0],
                           [ 0.0,  0.0,  0.0],
                           [-1.0, -1.0, -1.0]])
vertical = torch.tensor([[-1.0, 0.0, 1.0],
                         [-1.0, 0.0, 1.0],
                         [-1.0, 0.0, 1.0]])
filters = torch.stack([horizontal, vertical])[:, None]                   # shape (2, 1, 3, 3)
edges = F.conv2d(gray, filters, padding=1)[0]                            # shape (2, 32, 32)

fig, axes = plt.subplots(1, 3, figsize=(8.4, 3.0))
axes[0].imshow(gray[0, 0], cmap="gray")
axes[0].set_title("Input", pad=8)
for ax, out, title in zip(axes[1:], edges, ["Horizontal edge filter", "Vertical edge filter"]):
    limit = out.abs().max()
    ax.imshow(out, cmap="RdBu_r", vmin=-limit, vmax=limit)
    ax.set_title(title, pad=8)
for ax in axes:
    ax.axis("off")
plt.tight_layout(); plt.show()

In a network the filter weights start random and are learned. A layer with 32 filters produces 32 output images, or channels, and on a color image each filter has three channels of its own. Padding with zeros (padding=1 here, "same" in Keras) keeps the output the same size as the input. A stride of 2 moves the filter two pixels at a time and halves the height and width.

2.4 Batch normalization and dropout

The book's improved network adds two more layer types.

Batch normalization. As training changes the weights of early layers, the inputs to later layers drift, and training can become unstable. The book calls this covariate shift and compares it to carrying a tall pile of books in gusts of wind: each correction makes the pile a little less stable, until it falls. A batch normalization layer subtracts the batch mean of each channel and divides by the batch standard deviation, then applies a learned scale $\gamma$ and shift $\beta$. It also keeps running averages of the mean and standard deviation to use at test time, when there may be no batch. That gives four numbers per channel, two learned and two computed.

Dropout. A student who memorizes the answers to past exam papers fails on new questions. A network can do the same with its training data. A dropout layer sets a random set of units to zero at each training step, so the network cannot rely on any single unit and has to spread what it knows. At test time nothing is dropped. Here model.eval() switches this off.

The activations are leaky ReLUs. They pass a small slope for negative inputs, so a unit never has a gradient of exactly zero.

def conv_block(c_in, c_out, stride):
    return nn.Sequential(nn.Conv2d(c_in, c_out, 3, stride, padding=1),
                         nn.BatchNorm2d(c_out), nn.LeakyReLU())

torch.manual_seed(0)
cnn = nn.Sequential(
    conv_block(3, 32, stride=1),          # 32 x 32 x 32
    conv_block(32, 32, stride=2),         # 32 x 16 x 16
    conv_block(32, 64, stride=1),         # 64 x 16 x 16
    conv_block(64, 64, stride=2),         # 64 x 8 x 8
    nn.Flatten(),
    nn.Linear(64 * 8 * 8, 128), nn.BatchNorm1d(128), nn.LeakyReLU(), nn.Dropout(0.5),
    nn.Linear(128, 10),
).to(device)
print("parameters:", count_parameters(cnn), "  (MLP:", count_parameters(mlp), ")")
train(cnn, x_cifar, y_cifar, classification_loss, epochs=EPOCHS["cnn"], lr=5e-4)
print(f"CNN test accuracy: {test_accuracy(cnn):.1%}   (book: 71.5% after 10 epochs)")

parameters: 591914   (MLP: 646260 )

epoch  1/5   loss 1.3460   (16 s)

epoch  2/5   loss 0.9798   (16 s)

epoch  3/5   loss 0.8387   (14 s)

epoch  4/5   loss 0.7445   (13 s)

epoch  5/5   loss 0.6696   (13 s)

CNN test accuracy: 73.3%   (book: 71.5% after 10 epochs)

The CNN has fewer parameters than the MLP and is much more accurate. Only the architecture changed. Some test predictions:

idx = np.random.default_rng(2).choice(len(x_cifar_test), 10, replace=False)
predicted = apply(cnn, cnn, x_cifar_test[idx]).argmax(axis=1)
show_images(x_cifar_test[idx],
            titles=[f"{CLASSES[p]}\n(true: {CLASSES[t]})" for p, t in zip(predicted, y_cifar_test[idx])])

The book's closing advice for this chapter: it is the layers that are convolutional, not the network. Treat layers as building blocks that fit together when their shapes match, and design the architecture for the task. Part 3 uses the same blocks to build a network that generates images.

---

Part 3 · Variational autoencoders

---

3.1 The art exhibition

The chapter opens with a story. Two brothers, Mr. N. Coder and Mr. D. Coder, run an exhibition of paintings of single digits. The gallery has one wall and no artwork. When a painting arrives, Mr. N. Coder places a marker somewhere on the wall to stand for it and throws the painting away. When a visitor asks to see it, Mr. D. Coder repaints it using only the coordinates of its marker. Visitors who see a bad copy ask for a refund, and over the years the brothers adjust how they place markers and how they repaint to keep refunds low.

Then Mr. N. Coder has an idea: put markers on empty parts of the wall and let his brother paint whatever they stand for. That would be an exhibition of new paintings for free. It goes badly. The paintings lack variety, and some do not look like digits.

Story	Model
a painting	an image $x$
the wall	a two-dimensional latent space
Mr. N. Coder	the encoder, which maps $x$ to a point $z$
Mr. D. Coder	the decoder, which maps $z$ to a reconstruction $\hat x$
refunds	the reconstruction loss

Both networks are trained together to make $\hat x$ close to $x$. This is an autoencoder. The rest of Part 3 repeats the brothers' plan, finds out why it fails, and fixes it.

3.2 An autoencoder for digits

The architecture follows the book. The encoder has four convolutional layers with 32, 64, 64 and 64 filters and strides 1, 2, 2 and 1, then a dense layer down to $z \in \mathbb{R}^{2}$. The decoder mirrors it with transposed convolutions. With a stride of 2, a transposed convolution doubles the height and width instead of halving them, taking $7 \times 7$ back to $28 \times 28$. A sigmoid at the end keeps pixel values between 0 and 1.

For the loss the book prefers the root mean squared error over binary cross-entropy. Binary cross-entropy punishes confident mistakes heavily, which pushes pixel values toward the middle of the range and gives duller images. We use the mean squared error, which has the same minimum as its square root.

def encoder_body():
    """Image (1, 28, 28) to a vector of 64 * 7 * 7 = 3136 features."""
    return nn.Sequential(
        nn.Conv2d(1, 32, 3, 1, padding=1), nn.LeakyReLU(),
        nn.Conv2d(32, 64, 3, 2, padding=1), nn.LeakyReLU(),
        nn.Conv2d(64, 64, 3, 2, padding=1), nn.LeakyReLU(),
        nn.Conv2d(64, 64, 3, 1, padding=1), nn.LeakyReLU(),
        nn.Flatten())

def make_decoder(z_dim):
    """Point z to an image (1, 28, 28)."""
    return nn.Sequential(
        nn.Linear(z_dim, 64 * 7 * 7), nn.Unflatten(1, (64, 7, 7)),
        nn.ConvTranspose2d(64, 64, 3, 1, padding=1), nn.LeakyReLU(),
        nn.ConvTranspose2d(64, 64, 3, 2, padding=1, output_padding=1), nn.LeakyReLU(),
        nn.ConvTranspose2d(64, 32, 3, 2, padding=1, output_padding=1), nn.LeakyReLU(),
        nn.ConvTranspose2d(32, 1, 3, 1, padding=1), nn.Sigmoid())

class Autoencoder(nn.Module):
    def __init__(self, z_dim=2):
        super().__init__()
        self.encoder = nn.Sequential(encoder_body(), nn.Linear(64 * 7 * 7, z_dim))
        self.decoder = make_decoder(z_dim)

    def forward(self, x):
        return self.decoder(self.encoder(x))

def ae_loss(model, x, _labels):
    return F.mse_loss(model(x), x)

torch.manual_seed(0)
ae = Autoencoder().to(device)
print("parameters:", count_parameters(ae))
ae_history = train(ae, x_mnist, y_mnist, ae_loss, epochs=EPOCHS["ae"], lr=5e-4)

parameters: 200963

epoch  1/5   loss 0.0526   (25 s)

epoch  2/5   loss 0.0446   (25 s)

epoch  3/5   loss 0.0423   (26 s)

epoch  4/5   loss 0.0411   (26 s)

epoch  5/5   loss 0.0404   (25 s)

Test digits (top) and their reconstructions (bottom). Each digit passed through just two numbers.

originals = x_mnist_test[:10]
show_images(np.concatenate([originals, apply(ae, ae, originals)]), rows=2)

3.3 The wall, and what is wrong with it

Encode 5,000 test digits and plot where they land, colored by digit. The labels were never used in training, yet similar digits end up together. Then repeat the brothers' plan: pick points at random inside the region the encodings cover and decode them.

def plot_latent(ax, z, labels, title):
    sc = ax.scatter(z[:, 0], z[:, 1], c=labels, cmap="tab10", s=3, alpha=0.7, vmin=-0.5, vmax=9.5)
    ax.set_title(title, pad=8)
    return sc

z_ae = apply(ae.encoder, ae, x_mnist_test[:5000])
lo, hi = z_ae.min(axis=0), z_ae.max(axis=0)
random_z = np.random.default_rng(0).uniform(lo, hi, size=(20, 2)).astype(np.float32)

fig, ax = plt.subplots(figsize=(6.0, 5.0))
sc = plot_latent(ax, z_ae, y_mnist_test[:5000], "Autoencoder latent space")
ax.scatter(random_z[:, 0], random_z[:, 1], s=60, marker="x", color=C["ink"], label="random points")
ax.legend(loc="upper right", fontsize=9)
fig.colorbar(sc, ax=ax, ticks=range(10), label="digit")
plt.tight_layout(); plt.show()

show_images(apply(ae.decoder, ae, random_z), rows=2)

The book makes three observations about this plot, and each explains one way the brothers' exhibition failed.

1. The points are not centered on the origin, and they have no fixed bounds. There is no distribution to sample from, so it is not even clear how to choose a random point.
2. Some digits cover a small area and others a large one. In the book the 1s take up far more room than the 8s. Random points land on the large areas more often, so the generated digits lack variety.
3. There are gaps between the groups. Few images are encoded there, so the decoder has no reason to produce a good digit there.

Behind the third point is a deeper problem: nothing makes the space continuous. If one point decodes to a good 4, nothing guarantees that a point slightly next to it does. The book notes that this gets worse with more latent dimensions, which complex images such as faces need.

3.4 The variational art exhibition

To save the exhibition, Mr. N. Coder brings in his daughter, Epsilon. He still decides where a marker should go, but now he only tells Epsilon his suggestion and how confident he is in it. Epsilon places the marker somewhere near the suggestion, closer when he is more confident. Mr. D. Coder sees only her marker.

Epsilon has two complaints. She dislikes moving the ladder far from the center of the wall, and she dislikes being told too precisely or too vaguely where to put the markers. Her father pays her extra whenever he breaks these rules, an expense the book calls the "kitty-loss (KL) divulgence". He now has to keep both the refunds and these payments low.

A variational autoencoder (Kingma and Welling, 2013) makes exactly these two changes to the autoencoder: one to the encoder and one to the loss.

The encoder. Each image is mapped to a normal distribution in the latent space instead of a single point. The encoder outputs a mean $\mu$ (the suggestion) and a log-variance $\log \sigma^2$ (the confidence) for each latent dimension. The dimensions are assumed independent, so no covariances are needed. The log-variance is used because it can be any real number, while a variance must be positive. A point is then sampled as

$$z = \mu + \sigma \varepsilon, \qquad \sigma = e^{\frac{1}{2}\log \sigma^2}, \qquad \varepsilon \sim \mathcal{N}(0, I)$$

and $\varepsilon$ is Epsilon's random choice. Since the decoder must reconstruct the image from any point near $\mu$, all points in a neighborhood have to decode to similar images. This makes the space continuous.

The loss. The Kullback-Leibler divergence measures how far each encoded distribution is from the standard normal:

$$\mathrm{KL}\big(\mathcal{N}(\mu, \sigma^2) \,\|\, \mathcal{N}(0, 1)\big) = -\frac{1}{2} \sum_{k} \left(1 + \log \sigma_k^2 - \mu_k^2 - \sigma_k^2\right)$$

It is zero only when $\mu = 0$ and $\log \sigma^2 = 0$. It pays for moving away from the center and for confidence that is not just right, Epsilon's two complaints. As a result, the encodings fill the space around the origin with few gaps, and $\mathcal{N}(0, I)$ becomes a sensible distribution to sample from.

The total loss is the reconstruction error times a weight R_LOSS_FACTOR, plus the KL term. The book uses 1000 for digits. If the weight is too large, the KL term has little effect and the problems of Part 3.3 come back. If it is too small, the KL term dominates and the reconstructions are poor.

class VAE(nn.Module):
    def __init__(self, z_dim=2):
        super().__init__()
        self.body = encoder_body()
        self.mu = nn.Linear(64 * 7 * 7, z_dim)
        self.log_var = nn.Linear(64 * 7 * 7, z_dim)
        self.decoder = make_decoder(z_dim)

    def encode(self, x):
        h = self.body(x)
        return self.mu(h), self.log_var(h)

    def forward(self, x):
        mu, log_var = self.encode(x)
        z = mu + torch.exp(0.5 * log_var) * torch.randn_like(mu)
        return self.decoder(z), mu, log_var

R_LOSS_FACTOR = 1000

def vae_loss(model, x, _labels):
    x_hat, mu, log_var = model(x)
    reconstruction = ((x_hat - x) ** 2).mean(dim=(1, 2, 3))
    kl = -0.5 * torch.sum(1 + log_var - mu ** 2 - log_var.exp(), dim=1)
    return (R_LOSS_FACTOR * reconstruction + kl).mean()

torch.manual_seed(0)
vae = VAE().to(device)
vae_history = train(vae, x_mnist, y_mnist, vae_loss, epochs=EPOCHS["vae"], lr=5e-4)

epoch  1/10   loss 54.7539   (26 s)

epoch  2/10   loss 49.1927   (24 s)

epoch  3/10   loss 47.8510   (25 s)

epoch  4/10   loss 46.9806   (25 s)

epoch  5/10   loss 46.4027   (26 s)

epoch  6/10   loss 46.0281   (25 s)

epoch  7/10   loss 45.7120   (24 s)

epoch  8/10   loss 45.4013   (23 s)

epoch  9/10   loss 45.1908   (25 s)

epoch 10/10   loss 44.9826   (24 s)

3.5 The new exhibition

The left plot shows the means $\mu$ of the VAE encodings, with points sampled from $\mathcal{N}(0, I)$ as crosses. The right plot follows the book: each coordinate is passed through the standard normal CDF $\Phi$, which maps $\mathcal{N}(0, 1)$ to the uniform distribution on $[0, 1]$. If the encodings follow $\mathcal{N}(0, I)$, they fill the unit square evenly.

mu_vae = apply(lambda x: vae.encode(x)[0], vae, x_mnist_test[:5000])
normal_z = np.random.default_rng(0).normal(size=(20, 2)).astype(np.float32)
std_normal = torch.distributions.Normal(0, 1)

fig, axes = plt.subplots(1, 2, figsize=(11.5, 5.0))
plot_latent(axes[0], mu_vae, y_mnist_test[:5000], "VAE latent space")
axes[0].scatter(normal_z[:, 0], normal_z[:, 1], s=60, marker="x", color=C["ink"],
                label="samples from N(0, I)")
axes[0].legend(loc="upper right", fontsize=9)
sc = plot_latent(axes[1], std_normal.cdf(torch.from_numpy(mu_vae)).numpy(),
                 y_mnist_test[:5000], "After the normal CDF")
axes[1].set_aspect("equal")
fig.colorbar(sc, ax=axes, ticks=range(10), label="digit")
plt.show()

show_images(apply(vae.decoder, vae, normal_z), rows=2)

Compare this with Part 3.3. The encodings now form one compact cloud around the origin, so sampling from $\mathcal{N}(0, I)$ lands where the decoder has been trained. After the CDF the digits share the square much more evenly than they shared the autoencoder's space, so random samples give a mix of digits. Most of the decoded samples are readable digits. The labels were never used: the VAE organized the digits by itself to keep the reconstruction loss low.

3.6 Walking through the latent space

Because the space is continuous, a regular grid decodes to a map of how digits change from one region to the next. The grid is spaced by quantiles of the normal distribution, so every cell holds about the same probability.

n = 15
q = std_normal.icdf(torch.linspace(0.02, 0.98, n)).numpy()
grid = np.array([[a, b] for b in q[::-1] for a in q], dtype=np.float32)
tiles = apply(vae.decoder, vae, grid).reshape(n, n, 28, 28)
mosaic = tiles.transpose(0, 2, 1, 3).reshape(n * 28, n * 28)

fig, ax = plt.subplots(figsize=(6.5, 6.5))
ax.imshow(mosaic, cmap="gray_r")
ax.set_title("Decoded grid over the VAE latent space", pad=8)
ax.axis("off")
plt.tight_layout(); plt.show()

Walking in a straight line between the encodings of two images, $z = (1 - \alpha)\, z_A + \alpha\, z_B$ for $\alpha$ from 0 to 1, morphs one into the other. The book does this with faces. Below, a test 0 and a test 1 are shown at the ends, and the eight images between them are decoded from points on the line. In two dimensions the line between two digits passes through the regions of other digits, so those appear along the way.

pair = x_mnist_test[[np.argmax(y_mnist_test == 0), np.argmax(y_mnist_test == 1)]]
z_a, z_b = apply(lambda x: vae.encode(x)[0], vae, pair)
path = np.array([(1 - a) * z_a + a * z_b for a in np.linspace(0, 1, 8)], dtype=np.float32)
show_images(np.concatenate([pair[:1], apply(vae.decoder, vae, path), pair[1:]]),
            titles=["input"] + [""] * 8 + ["input"])

3.7 From digits to faces

The book ends the chapter by training a VAE on CelebA, over 200,000 color photos of celebrity faces, each labeled with attributes such as smiling or wearing glasses. The data set is too large to include here. The model follows the digit VAE with these changes:

• three color channels instead of one, and a 200-dimensional latent space, because faces need far more detail than digits;
• batch normalization after each convolution to speed up training, plus dropout;
• a reconstruction weight of 10,000 instead of 1,000.

A 200-dimensional space cannot be plotted, so the book checks each latent dimension separately. If a dimension's distribution differs clearly from $\mathcal{N}(0, 1)$, the reconstruction weight should be lowered, because the KL term is not having enough effect. The goal is good samples, not perfect reconstructions.

Three results come from the same ideas as above.

• New faces. Decoding samples from $\mathcal{N}(0, I)$ gives faces of people who do not exist. The book calls this a giant leap from the naive Bayes model of Part 1: the encoder turns pixels into high-level features, and the decoder turns them back, so dependencies between pixels are handled by the network.
• Latent space arithmetic. Subtract the mean encoding of non-smiling faces from the mean encoding of smiling faces to get a "smile" vector. Then $z_{\text{new}} = z + \alpha \cdot (\text{smile vector})$ decodes to the same face with more or less smile. The same works for blond hair or glasses, and the rest of the face barely changes.
• Morphing. The straight-line walk of Part 3.6 between two faces changes several features at once, and the transition is smooth.

---

Part 4 · Summary and exercises

---

4.1 Summary

Part	Idea	What we saw
1	A generative model estimates $p(x)$ and is sampled to make new data	A Gaussian cannot follow an arc. Naive Bayes on pixels gives scattered dots, because pixels depend on each other.
2	Deep networks learn features from raw pixels	On CIFAR-10 a CNN with batch normalization and dropout beats an MLP with more parameters.
3	An autoencoder learns a latent space, and a VAE makes it easy to sample from	Random points in the autoencoder's space often decode to unclear digits. Samples from $\mathcal{N}(0, I)$ in the VAE's space decode to digits.

The chain from Part 1 to Part 3: modeling pixels one at a time fails, because pixels depend on each other and good images are rare. So we describe each image by a point in a small latent space and let a deep network map points to images. The VAE's random sampling makes the space continuous, and its KL term centers the space on the origin, so a new image is a decoded sample from $\mathcal{N}(0, I)$.

4.2 Exercises

1. Naive Bayes in Part 1.3 multiplies 784 one-pixel models. Write $p(x)$ for a model that instead lets each pixel depend on the pixel to its left. Why is it still far from enough?
2. The MLP of Part 2.2 has 646,260 parameters. Compute this number from the layer sizes.
3. Suppose the encoder sends every image to $\mu = 0$ and $\log \sigma^2 = 0$. What is the KL term, and what can the decoder still do? What stops training from ending up there?
4. Set R_LOSS_FACTOR = 1 and retrain the VAE. Then try 100_000. Before running, predict how the reconstructions and the latent space plot will change.
5. Train the autoencoder and the VAE with z_dim=10. Which plots of Part 3 can you no longer draw directly, and how would you sample from the VAE?

Answers

1. $p(x) = p(x_1) \prod_{j=2}^{784} p(x_j \mid x_{j-1})$, with pixels taken row by row. Each pixel now depends on one neighbor, but strokes also run vertically and across the whole image, and a model of pairs cannot capture the shape of a digit.
2. $3072 \cdot 200 + 200 = 614{,}600$, then $200 \cdot 150 + 150 = 30{,}150$, then $150 \cdot 10 + 10 = 1{,}510$. The total is $646{,}260$.
3. The KL term is 0, its minimum. But every image gets the same distribution, so $z$ carries no information about the image, and the best the decoder can do is output one average image for everything. The reconstruction term stops this, because it grows when images are not reconstructed. The VAE loss is a balance between the two terms, which is why the weight R_LOSS_FACTOR matters.
4. With a factor of 1 the KL term dominates. The encodings are pulled close to $\mathcal{N}(0, I)$ whatever the image, the decoder gets little information, and the reconstructions are poor. With 100,000 the reconstruction term dominates, the KL term has little effect, and the VAE behaves like the plain autoencoder of Part 3.3: the latent space spreads out, with gaps, and is no longer close to $\mathcal{N}(0, I)$. These are the two failure modes the book describes, quoted at the end of Part 3.4.
5. The scatter plots and the decoded grid need a two-dimensional latent space. Sampling works the same way: draw $z \sim \mathcal{N}(0, I_{10})$ and decode it. This is how the face model of Part 3.7 generates images with 200 dimensions.

4.3 Sources

• David Foster, Generative Deep Learning, O'Reilly, 2019, chapters 1 to 3. Code: github.com/davidADSP/GDL_code.
• Kingma and Welling (2013), Auto-Encoding Variational Bayes — arXiv:1312.6114.
• Ioffe and Szegedy (2015), Batch Normalization — arXiv:1502.03167.
• Srivastava et al. (2014), Dropout: A Simple Way to Prevent Neural Networks from Overfitting, JMLR 15.