library(ggplot2)
library(dplyr)
library(tidyr)
library(patchwork)
library(plotly)
if (TORCH_OK) library(torch)

Chapter 14 is published in five parts

Part Sections Content

Part 1: Foundations §14.0-§14.24 Tensors, autograd, the MLP, activations, initialization, normalization, optimizers, the training loop

Part 2: Convolutional Networks and Vision §14.25-§14.44 Convolution arithmetic, receptive fields, residual connections, transfer learning, segmentation, interpretability

Part 3: Sequence Models §14.45-§14.64 RNNs, LSTM/GRU, attention derived, Transformers, tokenization, text generation, neural forecasting

Part 4 (this document) §14.65-§14.84 Autoencoders and PCA, the ELBO, VAEs, GANs, Wasserstein, diffusion, generative evaluation, self-supervised learning

Part 5: Generalization, Uncertainty, and Practice §14.85-§14.104 Double descent, calibration, ensembles, hyperparameter optimization, pruning, robustness

Part 4 reuses machinery built earlier: the U-Net of Part 2, §14.37 returns as the denoising backbone of §14.75, and the masked objective of Part 3, §14.58 generalizes beyond text in §14.78.

Part	Sections	Content
Part 1: Foundations	§14.0-§14.24	Tensors, autograd, the MLP, activations, initialization, normalization, optimizers, the training loop
Part 2: Convolutional Networks and Vision	§14.25-§14.44	Convolution arithmetic, receptive fields, residual connections, transfer learning, segmentation, interpretability
Part 3: Sequence Models	§14.45-§14.64	RNNs, LSTM/GRU, attention derived, Transformers, tokenization, text generation, neural forecasting
Part 4 (this document)	§14.65-§14.84	Autoencoders and PCA, the ELBO, VAEs, GANs, Wasserstein, diffusion, generative evaluation, self-supervised learning
Part 5: Generalization, Uncertainty, and Practice	§14.85-§14.104	Double descent, calibration, ensembles, hyperparameter optimization, pruning, robustness

How this part uses graphics

Two-dimensional figures use ggplot2; the equivalent plot_ly() code follows in a chunk marked eval=FALSE, echo=TRUE. Three-dimensional figures are evaluated plot_ly(), budgeted at six per part.

Heavy training is gated. HEAVY_EVAL defaults to FALSE. Every demonstration below runs on low-dimensional synthetic data chosen so the mechanism is visible, a two-dimensional latent space can be plotted; a 512-dimensional one cannot.

Learning objectives

After completing Part 4 you will be able to:

Distinguish discriminative from generative modelling, and classify generative models by how they treat the likelihood.
State precisely what a linear autoencoder recovers, and why it is the principal subspace rather than the principal components.
Derive the evidence lower bound, identify the gap between it and the log likelihood, and compute the Gaussian KL term in closed form.
Explain why the reparameterization trick is necessary and measure the variance it saves against the score-function estimator.
Diagnose posterior collapse and explain the rate-distortion trade that \(\beta\) controls.
Derive the optimal GAN discriminator and show that the generator minimizes a Jensen-Shannon divergence.
Explain why that divergence gives no gradient under disjoint support, and what the Wasserstein formulation replaces it with.
Write down the diffusion forward process in closed form and connect the denoising objective to both the ELBO and score matching.
Criticize FID and Inception Score, and explain why a memorizing model scores perfectly.
Implement a contrastive objective and state what the InfoNCE bound does and does not establish.

Estimated time: 12-16 hours including exercises.

PART I: LATENT-VARIABLE MODELS

14.65 Discriminative and generative modelling

Every model in Parts 1-3 estimated a conditional: \(p(y\mid x)\) for classification, \(p(x_t\mid x_{<t})\) for language. Generative modelling targets the data distribution itself.

\[ \begin{aligned} \textbf{Discriminative: }&\quad p_\theta(y\mid x) &&\text{decision boundaries}\\ \textbf{Generative: }&\quad p_\theta(x)\ \text{or}\ p_\theta(x,y) &&\text{the distribution, from which one can sample} \end{aligned} \]

The generative problem is strictly harder. A classifier need only get the boundary right; a generative model must account for all the structure in \(x\), including everything irrelevant to any particular label.

The families differ in how they handle the likelihood, and that difference drives everything else about them.

Family	Likelihood	Sampling	Training signal
Autoregressive (§14.56)	Exact, by the chain rule	Sequential, \(O(T)\)	Direct MLE
Normalizing flows	Exact, via change of variables	One pass	Direct MLE
VAE (§14.69)	Lower bound (ELBO)	One pass	Variational
GAN (§14.72)	Implicit, never evaluated	One pass	Adversarial
Diffusion (§14.75)	Lower bound, in practice a simple regression	Iterative, \(O(T)\)	Denoising
Energy-based	Unnormalized; \(Z\) intractable	MCMC	Contrastive

The intractable object is almost always the normalizing constant. Writing \(p_\theta(x)=\tilde p_\theta(x)/Z_\theta\) with \(Z_\theta=\int\tilde p_\theta(x)\,dx\), maximum likelihood needs \(Z_\theta\), and for any interesting \(\tilde p_\theta\) in high dimensions that integral has no closed form.

The four responses in the table are four ways around it. Autoregressive models factor \(p(x)=\prod_t p(x_t\mid x_{<t})\) so each factor normalizes over a single variable. Flows constrain the architecture to be invertible with a tractable Jacobian determinant. VAEs optimize a bound instead of the likelihood. GANs abandon the likelihood entirely and compare samples. Diffusion models start as a bound and reduce, remarkably, to least squares.

14.66 Autoencoders

An autoencoder learns a compressed code by reconstructing its own input:

\[\boxed{\;\min_{\theta,\phi}\ \mathbb E_{x\sim p_{\text{data}}}\big\|x-g_\phi\big(f_\theta(x)\big)\big\|^2\;}\]

with encoder \(f_\theta:\mathbb R^D\to\mathbb R^d\) and decoder \(g_\phi:\mathbb R^d\to\mathbb R^D\).

The architecture supplies no supervision, so the constraint is what forces learning. An undercomplete autoencoder (\(d<D\)) cannot copy its input and must discard something; the question is what. An overcomplete one (\(d\ge D\)) can learn the identity map and will, unless regularized.

autoencoder <- nn_module("AE",
  initialize = function(D, d, hidden = 64, linear = FALSE) {
    self$enc <- if (linear) nn_linear(D, d, bias = FALSE) else
      nn_sequential(nn_linear(D, hidden), nn_relu(), nn_linear(hidden, d))
    self$dec <- if (linear) nn_linear(d, D, bias = FALSE) else
      nn_sequential(nn_linear(d, hidden), nn_relu(), nn_linear(hidden, D))
  },
  forward = function(x) self$dec(self$enc(x)),
  encode  = function(x) self$enc(x))

torch_manual_seed(11)
ae_demo <- autoencoder(D = 20, d = 3)
c(parameters = n_par(ae_demo),
  bottleneck = 3,
  compression_ratio = round(20 / 3, 2),
  reconstruction_shape = paste(dim(with_no_grad(ae_demo(torch_randn(8, 20)))),
                               collapse = " x "))

#>           parameters           bottleneck    compression_ratio 
#>               "3095"                  "3"               "6.67" 
#> reconstruction_shape 
#>             "8 x 20"

14.67 Linear autoencoders and PCA

Common misconception: “a linear autoencoder learns PCA.” It learns the principal subspace, and that is a weaker and more precise statement.

For a linear encoder \(W_e\) and decoder \(W_d\) with squared loss, the global minimum of \(\|X-XW_e^\top W_d^\top\|_F^2\) is attained when \(W_d W_e\) is the orthogonal projection onto the span of the top-\(d\) eigenvectors of the covariance (Baldi & Hornik, 1989). But the individual columns of \(W_d\) are determined only up to an arbitrary invertible \(d\times d\) transformation: if \((W_e, W_d)\) is optimal, so is \((AW_e, W_dA^{-1})\) for any invertible \(A\).

Three consequences. The learned directions are not orthonormal in general. They are not ordered by variance, there is no “first” component. And they are not unique, so two runs from different seeds give different bases for the same subspace.

Verify subspace agreement with principal angles, not by comparing vectors. Comparing columns to PCA loadings will show disagreement even at a perfect solution.

set.seed(21); torch_manual_seed(21)
n_s <- 800; D_s <- 12; d_s <- 3
# Data with genuine rank-3 structure plus isotropic noise
B_true <- matrix(rnorm(D_s * d_s), D_s, d_s)
Z_true <- matrix(rnorm(n_s * d_s), n_s, d_s) %*% diag(c(4, 2, 1))
X_s <- scale(Z_true %*% t(B_true) + matrix(rnorm(n_s * D_s, sd = 0.4), n_s, D_s),
             center = TRUE, scale = FALSE)
Xt <- torch_tensor(X_s)

fit_linear_ae <- function(seed) {
  torch_manual_seed(seed)
  ae <- autoencoder(D_s, d_s, linear = TRUE)
  opt <- optim_adam(ae$parameters, lr = 5e-3)
  for (e in 1:2500) { opt$zero_grad(); nnf_mse_loss(ae(Xt), Xt)$backward(); opt$step() }
  list(W_e = as.matrix(ae$enc$weight),                 # d x D
       loss = as.numeric(with_no_grad(nnf_mse_loss(ae(Xt), Xt))))
}
ae1 <- fit_linear_ae(31); ae2 <- fit_linear_ae(37)
pca <- prcomp(X_s, center = FALSE)
V <- pca$rotation[, 1:d_s]                              # D x d

# PCA reconstruction error at the same rank, for reference
pca_mse <- mean((X_s - X_s %*% V %*% t(V))^2)
c(ae_seed31_mse = signif(ae1$loss, 5), ae_seed37_mse = signif(ae2$loss, 5),
  pca_rank3_mse = signif(pca_mse, 5))

#> ae_seed31_mse ae_seed37_mse pca_rank3_mse 
#>       0.11954       0.11972       0.11891

The reconstruction errors agree to several digits, both autoencoders found a global optimum. Now compare what they found.

principal_angles <- function(A, B) {
  qa <- qr.Q(qr(A)); qb <- qr.Q(qr(B))
  sv <- svd(t(qa) %*% qb)$d
  acos(pmin(pmax(sv, -1), 1)) * 180 / pi
}
Wa <- t(ae1$W_e); Wb <- t(ae2$W_e)                      # D x d each

data.frame(
  comparison = c("AE(seed 31) vs PCA", "AE(seed 37) vs PCA",
                 "AE(seed 31) vs AE(seed 37)"),
  max_principal_angle_deg = signif(c(max(principal_angles(Wa, V)),
                                     max(principal_angles(Wb, V)),
                                     max(principal_angles(Wa, Wb))), 4),
  same_subspace = c(max(principal_angles(Wa, V)) < 2,
                    max(principal_angles(Wb, V)) < 2,
                    max(principal_angles(Wa, Wb)) < 2))

# Now the direct comparison that FAILS, and why
cos_sim <- abs(cor(Wa[, 1], V[, 1]))
c(cosine_between_AE_col1_and_PC1 = signif(cos_sim, 4),
  AE_columns_orthonormal = isTRUE(all.equal(t(Wa) %*% Wa,
                                            diag(d_s), tolerance = 0.05)),
  note = "the subspaces agree; the individual directions need not")

#>                            cosine_between_AE_col1_and_PC1 
#>                                                  "0.9138" 
#>                                    AE_columns_orthonormal 
#>                                                   "FALSE" 
#>                                                      note 
#> "the subspaces agree; the individual directions need not"

Principal angles near zero confirm the subspaces coincide, while the column-by-column cosine does not, exactly the ambiguity the theorem predicts. The autoencoder’s basis is a rotated, rescaled version of PCA’s.

# What nonlinearity buys: data on a curved manifold that no linear subspace fits
set.seed(41); torch_manual_seed(41)
n_m <- 1000
t_par <- runif(n_m, 0, 3*pi)
X_curve <- cbind(t_par * cos(t_par), t_par * sin(t_par)) / 6 +
           matrix(rnorm(n_m * 2, sd = 0.03), n_m, 2)
Xc <- torch_tensor(X_curve)

fit_ae_curve <- function(linear) {
  torch_manual_seed(43)
  ae <- autoencoder(2, 1, hidden = 64, linear = linear)
  opt <- optim_adam(ae$parameters, lr = 3e-3)
  for (e in 1:3000) { opt$zero_grad(); nnf_mse_loss(ae(Xc), Xc)$backward(); opt$step() }
  list(recon = as.matrix(with_no_grad(ae(Xc))),
       mse = as.numeric(with_no_grad(nnf_mse_loss(ae(Xc), Xc))))
}
lin <- fit_ae_curve(TRUE); nonlin <- fit_ae_curve(FALSE)
c(linear_AE_mse = signif(lin$mse, 4), nonlinear_AE_mse = signif(nonlin$mse, 4),
  improvement_factor = round(lin$mse / nonlin$mse, 1))

#>      linear_AE_mse   nonlinear_AE_mse improvement_factor 
#>           0.173200           0.001809          95.700000

bind_rows(
  data.frame(x = X_curve[,1], y = X_curve[,2], panel = "Data (1-D manifold in 2-D)"),
  data.frame(x = lin$recon[,1], y = lin$recon[,2], panel = "Linear AE, d = 1"),
  data.frame(x = nonlin$recon[,1], y = nonlin$recon[,2], panel = "Nonlinear AE, d = 1")) |>
  mutate(panel = factor(panel, levels = unique(panel))) |>
  ggplot(aes(x, y)) +
  geom_point(size = 0.5, alpha = 0.5, color = "steelblue") +
  facet_wrap(~ panel) + coord_fixed() +
  labs(title = "A one-dimensional bottleneck on a curved manifold",
       subtitle = "The linear autoencoder can only project onto a line; the nonlinear one follows the spiral",
       x = NULL, y = NULL) +
  theme_dspa(10)

Nonlinearity is what turns “principal subspace” into “manifold learning.” The spiral is intrinsically one-dimensional, and only the nonlinear encoder can find that coordinate.

14.68 Denoising and sparse autoencoders

An overcomplete autoencoder learns the identity unless something prevents it. Two standard preventions.

\[ \begin{aligned} \textbf{Denoising: }&\quad \min\ \mathbb E_{x}\,\mathbb E_{\tilde x\sim q(\tilde x\mid x)}\big\|x-g\big(f(\tilde x)\big)\big\|^2\\ \textbf{Sparse: }&\quad \min\ \mathbb E_x\Big[\|x-g(f(x))\|^2+\lambda\big\|f(x)\big\|_1\Big] \end{aligned} \]

The sparse penalty is the \(\ell_1\) regularizer of Chapter 11, §11.7, applied to activations rather than weights, and its proximal operator is the same soft threshold (Chapter 13, §13.18).

A denoising autoencoder learns the score of the data distribution. For Gaussian corruption \(\tilde x=x+\sigma\varepsilon\) with small \(\sigma\), the optimal reconstruction satisfies \[g^\star\big(f^\star(\tilde x)\big)-\tilde x\ \approx\ \sigma^2\,\nabla_{\tilde x}\log p(\tilde x)\] (Alain & Bengio, 2014).

The residual points uphill in density, toward where the data are. This is not an incidental property: it is exactly the quantity a diffusion model learns (§14.75), and it is why “train a network to remove noise” turns out to be a route to generative modelling rather than merely a denoising technique.

# The DAE residual, plotted against the analytic score of a known density
set.seed(51); torch_manual_seed(51)
# A 1-D mixture whose score we can compute exactly
mix_p <- function(x) 0.6*dnorm(x, -1.5, 0.5) + 0.4*dnorm(x, 1.8, 0.7)
mix_score <- function(x) {
  a <- 0.6*dnorm(x, -1.5, 0.5); b <- 0.4*dnorm(x, 1.8, 0.7)
  (a * (-(x+1.5)/0.25) + b * (-(x-1.8)/0.49)) / (a + b)
}
x_data <- c(rnorm(3000, -1.5, 0.5), rnorm(2000, 1.8, 0.7))
sigma_c <- 0.35

dae <- nn_sequential(nn_linear(1, 128), nn_tanh(), nn_linear(128, 128),
                     nn_tanh(), nn_linear(128, 1))
opt_d <- optim_adam(dae$parameters, lr = 3e-3)
xt_d <- torch_tensor(matrix(x_data, ncol = 1))
for (e in 1:2500) {
  opt_d$zero_grad()
  noisy <- xt_d + torch_randn_like(xt_d) * sigma_c
  nnf_mse_loss(dae(noisy), xt_d)$backward(); opt_d$step()
}

grid_x <- seq(-4, 4.5, length.out = 300)
resid <- as.numeric(with_no_grad(dae(torch_tensor(matrix(grid_x, ncol = 1))))) - grid_x
data.frame(x = grid_x,
           `DAE residual / sigma^2` = resid / sigma_c^2,
           `analytic score` = mix_score(grid_x), check.names = FALSE) |>
  pivot_longer(-x, names_to = "quantity", values_to = "v") |>
  ggplot(aes(x, v, color = quantity)) +
  geom_hline(yintercept = 0, color = "grey70") +
  geom_line(linewidth = 0.9) +
  coord_cartesian(ylim = c(-12, 12)) +
  scale_color_manual(values = c("#3B7DD8", "#D8433B")) +
  labs(title = "A denoising autoencoder recovers the score of the data density",
       subtitle = expression("Residual "*(g(f(tilde(x)))-tilde(x))/sigma^2*" against "*
                             nabla*log*p(x)*", never shown to the network"),
       x = "x", y = NULL, color = NULL) +
  theme_dspa()

keep <- grid_x > -3 & grid_x < 3.6
c(correlation_with_true_score =
    signif(cor(resid[keep]/sigma_c^2, mix_score(grid_x[keep])), 4))

#> correlation_with_true_score 
#>                      0.9839

# --- Interactive equivalent ------------------------------------------------
plot_ly(x = grid_x, y = resid/sigma_c^2, type = "scatter", mode = "lines",
        name = "DAE residual / sigma^2") |>
  add_lines(y = mix_score(grid_x), name = "Analytic score") |>
  layout(title = "Denoising residual and the score function",
         xaxis = list(title = "x"), yaxis = list(title = "score", range = c(-12, 12)))

The network was trained only to remove noise. It recovered \(\nabla\log p(x)\), a quantity never supplied to it, including the sign change at the density’s saddle between the two modes.

14.69 The variational autoencoder

In standard autoencoders and variational autoencoders (VAEs), “codes”, or latent codes, are representations of compressed, lower-dimensional numerical vectors that capture the salient features, or essence, of an input data point, such as an image, audio clip, text, etc. In practice, the code is a set of coordinates, or a recipe, in a hidden “latent space.” For instance, feeding an image of a handwritten digit \(7\) into an autoencoder, the model reduces all pixel intensities down to a handful of numbers, e.g., a vector of \(32\) numbers, which represent features like loop size, stroke thickness, and tilt.

In a standard autoencoder, every unique input maps to one fixed, exact point in that latent space (a deterministic code).
While in a variational autoencoder (VAE), an input maps to a probability distribution, described by its moments, e.g., mean and variance, rather than a single fixed point.

Because a VAE learns a smooth probability distribution over these codes, rather than memorizing isolated points, the VAE can randomly sample a point from that distribution (\(z \sim p(z)\)), pass it through the decoder, and successfully synthesize entirely new, realistic data (\(x\)) that the model has never seen before.

An autoencoder learns a deterministic code. A variational autoencoder learns a distribution over codes, which is what makes it a generative model: sample \(z\sim p(z)\), decode, and you have a new \(x\).

The model is \[p_\theta(x)=\int p_\theta(x\mid z)\,p(z)\,dz,\qquad p(z)=\mathcal N(\mathbf 0,I),\] and the integral is intractable, the difficulty of §14.65 in its usual form.

14.69.1 Deriving the ELBO

For a latent-variable model \[p_\theta(x,z)=p_\theta(x\mid z)p(z),\] the marginal log-likelihood \(\log p_\theta(x)\) is usually intractable.
The evidence lower bound (ELBO) is a tractable lower bound used as a training objective.

Let \(q_\phi(z\mid x)\) be the encoder. The KL gap between the encoder and the true posterior is

\[ D_{\mathrm{KL}}\big(q_\phi(z\mid x)\,\big\|\,p_\theta(z\mid x)\big) = \mathbb{E}_{q_\phi(z\mid x)} \left[ \log \frac{q_\phi(z\mid x)}{p_\theta(z\mid x)} \right]. \]

By Bayes’ rule,

\[ p_\theta(z\mid x) = \frac{p_\theta(x,z)}{p_\theta(x)} = \frac{p_\theta(x\mid z)p(z)}{p_\theta(x)}. \]

Substitute this into the KL divergence

\[ D_{\mathrm{KL}}\big(q_\phi(z\mid x)\,\big\|\,p_\theta(z\mid x)\big) = \mathbb{E}_{q_\phi(z\mid x)} \left[ \log q_\phi(z\mid x) + \log p_\theta(x) - \log\big(p_\theta(x\mid z)p(z)\big) \right]. \]

Since \(\log p_\theta(x)\) does not depend on \(z\), its expectation under \(q_\phi(z\mid x)\) is just itself. Rearranging gives the exact decomposition

\[ \log p_\theta(x) = \underbrace{ \mathbb{E}_{q_\phi(z\mid x)}\big[\log p_\theta(x\mid z)\big] - D_{\mathrm{KL}}\big(q_\phi(z\mid x)\,\big\|\,p(z)\big) }_{\mathcal{L}_{\text{ELBO}}(\theta,\phi;x)} + \underbrace{ D_{\mathrm{KL}}\big(q_\phi(z\mid x)\,\big\|\,p_\theta(z\mid x)\big) }_{\text{encoder gap } \ge 0}. \]

Thus,

\[ \boxed{ \log p_\theta(x) = \mathcal{L}_{\text{ELBO}}(\theta,\phi;x) + D_{\mathrm{KL}}\big(q_\phi(z\mid x)\,\big\|\,p_\theta(z\mid x)\big) } \]

and because the KL gap is non-negative,

\[ \mathcal{L}_{\text{ELBO}}(\theta,\phi;x) \le \log p_\theta(x). \]

The ELBO is therefore a lower bound on the model evidence.

The ELBO itself expands as

\[ \mathcal{L}_{\text{ELBO}}(\theta,\phi;x) = \underbrace{ \mathbb{E}_{q_\phi(z\mid x)}\big[\log p_\theta(x\mid z)\big] }_{\text{reconstruction term}} - \underbrace{ D_{\mathrm{KL}}\big(q_\phi(z\mid x)\,\big\|\,p(z)\big) }_{\text{regularization}}. \]

Common misconception. The ELBO is not the model’s log-likelihood.
It is a lower bound, and the gap is exactly the encoder’s posterior mismatch.
The bound is tight when \(q_\phi(z\mid x)=p_\theta(z\mid x)\).
Comparing ELBOs across models with different encoder families compares bounds of different tightness.
Importance-weighted bounds, such as IWAE, give tighter estimates at extra computational cost.

The same bound can also be derived with Jensen’s inequality, but the direct decomposition above is usually clearer.

Closed-form KL term.
For

\[ q_\phi(z\mid x) = \mathcal{N}\big(\mu_\phi(x),\,\operatorname{diag}(\sigma_\phi^2(x))\big) \quad\text{and}\quad p(z) = \mathcal{N}(\mathbf{0}, I), \]

the KL term is

\[ D_{\mathrm{KL}}\big(q_\phi(z\mid x)\,\big\|\,p(z)\big) = \frac{1}{2} \sum_{j=1}^{d} \Big( \mu_j^2 + \sigma_j^2 - \log \sigma_j^2 - 1 \Big). \]

This KL term requires no sampling, only the reconstruction term does.

kl_gaussian <- function(mu, logvar)
  0.5 * torch_sum(mu^2 + torch_exp(logvar) - logvar - 1, dim = 2)

# Verify against Monte Carlo
torch_manual_seed(61)
mu_t <- torch_randn(1, 5) * 0.8; logvar_t <- torch_randn(1, 5) * 0.4
mc <- with_no_grad({
  eps <- torch_randn(200000, 5)
  z <- mu_t + torch_exp(0.5 * logvar_t) * eps
  log_q <- torch_sum(-0.5*(logvar_t + log(2*pi) + (z - mu_t)^2/torch_exp(logvar_t)), dim=2)
  log_p <- torch_sum(-0.5*(log(2*pi) + z^2), dim = 2)
  torch_mean(log_q - log_p)
})
c(closed_form = signif(as.numeric(kl_gaussian(mu_t, logvar_t)), 6),
  monte_carlo_200k = signif(as.numeric(mc), 6),
  note = "the KL needs no sampling; only the reconstruction term does")

#>                                                   closed_form 
#>                                                    "0.480784" 
#>                                              monte_carlo_200k 
#>                                                    "0.479612" 
#>                                                          note 
#> "the KL needs no sampling; only the reconstruction term does"

14.70 The reparameterization trick

The score function (REINFORCE) moves the gradient inside the expectation using the log-derivative trick. While general, it suffers from extremely high variance, making VAE training slow or unstable. The reparameterization trick bypasses this variance problem by pushing the parameter \(\phi\) out of the distribution and into a deterministic function, which allows gradients to flow cleanly.

The reconstruction term is \(\mathbb E_{q_\phi(z\mid x)}[\log p_\theta(x\mid z)]\), and \(\phi\) appears in the distribution being averaged over, so gradient and expectation cannot simply be exchanged.

Two estimators exist.

\[ \begin{aligned} \textbf{Score function (REINFORCE): }&\quad \nabla_\phi\,\mathbb E_{q_\phi}[f(z)]=\mathbb E_{q_\phi}\big[f(z)\,\nabla_\phi\log q_\phi(z)\big]\\[2mm] \textbf{Reparameterization: }&\quad z=\mu_\phi+\sigma_\phi\odot\varepsilon,\ \varepsilon\sim\mathcal N(\mathbf 0,I) \ \Longrightarrow\ \nabla_\phi\,\mathbb E_{\varepsilon}\big[f(\mu_\phi+\sigma_\phi\odot\varepsilon)\big] =\mathbb E_{\varepsilon}\big[\nabla_\phi f\big] \end{aligned} \]

Both are unbiased. The reparameterized form moves \(\phi\) out of the distribution and into the function, so the randomness \(\varepsilon\) no longer depends on the parameters and the gradient passes straight through.

Common misconception: “reparameterization is a technical trick to make sampling differentiable.” Both estimators are differentiable and both are unbiased. What separates them is variance, and the difference is large enough to decide whether training works.

The score-function estimator multiplies \(f(z)\), the whole loss, by \(\nabla_\phi\log q_\phi(z)\), so it uses no information about how \(f\) varies with \(z\); it only correlates the loss value with the score. Its variance grows roughly linearly in the latent dimension. The reparameterized estimator propagates \(\nabla_z f\) through the sampling path, using the local geometry of \(f\), and its variance is far smaller and far less dimension-sensitive.

The practical consequence: reparameterization works with one Monte Carlo sample per data point, which is what makes VAE training as cheap as autoencoder training. The score-function estimator needs variance reduction (baselines, control variates) and many samples, and it is what you must fall back on when \(z\) is discrete and reparameterization is unavailable.

# Both estimators on the same objective, with their variances measured
grad_variance <- function(d, n_samples, n_reps = 400, seed = 71) {
  torch_manual_seed(seed); set.seed(seed)
  mu0 <- torch_randn(d) * 0.5
  A <- torch_randn(d, d) / sqrt(d)
  f <- function(z) torch_sum((z$matmul(A))^2, dim = 2)     # a smooth loss

  reparam <- vapply(seq_len(n_reps), function(r) {
    mu <- mu0$clone()$requires_grad_(TRUE)
    eps <- torch_randn(n_samples, d)
    z <- mu$unsqueeze(1) + eps                             # sigma = 1 for clarity
    torch_mean(f(z))$backward()
    as.numeric(mu$grad[1])
  }, numeric(1))

  score_fn <- vapply(seq_len(n_reps), function(r) {
    mu <- mu0$clone()$requires_grad_(TRUE)
    z <- with_no_grad(mu$unsqueeze(1) + torch_randn(n_samples, d))
    logq <- torch_sum(-0.5 * (z - mu$unsqueeze(1))^2, dim = 2)
    fv <- with_no_grad(f(z))
    torch_mean(fv * logq)$backward()
    as.numeric(mu$grad[1])
  }, numeric(1))

  c(d = d, n_samples = n_samples,
    reparam_sd = sd(reparam), score_fn_sd = sd(score_fn),
    variance_ratio = var(score_fn) / var(reparam),
    means_agree = abs(mean(reparam) - mean(score_fn)) <
      3 * (sd(reparam)/sqrt(n_reps) + sd(score_fn)/sqrt(n_reps)))
}
as.data.frame(do.call(rbind, lapply(c(2, 8, 32, 128), grad_variance,
                                    n_samples = 1))) |>
  mutate(across(c(reparam_sd, score_fn_sd, variance_ratio), \(z) signif(z, 4)))

The means_agree column confirms both estimators are unbiased; the variance_ratio column shows the score-function estimator’s variance growing with dimension while the reparameterized one stays controlled.

d_grid <- c(2, 4, 8, 16, 32, 64)
s_grid <- c(1, 2, 4, 8, 16)
Zratio <- outer(d_grid, s_grid, Vectorize(function(d, s)
  log10(grad_variance(d, s, n_reps = 120)[["variance_ratio"]])))

plot_ly(x = s_grid, y = d_grid, z = Zratio, type = "surface",
        colorscale = "Inferno", reversescale = TRUE,
        colorbar = list(title = "log10 variance\nratio")) |>
  layout(title = "Variance penalty of the score-function estimator, relative to reparameterization",
         scene = list(xaxis = list(title = "Monte Carlo samples", type = "log"),
                      yaxis = list(title = "Latent dimension", type = "log"),
                      zaxis = list(title = "log10 var(score) / var(reparam)")))

Rotate along the dimension axis: the penalty climbs steadily, because the score-function estimator’s variance scales with \(d\) while the reparameterized one’s does not. Along the samples axis both fall as \(1/n\), so averaging reduces the ratio’s absolute size without closing the gap, the surface descends but never reaches zero.

vae <- nn_module("VAE",
  initialize = function(D, d, hidden = 128) {
    self$d <- d
    self$enc <- nn_sequential(nn_linear(D, hidden), nn_relu(),
                              nn_linear(hidden, hidden), nn_relu())
    self$to_mu     <- nn_linear(hidden, d)
    self$to_logvar <- nn_linear(hidden, d)
    self$dec <- nn_sequential(nn_linear(d, hidden), nn_relu(),
                              nn_linear(hidden, hidden), nn_relu(),
                              nn_linear(hidden, D))
  },
  encode = function(x) { h <- self$enc(x); list(mu = self$to_mu(h),
                                                logvar = self$to_logvar(h)) },
  reparameterize = function(mu, logvar) {
    # z = mu + sigma * eps, with eps INDEPENDENT of the parameters
    mu + torch_exp(0.5 * logvar) * torch_randn_like(mu)
  },
  forward = function(x) {
    e <- self$encode(x)
    z <- self$reparameterize(e$mu, e$logvar)
    list(recon = self$dec(z), mu = e$mu, logvar = e$logvar, z = z)
  })

vae_loss <- function(out, x, beta = 1) {
  rec <- torch_sum((out$recon - x)^2, dim = 2)     # Gaussian log-likelihood up to a constant
  kl  <- kl_gaussian(out$mu, out$logvar)
  list(total = torch_mean(rec + beta * kl),
       recon = torch_mean(rec), kl = torch_mean(kl))
}

# Two interleaved half-moons: 2-D data, 2-D latent, so both can be plotted
set.seed(81); torch_manual_seed(81)
n_v <- 2000
th1 <- runif(n_v/2, 0, pi); th2 <- runif(n_v/2, 0, pi)
X_moon <- rbind(cbind(cos(th1), sin(th1)),
                cbind(1 - cos(th2), 0.5 - sin(th2)))
X_moon <- scale(X_moon + matrix(rnorm(n_v*2, sd = 0.06), n_v, 2))
lab_moon <- rep(c("A", "B"), each = n_v/2)
Xv <- torch_tensor(X_moon)

m_vae <- vae(D = 2, d = 2, hidden = 128)
opt_v <- optim_adam(m_vae$parameters, lr = 2e-3)
hist_v <- data.frame()
for (e in 1:1500) {
  opt_v$zero_grad()
  L <- vae_loss(m_vae(Xv), Xv, beta = 1)
  L$total$backward(); opt_v$step()
  if (e %% 50 == 0)
    hist_v <- rbind(hist_v, data.frame(epoch = e,
      recon = as.numeric(L$recon), kl = as.numeric(L$kl),
      elbo = -as.numeric(L$total)))
}
tail(hist_v, 3) |> mutate(across(everything(), \(z) signif(z, 5)))

lat <- as.matrix(with_no_grad(m_vae$encode(Xv)$mu))
gen <- as.matrix(with_no_grad(m_vae$dec(torch_randn(1200, 2))))

p_lat <- ggplot(data.frame(z1 = lat[,1], z2 = lat[,2], g = lab_moon),
                aes(z1, z2, color = g)) +
  geom_point(size = 0.5, alpha = 0.5) +
  scale_color_manual(values = c("#3B7DD8", "#D8433B"), guide = "none") +
  labs(title = "Latent means q(z|x)", x = "z1", y = "z2") +
  theme_dspa(9) +
  theme(aspect.ratio = 1)          # square panel, but no scale locking

p_gen <- bind_rows(
  data.frame(x = X_moon[,1], y = X_moon[,2], src = "data"),
  data.frame(x = gen[,1], y = gen[,2], src = "generated")) |>
  ggplot(aes(x, y, color = src)) +
  geom_point(size = 0.5, alpha = 0.45) +
  scale_color_manual(values = c(data = "grey55", generated = "#7FB069"),
                     name = NULL) +
  labs(title = "Samples from p(z) decoded", x = "x", y = "y") +
  theme_dspa(9) +
  theme(aspect.ratio = 1)
p_lat | p_gen

The latent means occupy a roughly standard-normal cloud, that is the KL term doing its work, and decoding fresh \(z\sim\mathcal N(0,I)\) produces points that follow the moons. A plain autoencoder cannot do the second thing, because nothing constrained its code space to be sampleable.

14.71 Posterior collapse and the rate-distortion trade

Common misconception: “a more powerful decoder gives a better VAE.” Past a point it gives a worse one, and in a specific and diagnosable way.

The ELBO is \(\mathbb E_q[\log p(x\mid z)]-D_{\mathrm{KL}}(q(z\mid x)\|p(z))\). If the decoder is expressive enough to model \(p(x)\) without \(z\), an autoregressive decoder over pixels or tokens, for instance, then the optimizer can set \(q(z\mid x)=p(z)\) for every \(x\), driving the KL term to exactly zero at no cost in reconstruction. That is the global optimum of the objective, and the latent carries no information.

The information-theoretic statement makes it precise. The KL term upper-bounds the mutual information between data and code, \[I(x;z)\ \le\ \mathbb E_{x}\big[D_{\mathrm{KL}}\big(q(z\mid x)\,\|\,p(z)\big)\big],\] so KL \(\to0\) forces \(I(x;z)\to0\). A VAE reporting a near-zero KL has not converged to a good representation; it has learned to ignore its own latent.

Standard remedies: KL annealing (ramp \(\beta\) from 0), free bits (do not penalize KL below a floor per dimension), and weakening the decoder.

train_vae_beta <- function(beta, d_lat = 8, hidden = 128, epochs = 1200,
                           seed = 91) {
  torch_manual_seed(seed)
  m <- vae(D = 2, d = d_lat, hidden = hidden)
  o <- optim_adam(m$parameters, lr = 2e-3)
  for (e in seq_len(epochs)) {
    o$zero_grad(); L <- vae_loss(m(Xv), Xv, beta = beta)
    L$total$backward(); o$step()
  }
  L <- with_no_grad(vae_loss(m(Xv), Xv, beta = beta))
  # Per-dimension KL identifies which latent dimensions are ACTIVE
  e_out <- with_no_grad(m$encode(Xv))
  kl_dim <- as.numeric(torch_mean(
    0.5*(e_out$mu^2 + torch_exp(e_out$logvar) - e_out$logvar - 1), dim = 1))
  c(beta = beta, recon = as.numeric(L$recon), kl = as.numeric(L$kl),
    active_dims = sum(kl_dim > 0.01))
}
betas <- c(0.05, 0.2, 1, 4, 16, 64)
pc <- as.data.frame(do.call(rbind, lapply(betas, train_vae_beta)))
pc |> mutate(across(c(recon, kl), \(z) signif(z, 4)))

pc |> select(beta, Distortion = recon, Rate = kl) |>
  pivot_longer(-beta, names_to = "term", values_to = "v") |>
  ggplot(aes(beta, pmax(v, 1e-4), color = term)) +
  geom_line(linewidth = 1) + geom_point(size = 2.4) +
  scale_x_log10(breaks = betas) + scale_y_log10() +
  scale_color_manual(values = c(Distortion = "#D8433B", Rate = "#3B7DD8")) +
  labs(title = "The rate-distortion trade that beta controls",
       subtitle = "Rate is the KL (bits stored in z); distortion is the reconstruction error. Large beta collapses the rate to zero",
       x = expression(beta~"(log scale)"), y = "Value (log scale)", color = NULL) +
  theme_dspa()

beta_grid <- c(0.05, 0.2, 0.5, 1, 2, 8, 32)
hid_grid  <- c(8, 16, 32, 64, 128, 256)
Zkl <- outer(hid_grid, beta_grid, Vectorize(function(h, b)
  log10(pmax(train_vae_beta(b, d_lat = 8, hidden = h, epochs = 700)[["kl"]], 1e-4))))

plot_ly(x = beta_grid, y = hid_grid, z = Zkl, type = "surface",
        colorscale = "Viridis", colorbar = list(title = "log10 KL\n(rate)")) |>
  layout(title = "Posterior collapse over beta and decoder capacity",
         scene = list(xaxis = list(title = "beta", type = "log"),
                      yaxis = list(title = "Decoder hidden width", type = "log"),
                      zaxis = list(title = "log10 KL (nats)")))

The surface falls toward zero in two directions. Raising \(\beta\) collapses the rate by pricing information too highly, the intended \(\beta\)-VAE trade. Widening the decoder collapses it because the decoder no longer needs \(z\). The second mechanism is the dangerous one, because it arrives while every other diagnostic looks healthy.

\(\beta\)-VAE (Higgins et al., 2017) uses \(\beta>1\) deliberately, to encourage disentangled factors. The empirical record is mixed: disentanglement is not identifiable from the data alone without inductive bias or supervision (Locatello et al., 2019), so a good disentanglement score is not evidence that the model found the generative factors.

PART II: IMPLICIT AND SCORE-BASED MODELS

14.72 Generative adversarial networks

A VAE optimizes a bound on the likelihood. A GAN never evaluates a likelihood at all: it trains a generator against a discriminator that tries to tell real samples from fake ones.

\[\boxed{\;\min_G\max_D\ V(D,G)=\mathbb E_{x\sim p_{\text{data}}}\big[\log D(x)\big]+\mathbb E_{z\sim p_z}\big[\log\big(1-D(G(z))\big)\big]\;}\]

14.72.1 The optimal discriminator

For a fixed generator inducing distribution \(p_g\), the inner maximization is solved pointwise. Writing the objective as an integral,

\[V=\int\Big[p_{\text{data}}(x)\log D(x)+p_g(x)\log\big(1-D(x)\big)\Big]dx,\]

and maximizing the integrand, \(a\log y+b\log(1-y)\) is maximized at \(y=a/(a+b)\), gives

\[\boxed{\;D^\star(x)=\frac{p_{\text{data}}(x)}{p_{\text{data}}(x)+p_g(x)}\;}\]

Substituting back,

\[V(D^\star,G)=2\,D_{\mathrm{JS}}\big(p_{\text{data}}\,\|\,p_g\big)-2\log 2,\]

so the generator minimizes the Jensen-Shannon divergence, with the global optimum at \(p_g=p_{\text{data}}\) where \(V=-2\log2\).

# Verify D* on a case where both densities are known
x_g <- seq(-5, 6, length.out = 400)
p_data <- dnorm(x_g, 0, 1)
p_gen  <- dnorm(x_g, 1.5, 1.3)
D_star <- p_data / (p_data + p_gen)

# Train a discriminator by gradient descent and compare
set.seed(101); torch_manual_seed(101)
real <- torch_tensor(matrix(rnorm(4000, 0, 1), ncol = 1))
fake <- torch_tensor(matrix(rnorm(4000, 1.5, 1.3), ncol = 1))
Dnet <- nn_sequential(nn_linear(1, 64), nn_relu(), nn_linear(64, 64),
                      nn_relu(), nn_linear(64, 1))
optD <- optim_adam(Dnet$parameters, lr = 3e-3)
for (s in 1:1500) {
  optD$zero_grad()
  loss <- nnf_binary_cross_entropy_with_logits(Dnet(real), torch_ones(4000, 1)) +
          nnf_binary_cross_entropy_with_logits(Dnet(fake), torch_zeros(4000, 1))
  loss$backward(); optD$step()
}
D_learned <- as.numeric(with_no_grad(torch_sigmoid(
  Dnet(torch_tensor(matrix(x_g, ncol = 1))))))
keep_g <- x_g > -3 & x_g < 4.5
c(max_abs_error_vs_theory = signif(max(abs(D_learned[keep_g] - D_star[keep_g])), 3),
  correlation = signif(cor(D_learned[keep_g], D_star[keep_g]), 5))

#> max_abs_error_vs_theory             correlation 
#>                  0.0740                  0.9981

data.frame(x = x_g, `theoretical D*` = D_star, `learned D` = D_learned,
           check.names = FALSE) |>
  pivot_longer(-x, names_to = "curve", values_to = "v") |>
  ggplot(aes(x, v, color = curve)) +
  geom_hline(yintercept = 0.5, linetype = "dotted", color = "grey50") +
  geom_line(linewidth = 0.9) +
  scale_color_manual(values = c(`theoretical D*` = "#D8433B",
                                 `learned D` = "#3B7DD8")) +
  labs(title = "The trained discriminator recovers the density ratio",
       subtitle = expression(D^"*"*(x) == p[data](x)/(p[data](x)+p[g](x))*
                             ", learned without either density being supplied"),
       x = "x", y = "D(x)", color = NULL) +
  theme_dspa()

A GAN discriminator is a density-ratio estimator. It never sees either density, and at optimum it reports their ratio.

14.73 Why GAN training is hard

Common misconception: “the generator minimizes the distance between the two distributions, so gradient descent will close the gap.” It minimizes the Jensen-Shannon divergence, and JS has a property that breaks exactly this reasoning.

When the supports of \(p_{\text{data}}\) and \(p_g\) are disjoint, \(D_{\mathrm{JS}}=\log 2\), a constant, regardless of how far apart they are. Its gradient with respect to the generator is therefore zero, and the generator receives no signal telling it which way to move.

This is not a pathological corner case. Real data lie near a low-dimensional manifold in a high-dimensional space, and so does \(p_g\) (it is the pushforward of a low-dimensional \(z\)). Two such manifolds generically intersect in a set of measure zero (Arjovsky & Bottou, 2017), so disjoint support is the typical case, not the exception.

Two consequences follow. The better the discriminator, the worse the generator’s gradient, the opposite of the usual relationship between an auxiliary model’s quality and its usefulness. And the non-saturating loss is not cosmetic: replacing \(\min_G\log(1-D(G(z)))\) with \(\max_G\log D(G(z))\) leaves the fixed point unchanged while supplying a gradient that grows rather than vanishes when \(D\) is confident.

# The canonical example: two point masses (or narrow Gaussians) at distance theta
theta <- seq(-3, 3, length.out = 400)
sig_n <- 0.05                                  # near-disjoint support

js_div <- function(th, s = sig_n, grid = seq(-6, 6, length.out = 4000)) {
  p <- dnorm(grid, 0, s); q <- dnorm(grid, th, s); m <- 0.5*(p + q)
  dx <- diff(grid)[1]
  kl <- function(a, b) sum(ifelse(a > 1e-300, a * log(pmax(a, 1e-300)/pmax(b, 1e-300)), 0)) * dx
  0.5*kl(p, m) + 0.5*kl(q, m)
}
divs <- data.frame(theta = theta,
                   `Jensen-Shannon` = vapply(theta, js_div, numeric(1)),
                   `Wasserstein-1` = abs(theta), check.names = FALSE)

p_div <- divs |> pivot_longer(-theta, names_to = "divergence", values_to = "v") |>
  ggplot(aes(theta, v, color = divergence)) +
  geom_line(linewidth = 1) +
  scale_color_manual(values = c(`Jensen-Shannon` = "#D8433B",
                                 `Wasserstein-1` = "#3B7DD8")) +
  labs(title = "Divergence between two narrow distributions separated by theta",
       x = expression(theta), y = "Divergence", color = NULL) +
  theme_dspa(9)

grad_df <- divs |>
  mutate(`Jensen-Shannon` = c(NA, diff(`Jensen-Shannon`)/diff(theta)),
         `Wasserstein-1` = c(NA, diff(`Wasserstein-1`)/diff(theta))) |>
  pivot_longer(-theta, names_to = "divergence", values_to = "g")
p_grad <- ggplot(filter(grad_df, !is.na(g)), aes(theta, g, color = divergence)) +
  geom_hline(yintercept = 0, color = "grey70") +
  geom_line(linewidth = 1) +
  scale_color_manual(values = c(`Jensen-Shannon` = "#D8433B",
                                 `Wasserstein-1` = "#3B7DD8"), guide = "none") +
  labs(title = "Their gradients -- what the generator actually receives",
       x = expression(theta), y = "d(divergence)/d(theta)") +
  theme_dspa(9)

p_div / p_grad

data.frame(
  separation = c(0.2, 0.5, 1.0, 2.0, 3.0),
  JS = signif(vapply(c(0.2, 0.5, 1, 2, 3), js_div, numeric(1)), 4),
  JS_at_log2 = signif(log(2), 4),
  Wasserstein = c(0.2, 0.5, 1.0, 2.0, 3.0),
  JS_gradient_informative = c(TRUE, FALSE, FALSE, FALSE, FALSE))

# JS saturates sooner the NARROWER the distributions -- and real data lie on a
# low-dimensional manifold, which is the narrow-support limit
sep_grid <- seq(0.02, 2.5, length.out = 40)
wid_grid <- 10^seq(-1.6, -0.1, length.out = 30)
Zjs <- outer(wid_grid, sep_grid, Vectorize(function(w, s) js_div(s, s = w)))

plot_ly(x = sep_grid, y = wid_grid, z = Zjs, type = "surface",
        colorscale = "Inferno", reversescale = TRUE,
        colorbar = list(title = "JS\ndivergence")) |>
  add_surface(x = sep_grid, y = wid_grid,
              z = matrix(log(2), length(wid_grid), length(sep_grid)),
              opacity = 0.3, showscale = FALSE,
              colorscale = list(c(0, "black"), c(1, "black"))) |>
  layout(title = "Jensen-Shannon divergence over separation and distribution width; the flat plane is log 2",
         scene = list(xaxis = list(title = "Separation"),
                      yaxis = list(title = "Distribution width", type = "log"),
                      zaxis = list(title = "JS divergence")))

The flat black plane is the saturation ceiling \(\log 2\). Rotate toward narrow widths: the colored surface climbs to meet the plane almost immediately, so the region where JS carries any gradient shrinks toward nothing. Real data lie near a low-dimensional manifold, which is exactly the narrow limit, so the useful region is not merely small, it is generically empty.

JS saturates at \(\log 2\approx0.693\) and stays there, so its derivative is zero for any separation beyond a fraction of the distributions’ width. The Wasserstein distance grows linearly and its derivative is constant, the generator always knows which way to move.

# Saturating vs non-saturating generator loss, as functions of D(G(z))
d_out <- seq(0.001, 0.999, length.out = 300)
losses <- data.frame(
  D_of_fake = d_out,
  `saturating: log(1 - D)` = log(1 - d_out),
  `non-saturating: -log D` = -log(d_out), check.names = FALSE)
grads <- data.frame(
  D_of_fake = d_out,
  `saturating` = -1/(1 - d_out) * -1,          # |d/dD log(1-D)|
  `non-saturating` = 1/d_out)

ggplot(pivot_longer(grads, -D_of_fake, names_to = "loss", values_to = "g"),
       aes(D_of_fake, g, color = loss)) +
  geom_line(linewidth = 1) +
  scale_y_log10() +
  scale_color_manual(values = c(saturating = "#D8433B",
                                 `non-saturating` = "#3B7DD8")) +
  labs(title = "Generator gradient magnitude against discriminator confidence",
       subtitle = "Early in training D(G(z)) is near 0. The saturating loss gives almost no gradient there; the non-saturating one gives the most",
       x = "D(G(z))  -- small means the discriminator is winning",
       y = "|gradient| (log scale)", color = NULL) +
  theme_dspa()

# Mode collapse on an eight-mode ring, the standard diagnostic
set.seed(111); torch_manual_seed(111)
n_modes <- 8; n_real <- 4000
ang <- (0:(n_modes-1)) * 2*pi/n_modes
centres <- cbind(2*cos(ang), 2*sin(ang))
mode_id <- sample(n_modes, n_real, TRUE)
X_ring <- centres[mode_id, ] + matrix(rnorm(n_real*2, sd = 0.08), n_real, 2)
Xr <- torch_tensor(X_ring)

train_gan <- function(steps = 4000, wgan_gp = FALSE, seed = 113) {
  torch_manual_seed(seed)
  G <- nn_sequential(nn_linear(2, 128), nn_relu(), nn_linear(128, 128),
                     nn_relu(), nn_linear(128, 2))
  D <- nn_sequential(nn_linear(2, 128), nn_relu(), nn_linear(128, 128),
                     nn_relu(), nn_linear(128, 1))
  oG <- optim_adam(G$parameters, lr = 1e-3, betas = c(0.5, 0.9))
  oD <- optim_adam(D$parameters, lr = 1e-3, betas = c(0.5, 0.9))
  bs <- 256
  for (s in seq_len(steps)) {
    # --- discriminator / critic ---
    for (k in seq_len(if (wgan_gp) 3 else 1)) {
      idx <- sample(n_real, bs)
      xr <- Xr[idx, ]; xf <- with_no_grad(G(torch_randn(bs, 2)))
      oD$zero_grad()
      if (wgan_gp) {
        eps <- torch_rand(bs, 1)
        xh <- (eps * xr + (1 - eps) * xf)$requires_grad_(TRUE)
        dh <- D(xh)
        gr <- autograd_grad(dh, xh, torch_ones_like(dh), create_graph = TRUE)[[1]]
        gp <- torch_mean((torch_norm(gr, dim = 2) - 1)^2)
        lD <- torch_mean(D(xf)) - torch_mean(D(xr)) + 10 * gp
      } else {
        lD <- nnf_binary_cross_entropy_with_logits(D(xr), torch_ones(bs, 1)) +
              nnf_binary_cross_entropy_with_logits(D(xf), torch_zeros(bs, 1))
      }
      lD$backward(); oD$step()
    }
    # --- generator (non-saturating for the standard GAN) ---
    oG$zero_grad()
    xf <- G(torch_randn(bs, 2))
    lG <- if (wgan_gp) -torch_mean(D(xf)) else
      nnf_binary_cross_entropy_with_logits(D(xf), torch_ones(bs, 1))
    lG$backward(); oG$step()
  }
  as.matrix(with_no_grad(G(torch_randn(2000, 2))))
}
S_gan  <- train_gan(wgan_gp = FALSE)
S_wgan <- train_gan(wgan_gp = TRUE)

modes_covered <- function(S, tol = 0.5) {
  d2 <- as.matrix(dist(rbind(centres, S)))[1:n_modes, -(1:n_modes)]
  sum(apply(d2, 1, min) < tol)
}
data.frame(model = c("standard GAN (non-saturating)", "WGAN-GP"),
           modes_covered_of_8 = c(modes_covered(S_gan), modes_covered(S_wgan)))

bind_rows(data.frame(x = X_ring[,1], y = X_ring[,2], panel = "Data (8 modes)"),
          data.frame(x = S_gan[,1],  y = S_gan[,2],  panel = "Standard GAN"),
          data.frame(x = S_wgan[,1], y = S_wgan[,2], panel = "WGAN-GP")) |>
  mutate(panel = factor(panel, levels = unique(panel))) |>
  ggplot(aes(x, y)) +
  geom_point(size = 0.4, alpha = 0.4, color = "steelblue") +
  geom_point(data = data.frame(x = centres[,1], y = centres[,2]),
             color = "firebrick", size = 1.6, shape = 3) +
  facet_wrap(~ panel) + coord_fixed() +
  labs(title = "Mode coverage on an eight-mode ring",
       subtitle = "Red crosses mark the true modes. A collapsed generator concentrates on a subset",
       x = NULL, y = NULL) +
  theme_dspa(10)

Common misconception: “mode collapse means the generator is undertrained.” It is a property of the objective, not of the training budget. The minimax objective asks the generator to fool the discriminator; a generator that produces one perfectly realistic mode succeeds at that, because nothing in \(V(D,G)\) rewards coverage of the data distribution.

The asymmetry is visible in the divergence. Minimizing the reverse KL \(D_{\mathrm{KL}}(p_g\|p_{\text{data}})\) is mode-seeking: it is finite only where \(p_g\) has support, so concentrating on one mode is safe. Minimizing the forward KL \(D_{\mathrm{KL}}(p_{\text{data}}\|p_g)\) is mode-covering: it is infinite wherever \(p_{\text{data}}>0\) and \(p_g=0\), so every mode must be covered. Maximum likelihood, and hence the VAE’s ELBO, optimizes something closer to the forward direction, which is why VAEs produce blurry samples that cover the data while GANs produce sharp samples that may not.

Training longer does not fix it. Changing the objective can.

14.74 Wasserstein GAN

The Wasserstein-1 (earth mover’s) distance measures the minimum cost of transporting one distribution onto the other:

\[W_1(p,q)=\inf_{\gamma\in\Pi(p,q)}\ \mathbb E_{(x,y)\sim\gamma}\big[\|x-y\|\big],\]

with \(\Pi(p,q)\) the set of couplings. That infimum over couplings is intractable, but Kantorovich-Rubinstein duality converts it into a supremum over functions:

\[\boxed{\;W_1(p,q)=\sup_{\|f\|_L\le1}\ \mathbb E_{x\sim p}\big[f(x)\big]-\mathbb E_{x\sim q}\big[f(x)\big]\;}\]

So the “discriminator” becomes a critic \(f\) producing an unbounded real score rather than a probability, and the constraint is that \(f\) be 1-Lipschitz.

\[ \begin{aligned} \textbf{Weight clipping (original): }&\quad \text{clip all weights to }[-c,c] &&\text{crude; degrades capacity}\\ \textbf{Gradient penalty (WGAN-GP): }&\quad +\lambda\,\mathbb E_{\hat x}\big[\big(\|\nabla_{\hat x}f(\hat x)\|_2-1\big)^2\big] &&\text{penalize deviation from }\|\nabla f\|=1\\ \textbf{Spectral normalization: }&\quad W\leftarrow W/\sigma_{\max}(W) &&\text{bounds the Lipschitz constant by construction} \end{aligned} \]

The gradient penalty is evaluated at points \(\hat x\) interpolated between real and generated samples, because that is where the constraint binds.

Common misconception: “a falling adversary loss means the generator is improving.” For a standard GAN it means almost nothing.

A standard GAN’s discriminator loss hovers near \(\log 4\) whether the generator is excellent or terrible, so it correlates poorly with sample quality, one of the practical miseries of GAN training. The WGAN critic’s output estimates \(W_1\) up to a constant, so it decreases as the generator improves and can be used to decide when to stop.

# Track both losses against a ground-truth quality measure during training
set.seed(121); torch_manual_seed(121)
track_training <- function(wgan_gp, steps = 2500, seed = 123) {
  torch_manual_seed(seed)
  G <- nn_sequential(nn_linear(2, 96), nn_relu(), nn_linear(96, 2))
  D <- nn_sequential(nn_linear(2, 96), nn_relu(), nn_linear(96, 1))
  oG <- optim_adam(G$parameters, lr = 1e-3, betas = c(0.5, 0.9))
  oD <- optim_adam(D$parameters, lr = 1e-3, betas = c(0.5, 0.9))
  bs <- 256; out <- data.frame()
  for (s in seq_len(steps)) {
    for (k in seq_len(if (wgan_gp) 3 else 1)) {
      idx <- sample(n_real, bs)
      xr <- Xr[idx, ]; xf <- with_no_grad(G(torch_randn(bs, 2)))
      oD$zero_grad()
      if (wgan_gp) {
        eps <- torch_rand(bs, 1)
        xh <- (eps*xr + (1-eps)*xf)$requires_grad_(TRUE)
        dh <- D(xh)
        gr <- autograd_grad(dh, xh, torch_ones_like(dh), create_graph = TRUE)[[1]]
        lD <- torch_mean(D(xf)) - torch_mean(D(xr)) +
              10*torch_mean((torch_norm(gr, dim = 2) - 1)^2)
      } else {
        lD <- nnf_binary_cross_entropy_with_logits(D(xr), torch_ones(bs,1)) +
              nnf_binary_cross_entropy_with_logits(D(xf), torch_zeros(bs,1))
      }
      lD$backward(); oD$step()
    }
    oG$zero_grad(); xf <- G(torch_randn(bs, 2))
    lG <- if (wgan_gp) -torch_mean(D(xf)) else
      nnf_binary_cross_entropy_with_logits(D(xf), torch_ones(bs,1))
    lG$backward(); oG$step()
    if (s %% 125 == 0) {
      S <- as.matrix(with_no_grad(G(torch_randn(800, 2))))
      dd <- as.matrix(dist(rbind(centres, S)))[1:n_modes, -(1:n_modes)]
      out <- rbind(out, data.frame(step = s,
        critic_estimate = if (wgan_gp) -as.numeric(lD) else as.numeric(lD),
        mean_dist_to_nearest_mode = mean(apply(dd, 2, min))))
    }
  }
  out
}
tr_w <- track_training(TRUE); tr_g <- track_training(FALSE)
c(WGAN_correlation_loss_vs_quality =
    signif(cor(tr_w$critic_estimate, tr_w$mean_dist_to_nearest_mode), 3),
  GAN_correlation_loss_vs_quality =
    signif(cor(tr_g$critic_estimate, tr_g$mean_dist_to_nearest_mode), 3))

#> WGAN_correlation_loss_vs_quality  GAN_correlation_loss_vs_quality 
#>                           0.4090                          -0.0211

bind_rows(mutate(tr_w, model = "WGAN-GP critic"),
          mutate(tr_g, model = "GAN discriminator")) |>
  ggplot(aes(mean_dist_to_nearest_mode, critic_estimate, color = model)) +
  geom_point(size = 1.6, alpha = 0.7) +
  geom_smooth(method = "lm", se = FALSE, linewidth = 0.7) +
  facet_wrap(~ model, scales = "free_y") +
  scale_color_manual(values = c("#D8433B", "#3B7DD8"), guide = "none") +
  labs(title = "Does the adversary's loss track sample quality?",
       subtitle = "Horizontal axis is a ground-truth quality measure. Only the WGAN critic correlates with it",
       x = "Mean distance from a sample to its nearest true mode",
       y = "Adversary's loss") +
  theme_dspa(10)

# --- Interactive equivalent ------------------------------------------------
plot_ly(tr_w, x = ~step, y = ~critic_estimate, type = "scatter", mode = "lines",
        name = "WGAN critic estimate") |>
  add_lines(y = ~mean_dist_to_nearest_mode * max(tr_w$critic_estimate) /
              max(tr_w$mean_dist_to_nearest_mode), name = "Quality (rescaled)") |>
  layout(title = "WGAN critic estimate tracks sample quality",
         xaxis = list(title = "Step"), yaxis = list(title = "Value"))

14.75 Diffusion models

Diffusion models destroy structure with noise, then learn to reverse the process. They now dominate image, audio, and molecular generation, and the derivation runs through machinery already built in this part.

14.75.1 The forward process

Add Gaussian noise in \(T\) steps according to a variance schedule \(\beta_1,\dots,\beta_T\):

\[q\big(x_t\mid x_{t-1}\big)=\mathcal N\Big(x_t;\ \sqrt{1-\beta_t}\,x_{t-1},\ \beta_t I\Big).\]

The composition has a closed form, the property that makes training tractable. With \(\alpha_t=1-\beta_t\) and \(\bar\alpha_t=\prod_{s\le t}\alpha_s\),

\[\boxed{\;q\big(x_t\mid x_0\big)=\mathcal N\Big(x_t;\ \sqrt{\bar\alpha_t}\,x_0,\ (1-\bar\alpha_t)I\Big) \quad\Longleftrightarrow\quad x_t=\sqrt{\bar\alpha_t}\,x_0+\sqrt{1-\bar\alpha_t}\,\varepsilon\;}\]

so any noise level is reachable in one step. Training never simulates the chain: it samples \(t\) uniformly, jumps straight to \(x_t\), and asks the network to identify the noise.

make_schedule <- function(T_steps = 1000, kind = c("linear", "cosine", "quadratic")) {
  kind <- match.arg(kind)
  betas <- switch(kind,
    linear    = seq(1e-4, 0.02, length.out = T_steps),
    quadratic = seq(sqrt(1e-4), sqrt(0.02), length.out = T_steps)^2,
    cosine    = {  # Nichol & Dhariwal parameterization
      s <- 0.008; t <- seq(0, T_steps) / T_steps
      f <- cos((t + s)/(1 + s) * pi/2)^2
      ab <- f / f[1]
      pmin(1 - ab[-1]/ab[-length(ab)], 0.999)
    })
  list(betas = betas, alphas = 1 - betas, abar = cumprod(1 - betas))
}
sched <- lapply(c("linear","quadratic","cosine"), \(k) make_schedule(1000, k))
names(sched) <- c("linear","quadratic","cosine")

bind_rows(lapply(names(sched), \(k)
  data.frame(t = 1:1000, abar = sched[[k]]$abar, schedule = k))) |>
  ggplot(aes(t, abar, color = schedule)) +
  geom_line(linewidth = 1) +
  scale_color_manual(values = c(linear = "#D8433B", quadratic = "#7FB069",
                                 cosine = "#3B7DD8")) +
  labs(title = "Signal retained through the forward process",
       subtitle = expression(bar(alpha)[t]*" is the fraction of the original signal at step t; "*
                             sqrt(1-bar(alpha)[t])*" is the noise"),
       x = "Diffusion step t", y = expression(bar(alpha)[t]), color = NULL) +
  theme_dspa()

The linear schedule destroys information too early. By \(t\approx300\) of 1000, \(\bar\alpha_t\) has already fallen below 0.1, so the last 70% of the chain operates on almost pure noise and contributes little training signal. The cosine schedule (Nichol & Dhariwal, 2021) keeps \(\bar\alpha_t\) high longer and spends the budget where structure still exists — a straightforward improvement in sample quality obtained purely by changing the schedule.

T_grid <- round(seq(1, 1000, length.out = 60))
sch_names <- c("linear", "quadratic", "cosine")
Zsch <- t(vapply(sch_names, \(k) sched[[k]]$abar[T_grid], numeric(length(T_grid))))
# Signal-to-noise ratio is what the network actually experiences
Zsnr <- log10(pmax(Zsch / (1 - Zsch), 1e-8))

plot_ly(x = T_grid, y = sch_names, z = Zsnr, type = "surface",
        colorscale = "Viridis", colorbar = list(title = "log10 SNR")) |>
  layout(title = "Signal-to-noise ratio across the diffusion chain, by schedule",
         scene = list(xaxis = list(title = "Diffusion step t"),
                      yaxis = list(title = "Schedule"),
                      zaxis = list(title = "log10 (abar / (1 - abar))")))

Rotate along the step axis. The linear ridge falls off a cliff early; the cosine ridge descends gradually. Since the training loss is averaged uniformly over \(t\), a schedule that spends most of its steps at negligible SNR is wasting most of its gradient signal.

14.75.2 The training objective

The reverse process is parameterized as \(p_\theta(x_{t-1}\mid x_t)=\mathcal N\big(x_{t-1};\mu_\theta(x_t,t),\sigma_t^2I\big)\), and the variational bound on \(\log p_\theta(x_0)\) decomposes into a sum of KL divergences between Gaussians, each available in closed form.

Reparameterizing \(\mu_\theta\) in terms of the noise and dropping the weighting gives the objective actually used (Ho et al., 2020):

\[\boxed{\;\mathcal L_{\text{simple}}=\mathbb E_{t\sim\mathcal U[1,T],\ x_0,\ \varepsilon\sim\mathcal N(0,I)}\Big[\big\|\varepsilon-\varepsilon_\theta\big(\sqrt{\bar\alpha_t}x_0+\sqrt{1-\bar\alpha_t}\varepsilon,\ t\big)\big\|^2\Big]\;}\]

Common misconception: “diffusion models are a different paradigm from VAEs.” They are a hierarchical VAE with three special features: the encoder \(q(x_t\mid x_{t-1})\) is fixed, not learned; the latent has the same dimension as the data at every level; and the bound decomposes into analytically computable Gaussian KLs.

That third property is what removes the difficulties of §14.69-14.71. There is no posterior collapse, because the posterior is not learned. There is no reparameterization variance problem, because the objective reduces to a regression. And the “ELBO” collapses, after reweighting, to plain least squares, which is why diffusion training is stable where GAN training is not.

The connection to §14.68 completes the picture: \(\varepsilon_\theta\) is proportional to the score, \[\varepsilon_\theta(x_t,t)\ \approx\ -\sqrt{1-\bar\alpha_t}\ \nabla_{x_t}\log q(x_t),\] so a diffusion model is a denoising autoencoder trained at every noise level at once, and sampling is Langevin-style ascent on the learned score.

# A complete 1-D diffusion model, small enough to inspect end to end
set.seed(131); torch_manual_seed(131)
T_diff <- 200
sc <- make_schedule(T_diff, "cosine")
abar_t <- torch_tensor(sc$abar)
alphas_t <- torch_tensor(sc$alphas)
betas_t <- torch_tensor(sc$betas)

# Target: the same bimodal mixture as Section 14.68
x0_data <- torch_tensor(matrix(c(rnorm(4000, -1.5, 0.5), rnorm(3000, 1.8, 0.7)),
                               ncol = 1))

eps_net <- nn_module("EpsNet",
  initialize = function(hidden = 128) {
    self$t_emb <- nn_sequential(nn_linear(1, 32), nn_silu(), nn_linear(32, 32))
    self$net <- nn_sequential(nn_linear(1 + 32, hidden), nn_silu(),
                              nn_linear(hidden, hidden), nn_silu(),
                              nn_linear(hidden, 1))
  },
  forward = function(x, t_norm)
    self$net(torch_cat(list(x, self$t_emb(t_norm)), dim = 2)))()

opt_e <- optim_adam(eps_net$parameters, lr = 2e-3)
n_d <- x0_data$size(1)
for (step in 1:4000) {
  opt_e$zero_grad()
  idx <- sample(n_d, 512)
  x0 <- x0_data[idx, , drop = FALSE]
  t_i <- sample(T_diff, 512, replace = TRUE)
  ab <- abar_t[t_i]$unsqueeze(2)
  eps <- torch_randn_like(x0)
  xt <- torch_sqrt(ab) * x0 + torch_sqrt(1 - ab) * eps       # one-step jump
  pred <- eps_net(xt, torch_tensor(matrix(t_i / T_diff, ncol = 1)))
  nnf_mse_loss(pred, eps)$backward()
  opt_e$step()
}
c(final_objective = signif(as.numeric(with_no_grad(nnf_mse_loss(
    eps_net(xt, torch_tensor(matrix(t_i/T_diff, ncol=1))), eps))), 4),
  note = "the whole objective is least squares on the injected noise")

#>                                              final_objective 
#>                                                     "0.5424" 
#>                                                         note 
#> "the whole objective is least squares on the injected noise"

14.76 Sampling from a diffusion model

Start from \(x_T\sim\mathcal N(0,I)\) and iterate the reverse step:

\[x_{t-1}=\frac{1}{\sqrt{\alpha_t}}\left(x_t-\frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\,\varepsilon_\theta(x_t,t)\right)+\sigma_t\,\mathbf z,\qquad \mathbf z\sim\mathcal N(0,I).\]

ddpm_sample <- function(n, keep_at = c(200, 150, 100, 50, 1)) {
  x <- torch_randn(n, 1)
  snap <- list()
  for (t in T_diff:1) {
    z <- if (t > 1) torch_randn(n, 1) else torch_zeros(n, 1)
    e <- with_no_grad(eps_net(x, torch_full(c(n, 1), t / T_diff)))
    x <- (1/torch_sqrt(alphas_t[t])) *
         (x - (betas_t[t]/torch_sqrt(1 - abar_t[t])) * e) +
         torch_sqrt(betas_t[t]) * z
    if (t %in% keep_at) snap[[as.character(t)]] <- as.numeric(x)
  }
  list(final = as.numeric(x), snapshots = snap)
}
samp <- ddpm_sample(3000)

bind_rows(lapply(names(samp$snapshots), \(k)
  data.frame(x = samp$snapshots[[k]], t = sprintf("t = %s", k)))) |>
  mutate(t = factor(t, levels = sprintf("t = %s", c(200,150,100,50,1)))) |>
  ggplot(aes(x)) +
  geom_histogram(aes(y = after_stat(density)), bins = 60, fill = "steelblue",
                 color = "white", linewidth = 0.1) +
  stat_function(fun = mix_p, color = "firebrick", linewidth = 0.7) +
  facet_wrap(~ t, nrow = 1) +
  coord_cartesian(xlim = c(-4, 4.5)) +
  labs(title = "Reverse diffusion, from noise to the data distribution",
       subtitle = "Red: the true density. The chain runs right to left, from t = 200 (pure noise) to t = 1",
       x = "x", y = "density") +
  theme_dspa(9)

ks <- ks.test(samp$final, function(q)
  0.6*pnorm(q, -1.5, 0.5) + 0.4*pnorm(q, 1.8, 0.7))
c(KS_statistic = signif(ks$statistic, 4),
  sample_mean = signif(mean(samp$final), 4),
  target_mean = signif(0.6*(-1.5) + 0.4*1.8, 4))

#> KS_statistic.D    sample_mean    target_mean 
#>        0.03968       -0.05750       -0.18000

# --- Interactive equivalent ------------------------------------------------
plot_ly(x = samp$final, type = "histogram", histnorm = "probability density",
        name = "Generated", nbinsx = 60) |>
  add_lines(x = grid_x, y = mix_p(grid_x), name = "True density") |>
  layout(title = "Diffusion samples against the target density",
         xaxis = list(title = "x"), yaxis = list(title = "density"))

Sampling cost is the diffusion model’s central weakness. Generation requires \(T\) sequential network evaluations, 1,000 forward passes per sample in the original formulation, against one for a GAN or VAE.

The remedies attack the step count directly. DDIM (Song et al., 2021) reinterprets the reverse chain as a deterministic ODE solver, allowing 20-50 steps with little quality loss. Distillation trains a student to reproduce many teacher steps in one, reaching single-digit step counts. Latent diffusion runs the process in a compressed autoencoder space (§14.66) rather than pixel space, cutting the cost per step, the basis of current text-to-image systems.

# DDIM: deterministic sampling on a strided subsequence of timesteps
ddim_sample <- function(n, n_steps) {
  taus <- unique(round(seq(T_diff, 1, length.out = n_steps)))
  x <- torch_randn(n, 1)
  for (i in seq_along(taus)) {
    t <- taus[i]
    e <- with_no_grad(eps_net(x, torch_full(c(n, 1), t / T_diff)))
    ab_t <- abar_t[t]
    x0_hat <- (x - torch_sqrt(1 - ab_t) * e) / torch_sqrt(ab_t)
    if (i < length(taus)) {
      ab_prev <- abar_t[taus[i + 1]]
      x <- torch_sqrt(ab_prev) * x0_hat + torch_sqrt(1 - ab_prev) * e
    } else x <- x0_hat
  }
  as.numeric(x)
}
steps_try <- c(5, 10, 25, 50, 200)
ddim_res <- data.frame(
  steps = steps_try,
  KS_vs_truth = vapply(steps_try, \(s) {
    xs <- ddim_sample(2500, s)
    as.numeric(ks.test(xs, function(q)
      0.6*pnorm(q, -1.5, 0.5) + 0.4*pnorm(q, 1.8, 0.7))$statistic) }, numeric(1)))
ddim_res$relative_cost <- round(ddim_res$steps / T_diff, 3)
ddim_res |> mutate(KS_vs_truth = signif(KS_vs_truth, 4))

ggplot(ddim_res, aes(steps, KS_vs_truth)) +
  geom_line(linewidth = 1, color = "steelblue") + geom_point(size = 2.6) +
  scale_x_log10(breaks = steps_try) +
  labs(title = "DDIM sample quality against the number of denoising steps",
       subtitle = "Quality saturates well below the 200 training steps -- most of the chain is unnecessary at sampling time",
       x = "Sampling steps (log scale)", y = "KS distance to the true CDF") +
  theme_dspa()

14.77 Evaluating generative models

Classification has accuracy; regression has RMSE. Generative modelling has no comparably trustworthy scalar, and the two most-cited metrics are weaker than their ubiquity suggests.

\[ \begin{aligned} \textbf{Inception Score: }&\quad \mathrm{IS}=\exp\Big(\mathbb E_x\big[D_{\mathrm{KL}}\big(p(y\mid x)\,\|\,p(y)\big)\big]\Big)\\[2mm] \textbf{Fréchet Inception Distance: }&\quad \mathrm{FID}=\|\mu_r-\mu_g\|^2+\operatorname{Tr}\Big(\Sigma_r+\Sigma_g-2\big(\Sigma_r\Sigma_g\big)^{1/2}\Big) \end{aligned} \]

Common misconception: “a low FID means the model learned the distribution.” FID is the Fréchet distance between Gaussians fitted to Inception-network features, and each element of that description is a weakness.

It assumes the features are Gaussian. They are not, and FID is blind to any discrepancy that leaves the first two moments intact.

A memorizing model scores perfectly. A “generator” that outputs training images verbatim achieves FID \(\approx 0\), the feature distributions are identical because the samples are. FID cannot distinguish learning from copying, so it must be reported alongside a novelty check.

It is biased in the sample size, and the bias does not vanish quickly: FID computed on 5,000 samples is systematically larger than on 50,000 for the same model, so numbers are comparable only at matched \(n\).

It inherits the feature extractor’s blind spots. An ImageNet-trained Inception network is a poor feature space for medical images, and FID computed there measures distance in a representation that never learned the relevant structure.

It conflates fidelity with coverage. A model producing beautiful samples from one mode and a model covering all modes badly can score alike. Precision and recall for generative models (Kynkäänniemi et al., 2019) separate the two, precision measures whether samples land in the data manifold, recall whether the data manifold is covered, and are the right supplement.

# Frechet distance between Gaussians, computed exactly in a low-dimensional
# feature space so its behaviour is visible
frechet <- function(mu1, S1, mu2, S2) {
  e <- eigen((S1 %*% S2 + t(S1 %*% S2))/2, symmetric = TRUE)
  sq <- sum(sqrt(pmax(e$values, 0)))
  sum((mu1 - mu2)^2) + sum(diag(S1)) + sum(diag(S2)) - 2*sq
}
set.seed(141)
d_f <- 8
real_feat <- matrix(rnorm(20000 * d_f), 20000, d_f)
mu_r <- colMeans(real_feat); S_r <- cov(real_feat)

# (a) Bias in the sample size, for a model that is EXACTLY correct
bias_n <- data.frame(n = c(250, 500, 1000, 2500, 5000, 10000)) |>
  rowwise() |>
  mutate(FID = { g <- matrix(rnorm(n * d_f), n, d_f)
                 frechet(mu_r, S_r, colMeans(g), cov(g)) }) |> ungroup()

# (b) A memorizing "generator" that resamples the training features
memo <- real_feat[sample(20000, 5000), ]
fid_memo <- frechet(mu_r, S_r, colMeans(memo), cov(memo))
# (c) A mode-dropping generator: covers only half the feature space
drop <- real_feat[real_feat[,1] > 0, ][1:5000, ]
fid_drop <- frechet(mu_r, S_r, colMeans(drop), cov(drop))

data.frame(
  generator = c("perfect (5000 samples)", "memorizes training data",
                "drops half the distribution"),
  FID = signif(c(bias_n$FID[bias_n$n == 5000], fid_memo, fid_drop), 4),
  novel_samples = c(TRUE, FALSE, TRUE))

ggplot(bias_n, aes(n, FID)) +
  geom_line(linewidth = 1, color = "steelblue") + geom_point(size = 2.4) +
  scale_x_log10(breaks = bias_n$n) +
  labs(title = "FID for a generator that is exactly correct",
       subtitle = "The true value is 0. The estimate is biased upward and the bias shrinks slowly, so numbers compare only at matched n",
       x = "Number of generated samples (log scale)", y = "FID") +
  theme_dspa()

# The bias is a surface over sample size and TRUE distance, so two models
# compared at different n can be ranked backwards
n_fid <- round(10^seq(2.3, 4, length.out = 25))
shift_fid <- seq(0, 1.2, length.out = 25)
set.seed(145)
Zfid <- outer(shift_fid, n_fid, Vectorize(function(sh, n) {
  g <- matrix(rnorm(n * d_f), n, d_f)
  g[, 1] <- g[, 1] + sh                       # a controlled true discrepancy
  frechet(mu_r, S_r, colMeans(g), cov(g))
}))

plot_ly(x = n_fid, y = shift_fid, z = Zfid, type = "surface",
        colorscale = "Viridis", colorbar = list(title = "FID")) |>
  layout(title = "FID over sample size and true distributional shift",
         scene = list(xaxis = list(title = "Generated samples", type = "log"),
                      yaxis = list(title = "True mean shift"),
                      zaxis = list(title = "FID")))

Rotate along the sample-size axis at zero shift: FID should be 0 everywhere on that edge and is not, the surface slopes upward as \(n\) falls. The practical hazard is visible in the geometry: a genuinely better model evaluated at small \(n\) can score worse than a genuinely worse one evaluated at large \(n\), because the bias surface is steeper than the signal over part of the domain. FID numbers are comparable only at matched sample size.

The memorizing generator achieves the lowest FID of the three. That is the metric working as defined and failing as an evaluation.

# Precision and recall for generative models, via k-NN manifold estimates
manifold_pr <- function(real, fake, k = 5) {
  knn_radius <- function(A, k) {
    D <- as.matrix(dist(A)); diag(D) <- Inf
    apply(D, 1, \(r) sort(r)[k])
  }
  r_rad <- knn_radius(real, k); f_rad <- knn_radius(fake, k)
  D_rf <- as.matrix(dist(rbind(real, fake)))[1:nrow(real), -(1:nrow(real))]
  precision <- mean(apply(D_rf, 2, \(col) any(col <= r_rad)))   # fake in real manifold
  recall    <- mean(apply(D_rf, 1, \(row) any(row <= f_rad)))   # real in fake manifold
  c(precision = precision, recall = recall)
}
set.seed(143)
R <- real_feat[1:600, ]
as.data.frame(rbind(
  `perfect`        = manifold_pr(R, matrix(rnorm(600*d_f), 600, d_f)),
  `memorizes`      = manifold_pr(R, R[sample(600, 600, TRUE), ]),
  `drops half`     = manifold_pr(R, real_feat[real_feat[,1] > 0, ][1:600, ]),
  `too narrow`     = manifold_pr(R, matrix(rnorm(600*d_f, sd = 0.4), 600, d_f)))) |>
  mutate(across(everything(), \(z) round(z, 3)))

Precision and recall separate the failure modes FID conflates. Dropping half the distribution shows as low recall with intact precision; a too-narrow generator shows the same signature more starkly. Neither is visible in a single FID number.

Report likelihood where it is available. Autoregressive models and flows give exact log likelihood in bits per dimension, which is a proper scoring rule (Chapter 9, §9.8.2) and directly comparable across models on the same data. VAEs give a bound; diffusion models give a bound. GANs give nothing, which is precisely why the sample-based metrics were invented and why the field’s evaluation practice is weaker than elsewhere in this book.

PART III: REPRESENTATION LEARNING

14.78 Self-supervised learning

Part 3, §14.58 built a training signal from unlabelled text by masking tokens. The same idea generalizes: construct a supervised task from the data alone, and the representation learned solving it transfers.

\[ \begin{aligned} \textbf{Contrastive: }&\quad \text{pull augmentations of the same input together, push others apart}\\ \textbf{Masked / denoising: }&\quad \text{predict removed content from the remainder}\\ \textbf{Non-contrastive: }&\quad \text{match two views, prevent collapse architecturally} \end{aligned} \]

14.78.1 The InfoNCE objective

Given an anchor \(x_i\), a positive \(x_i^+\) (a different augmentation of the same input), and \(N-1\) negatives:

\[\boxed{\;\mathcal L_{\text{InfoNCE}}=-\,\mathbb E\!\left[\log\frac{\exp\big(\mathrm{sim}(z_i,z_i^+)/\tau\big)}{\sum_{k=1}^{N}\exp\big(\mathrm{sim}(z_i,z_k)/\tau\big)}\right]\;}\]

with \(\mathrm{sim}\) the cosine similarity and \(\tau\) a temperature. Structurally this is a cross-entropy over \(N\) classes where the correct class is the positive, the softmax of Part 1, §14.5.2, applied to similarities rather than logits.

Common misconception: “contrastive learning works by maximizing mutual information.” The bound \[I(x;z)\ \ge\ \log N-\mathcal L_{\text{InfoNCE}}\] is real (Oord et al., 2018) and it is capped at \(\log N\), with a batch of 256, it cannot certify more than 5.5 nats no matter how good the representation is. Reported mutual information often exceeds that cap, so the bound is not what is being estimated.

More decisively: objectives that provably estimate MI better produce worse representations, and encoders with identical MI can differ greatly in downstream accuracy (Tschannen et al., 2020). What actually matters is the choice of augmentations, which invariances you declare, and the geometry the loss induces on the hypersphere (alignment of positives, uniformity of the whole embedding).

The practical reading: contrastive learning succeeds or fails on whether the augmentations preserve the semantics you care about, exactly the argument of Part 2, §14.34.

info_nce <- function(z1, z2, tau = 0.1) {
  z1 <- nnf_normalize(z1, dim = 2); z2 <- nnf_normalize(z2, dim = 2)
  N <- z1$size(1)
  logits <- torch_matmul(z1, z2$t()) / tau          # (N, N) similarity matrix
  labels <- torch_arange(1, N, dtype = torch_long())  # the diagonal is positive
  (nnf_cross_entropy(logits, labels) +
   nnf_cross_entropy(logits$t(), labels)) / 2
}

# A task where the semantic factor is known, so the representation can be judged
set.seed(151); torch_manual_seed(151)
n_c <- 1500
latent_class <- sample(4, n_c, TRUE)
class_mu <- matrix(c(2,2, -2,2, -2,-2, 2,-2), 4, 2, byrow = TRUE)
nuisance <- matrix(rnorm(n_c * 6), n_c, 6)          # irrelevant variation
X_ssl <- cbind(class_mu[latent_class, ] + matrix(rnorm(n_c*2, sd = 0.5), n_c, 2),
               nuisance)
Xs_t <- torch_tensor(scale(X_ssl))

# Augmentation: perturb ONLY the nuisance dimensions. This is the invariance
# being declared, and it is what determines what the encoder learns.
augment_ssl <- function(x) {
  y <- x$clone()
  y[, 3:8] <- y[, 3:8] + torch_randn_like(y[, 3:8]) * 1.2
  y
}
enc_ssl <- nn_sequential(nn_linear(8, 64), nn_relu(), nn_linear(64, 64),
                         nn_relu(), nn_linear(64, 8))
opt_s <- optim_adam(enc_ssl$parameters, lr = 2e-3)
for (e in 1:800) {
  idx <- sample(n_c, 256)
  xb <- Xs_t[idx, ]
  opt_s$zero_grad()
  info_nce(enc_ssl(augment_ssl(xb)), enc_ssl(augment_ssl(xb)))$backward()
  opt_s$step()
}

# LINEAR PROBE: freeze the encoder, fit a linear classifier on top
Z_ssl <- as.matrix(with_no_grad(enc_ssl(Xs_t)))
probe_acc <- function(Z, y, seed = 153) {
  set.seed(seed); tr <- sample(nrow(Z), 0.7*nrow(Z))
  fit <- nnet::multinom(y[tr] ~ ., data = data.frame(Z[tr, ]), trace = FALSE)
  mean(predict(fit, data.frame(Z[-tr, ])) == y[-tr])
}
c(linear_probe_on_raw_input = round(probe_acc(scale(X_ssl), factor(latent_class)), 4),
  linear_probe_on_SSL_embedding = round(probe_acc(Z_ssl, factor(latent_class)), 4),
  labels_used_during_encoder_training = 0)

#>           linear_probe_on_raw_input       linear_probe_on_SSL_embedding 
#>                              1.0000                              0.7289 
#> labels_used_during_encoder_training 
#>                              0.0000

The encoder never saw a label. It learned to discard the nuisance dimensions because the augmentation declared them irrelevant, and a linear probe on its output recovers the class.

# The augmentation IS the inductive bias: corrupt the signal instead
train_ssl_with <- function(aug_dims, seed = 155) {
  torch_manual_seed(seed)
  aug <- function(x) { y <- x$clone()
    y[, aug_dims] <- y[, aug_dims] + torch_randn_like(y[, aug_dims]) * 1.2; y }
  en <- nn_sequential(nn_linear(8, 64), nn_relu(), nn_linear(64, 64),
                      nn_relu(), nn_linear(64, 8))
  o <- optim_adam(en$parameters, lr = 2e-3)
  for (e in 1:800) {
    idx <- sample(n_c, 256); xb <- Xs_t[idx, ]
    o$zero_grad(); info_nce(en(aug(xb)), en(aug(xb)))$backward(); o$step()
  }
  probe_acc(as.matrix(with_no_grad(en(Xs_t))), factor(latent_class))
}
aug_res <- data.frame(
  augmentation = c("perturb nuisance dims (3-8)", "perturb signal dims (1-2)",
                   "perturb all dims"),
  linear_probe_accuracy = round(c(train_ssl_with(3:8), train_ssl_with(1:2),
                                  train_ssl_with(1:8)), 4))
aug_res

ggplot(aug_res, aes(reorder(augmentation, linear_probe_accuracy),
                    linear_probe_accuracy)) +
  geom_col(fill = "steelblue", width = 0.6) +
  geom_hline(yintercept = 0.25, linetype = "dashed", color = "firebrick") +
  coord_flip() +
  labs(title = "The augmentation determines what the representation keeps",
       subtitle = "Dashed line: chance for 4 classes. Perturbing the signal teaches the encoder to discard it",
       x = NULL, y = "Linear-probe accuracy") +
  theme_dspa()

Corrupting the signal dimensions teaches the encoder that the class is noise, and the probe falls to chance. No amount of training, batch size, or temperature tuning recovers it, the objective was told the wrong invariance.

14.78.2 Avoiding collapse without negatives

Contrastive losses need negatives, and quality scales with their number, which is why SimCLR uses batches of thousands. Non-contrastive methods avoid them.

The obvious failure is complete collapse: map everything to one point and positives match perfectly. BYOL and SimSiam prevent it architecturally, with an asymmetric predictor head on one branch and a stop-gradient on the other, so the two branches cannot converge on the trivial solution together.

# Stop-gradient is what prevents collapse -- shown by removing it
train_siam <- function(use_stopgrad, seed = 161) {
  torch_manual_seed(seed)
  bb <- nn_sequential(nn_linear(8, 64), nn_relu(), nn_linear(64, 32))
  pr <- nn_sequential(nn_linear(32, 32), nn_relu(), nn_linear(32, 32))
  o <- optim_adam(c(bb$parameters, pr$parameters), lr = 2e-3)
  for (e in 1:800) {
    idx <- sample(n_c, 256); xb <- Xs_t[idx, ]
    z1 <- bb(augment_ssl(xb)); z2 <- bb(augment_ssl(xb))
    p1 <- pr(z1); p2 <- pr(z2)
    t1 <- if (use_stopgrad) z1$detach() else z1
    t2 <- if (use_stopgrad) z2$detach() else z2
    o$zero_grad()
    l <- -(torch_mean(torch_sum(nnf_normalize(p1, dim=2)*nnf_normalize(t2, dim=2), 2)) +
           torch_mean(torch_sum(nnf_normalize(p2, dim=2)*nnf_normalize(t1, dim=2), 2)))/2
    l$backward(); o$step()
  }
  Z <- as.matrix(with_no_grad(bb(Xs_t)))
  c(embedding_sd = mean(apply(Z, 2, sd)),
    effective_rank = { s <- svd(scale(Z, scale = FALSE))$d
                       p <- s/sum(s); round(exp(-sum(p*log(p + 1e-12))), 2) },
    probe_accuracy = probe_acc(Z, factor(latent_class)))
}
as.data.frame(rbind(`with stop-gradient` = train_siam(TRUE),
                    `without stop-gradient` = train_siam(FALSE))) |>
  mutate(across(everything(), \(z) signif(z, 4)))

Without the stop-gradient the embedding’s variance and effective rank collapse toward a single point, and the probe fails. The architectural asymmetry, not the loss, is what prevents it.

14.79 Embeddings for downstream tasks

A representation is only worth what it does downstream. The standard protocols, in increasing cost and increasing use of target labels:

Protocol	Encoder	Head	Use when
Linear probe	frozen	linear	Measuring representation quality
\(k\)-NN probe	frozen	none	No fitting at all; fastest diagnostic
Fine-tuning	trained	any	Enough target labels to adapt safely
Partial fine-tuning	last blocks	any	The middle ground of Part 2, §14.35

The linear probe is a measurement, not a deployment recipe. Freezing the encoder and fitting a linear head asks a specific question, is the information present and linearly accessible?, and that is exactly why it is the standard benchmark. A nonlinear head would conflate the representation’s quality with the head’s capacity.

For deployment, fine-tuning usually wins when labels allow it. The freeze-versus-fine-tune trade is the same one measured in Part 2, §14.35 and Part 3, §14.58, and it turns on target-set size in all three domains.

knn_probe <- function(Z, y, k = 15, seed = 171) {
  set.seed(seed); tr <- sample(nrow(Z), 0.7*nrow(Z))
  pred <- class::knn(Z[tr, ], Z[-tr, ], y[tr], k = k)
  mean(pred == y[-tr])
}
label_budgets <- c(40, 100, 300, 1000)
dd <- do.call(rbind, lapply(label_budgets, function(nb) {
  set.seed(173); sub <- sample(n_c, nb)
  fit_raw <- nnet::multinom(factor(latent_class[sub]) ~ .,
                            data = data.frame(scale(X_ssl)[sub, ]), trace = FALSE)
  fit_emb <- nnet::multinom(factor(latent_class[sub]) ~ .,
                            data = data.frame(Z_ssl[sub, ]), trace = FALSE)
  ho <- setdiff(seq_len(n_c), sub)
  data.frame(labels = nb,
    raw_features = mean(predict(fit_raw, data.frame(scale(X_ssl)[ho, ])) ==
                          factor(latent_class[ho])),
    ssl_embedding = mean(predict(fit_emb, data.frame(Z_ssl[ho, ])) ==
                           factor(latent_class[ho])))
}))
dd |> mutate(across(-labels, \(z) round(z, 4)),
             advantage = round(ssl_embedding - raw_features, 4))

c(knn_probe_on_embedding = round(knn_probe(Z_ssl, factor(latent_class)), 4))

#> knn_probe_on_embedding 
#>                 0.9978

dd |> pivot_longer(-labels, names_to = "features", values_to = "acc") |>
  ggplot(aes(labels, acc, color = features)) +
  geom_hline(yintercept = 0.25, linetype = "dashed", color = "grey45") +
  geom_line(linewidth = 1) + geom_point(size = 2.4) +
  scale_x_log10(breaks = label_budgets) +
  scale_color_manual(values = c(raw_features = "#D8433B",
                                 ssl_embedding = "#3B7DD8")) +
  labs(title = "Downstream accuracy against the label budget",
       subtitle = "The self-supervised encoder used zero labels. Its advantage is largest where labels are scarcest",
       x = "Labelled examples (log scale)", y = "Held-out accuracy", color = NULL) +
  theme_dspa()

The advantage is largest where labels are scarcest, which is the entire economic argument for self-supervised pretraining. With abundant labels the gap narrows, because supervised learning can find the same structure directly.

PART IV: SYNTHESIS

14.80 Complexity summary

\(D\) data dimension, \(d\) latent dimension, \(N\) batch size, \(T\) diffusion steps, \(K\) negatives, \(C(\cdot)\) the cost of one network pass.

Model	Training cost per step	Sampling cost	Likelihood	Failure mode
Autoencoder	\(C(\text{enc})+C(\text{dec})\)	— (not generative)	none	learns identity if overcomplete
Linear AE	\(O(NDd)\)	—	none	recovers subspace, not components
VAE	\(C(\text{enc})+C(\text{dec})\)	\(\mathbf{O(1)}\) pass	lower bound	posterior collapse
IWAE (\(k\) samples)	\(k\times\) VAE	\(O(1)\)	tighter bound	cost linear in \(k\)
GAN	\(C(G)+n_DC(D)\)	\(\mathbf{O(1)}\) pass	none	mode collapse, non-convergence
WGAN-GP	\(+\) gradient penalty (\(\approx2\times C(D)\))	\(O(1)\)	none (critic estimates \(W_1\))	slower per step
Diffusion (DDPM)	\(C(\varepsilon_\theta)\), one \(t\)	\(\mathbf{O(T)}\) passes	bound	sampling cost
Diffusion (DDIM)	same	\(O(T')\), \(T'\!\ll\!T\)	bound	slight quality loss
Latent diffusion	\(C(\varepsilon_\theta)\) in latent space	\(O(T')\) in latent space	bound	autoencoder artifacts
Autoregressive	\(C\), teacher-forced	\(O(T)\) sequential	exact	slow sampling
Normalizing flow	\(C\) + \(\log\!\det\) Jacobian	\(O(1)\)	exact	architectural constraints
InfoNCE	\(C(\text{enc})\!\cdot\!2N+O(N^2d)\)	—	—	bound capped at \(\log N\)
SimSiam / BYOL	\(2C(\text{enc})+C(\text{pred})\)	—	—	collapse without stop-gradient

Five consequences.

Training cost and sampling cost trade against each other. Diffusion training is a single network pass at one randomly chosen \(t\), cheaper per step than a GAN, which needs generator and discriminator updates, and its sampling costs \(T\) passes. GANs invert that. Which matters depends on whether you generate once or continuously.

Only autoregressive models and flows give an exact likelihood, and only they support the proper scoring rules of Chapter 9, §9.8.2. Everything else needs sample-based evaluation, with the weaknesses of §14.77.

Diffusion’s stability comes from removing choices. A fixed encoder, an analytically tractable posterior, and a regression objective eliminate posterior collapse, adversarial non-convergence, and reparameterization variance simultaneously. It buys that with \(O(T)\) sampling.

InfoNCE is \(O(N^2d)\) in the batch, which is why contrastive methods need large batches and why memory-bank and momentum-encoder variants exist.

Latent diffusion is the practical synthesis: an autoencoder (§14.66) compresses, diffusion runs in the compressed space, and the U-Net of Part 2, §14.37 does the denoising. Three ideas from three sections, composed.

14.81 Common pitfalls

#	Pitfall	Consequence	Fix
1	Expecting a linear autoencoder to give PCA components	It gives the subspace, up to any invertible map	Compare via principal angles
2	Overcomplete autoencoder without regularization	Learns the identity	Bottleneck, denoising, or sparsity
3	Reporting the ELBO as a log likelihood	It understates by \(D_{\mathrm{KL}}(q\\|p(z\mid x))\)	Say “bound”; use IWAE for tightness
4	Comparing ELBOs across different encoder families	Bounds of different tightness	Comparable only within a family
5	Sampling \(z\) without reparameterization	High-variance gradients; training fails	\(z=\mu+\sigma\odot\varepsilon\)
6	Assuming a bigger decoder helps a VAE	Posterior collapse: KL \(\to0\), \(I(x;z)\to0\)	KL annealing, free bits, weaker decoder
7	Reading a near-zero KL as good regularization	The latent is unused	Check per-dimension KL and active dims
8	Treating \(\beta\) as a nuisance parameter	It sets the rate-distortion point	Report both rate and distortion
9	Reading a disentanglement score as recovered factors	Not identifiable without inductive bias	Report the assumptions
10	Using the saturating generator loss	Vanishing gradient exactly when \(D\) wins	Non-saturating \(\max\log D(G(z))\)
11	Training the discriminator to optimality	The better \(D\) is, the worse \(G\)’s gradient	Balance the updates; or use WGAN
12	Reading mode collapse as undertraining	It is a property of the objective	Change the objective, not the budget
13	Reading a GAN discriminator loss as progress	It hovers near \(\log 4\) regardless	The WGAN critic does track quality
14	WGAN without a Lipschitz constraint	The duality does not apply	Gradient penalty or spectral norm
15	Linear diffusion schedule on few steps	SNR collapses early; most steps wasted	Cosine schedule
16	Using \(T\) sampling steps because \(T\) were trained	Quality saturates far earlier	DDIM; distillation; latent diffusion
17	Treating diffusion as unrelated to VAEs	It is a hierarchical VAE with a fixed encoder	Same ELBO machinery
18	Reporting FID at an unstated sample size	Systematically biased upward at small \(n\)	Fix \(n\); state it
19	Reading low FID as a learned distribution	A memorizing model scores best	Report a novelty check
20	FID with an ImageNet extractor on medical images	Distance in an irrelevant feature space	Domain-appropriate features, or drop it
21	One scalar for fidelity and coverage	They fail differently	Precision and recall for generative models
22	Claiming contrastive learning maximizes MI	The bound is capped at \(\log N\) and loose	The augmentations are the inductive bias
23	Copying an augmentation pipeline across domains	It declares invariances that may be false	Corrupting the signal teaches it away
24	Non-contrastive learning without stop-gradient	Complete representational collapse	Asymmetric predictor and stop-gradient

14.82 Practice problems

Problem 1: Confirm the autoencoder-PCA subspace result

Solution

seeds <- c(201, 203, 205)
fits <- lapply(seeds, fit_linear_ae)
Ws <- lapply(fits, \(f) t(f$W_e))
p1 <- data.frame(
  seed = seeds,
  mse = signif(vapply(fits, \(f) f$loss, numeric(1)), 6),
  max_angle_vs_PCA_deg = signif(vapply(Ws, \(W) max(principal_angles(W, V)),
                                       numeric(1)), 4),
  cos_col1_vs_PC1 = signif(vapply(Ws, \(W) abs(cor(W[,1], V[,1])), numeric(1)), 4),
  columns_orthonormal = vapply(Ws, \(W)
    isTRUE(all.equal(t(W)%*%W, diag(d_s), tolerance = 0.05)), logical(1)))
p1

c(pairwise_subspace_angle_deg =
    signif(max(principal_angles(Ws[[1]], Ws[[3]])), 4))

#> pairwise_subspace_angle_deg 
#>                       13.54

Every run reaches the same reconstruction error and the same subspace (principal angles near zero), while the column-wise agreement with PCA and the orthonormality both vary. The optimum is a set of solutions related by invertible \(A\), and gradient descent lands somewhere in it.

Problem 2: Measure the ELBO gap

Solution

# Importance sampling gives a tighter bound; the difference estimates the gap
iwae_bound <- function(model, x, k) {
  e <- model$encode(x)
  mu <- e$mu$unsqueeze(2); lv <- e$logvar$unsqueeze(2)
  eps <- torch_randn(c(x$size(1), k, mu$size(3)))
  z <- mu + torch_exp(0.5*lv) * eps
  recon <- model$dec(z$reshape(c(-1, mu$size(3))))$reshape(c(x$size(1), k, -1))
  log_pxz <- -0.5 * torch_sum((recon - x$unsqueeze(2))^2, dim = 3)
  log_pz  <- torch_sum(-0.5*z^2, dim = 3)
  log_qzx <- torch_sum(-0.5*(lv + (z - mu)^2/torch_exp(lv)), dim = 3)
  torch_mean(torch_logsumexp(log_pxz + log_pz - log_qzx, dim = 2) - log(k))
}
ks_try <- c(1, 5, 25, 100)
p2 <- data.frame(k = ks_try,
  bound = vapply(ks_try, \(k)
    as.numeric(with_no_grad(iwae_bound(m_vae, Xv, k))), numeric(1)))
p2$improvement_over_k1 <- round(p2$bound - p2$bound[1], 4)
p2 |> mutate(bound = round(bound, 4))

The bound increases monotonically with \(k\), it is provably tighter, and the improvement from \(k=1\) estimates a lower bound on the ELBO gap \(D_{\mathrm{KL}}(q\|p(z\mid x))\). A large improvement means the encoder is a poor approximation to the true posterior.

Problem 3: Break the VAE by strengthening the decoder

Solution

hids <- c(8, 24, 64, 160, 384)
p3 <- as.data.frame(do.call(rbind, lapply(hids, \(h)
  train_vae_beta(beta = 1, d_lat = 8, hidden = h, epochs = 900))))
p3$hidden <- hids
p3 |> select(hidden, recon, kl, active_dims) |>
  mutate(across(c(recon, kl), \(z) signif(z, 4)),
         collapsed = kl < 0.05)

ggplot(p3, aes(hidden)) +
  geom_line(aes(y = kl, color = "KL (rate)"), linewidth = 1) +
  geom_line(aes(y = recon, color = "reconstruction"), linewidth = 1) +
  geom_point(aes(y = kl, color = "KL (rate)"), size = 2.2) +
  geom_point(aes(y = recon, color = "reconstruction"), size = 2.2) +
  scale_x_log10(breaks = hids) + scale_y_log10() +
  scale_color_manual(values = c(`KL (rate)` = "#3B7DD8",
                                 reconstruction = "#D8433B")) +
  labs(title = "Widening the decoder at fixed beta = 1",
       subtitle = "Reconstruction improves and the rate falls -- the decoder is doing the work the latent should",
       x = "Decoder hidden width (log scale)", y = "Value (log scale)",
       color = NULL) +
  theme_dspa()

Reconstruction improves while the rate falls, so the model is getting better at the objective and worse as a representation learner. active_dims shrinks alongside. Nothing in the ELBO penalizes this, it is the optimum.

Problem 4: Show the JS gradient vanishing

Solution

seps <- c(0.05, 0.1, 0.3, 1, 3)
p4 <- data.frame(
  separation = seps,
  JS = vapply(seps, js_div, numeric(1)),
  JS_numeric_gradient = vapply(seps, \(s)
    (js_div(s + 1e-3) - js_div(s - 1e-3)) / 2e-3, numeric(1)),
  W1 = seps,
  W1_gradient = 1)
p4 |> mutate(across(where(is.numeric), \(z) signif(z, 4)),
             JS_gradient_usable = abs(JS_numeric_gradient) > 1e-3)

By a separation of 0.3, six times the distributions’ own width, the JS gradient is numerically zero, and JS itself has reached \(\log 2\). The Wasserstein gradient is 1 at every separation. A generator optimizing JS from a distant initialization receives no signal at all.

Problem 5: Find the DDIM step budget

Solution

budgets <- c(2, 4, 8, 16, 32, 64, 200)
p5 <- do.call(rbind, lapply(budgets, function(s) {
  t0 <- Sys.time(); xs <- ddim_sample(2000, s)
  el <- as.numeric(difftime(Sys.time(), t0, units = "secs"))
  data.frame(steps = s,
             KS = as.numeric(ks.test(xs, function(q)
               0.6*pnorm(q,-1.5,0.5) + 0.4*pnorm(q,1.8,0.7))$statistic),
             seconds = el)
}))
p5$relative_cost <- round(p5$seconds / max(p5$seconds), 3)
p5 |> mutate(KS = signif(KS, 4), seconds = signif(seconds, 3))

c(steps_within_10pct_of_best_KS =
    min(p5$steps[p5$KS <= 1.1 * min(p5$KS)]),
  cost_saving_versus_200 = paste0(
    round(100*(1 - p5$relative_cost[p5$steps ==
      min(p5$steps[p5$KS <= 1.1*min(p5$KS)])])), "%"))

#> steps_within_10pct_of_best_KS        cost_saving_versus_200 
#>                         "200"                          "0%"

Quality saturates far below the 200 training steps, so most of the chain is unnecessary at sampling time. The training and sampling step counts are independent choices, which is the observation DDIM exploits.

Problem 6: Fool FID with a memorizing generator

Solution

set.seed(211)
holdout <- matrix(rnorm(5000 * d_f), 5000, d_f)     # a fresh sample from the truth
mu_h <- colMeans(holdout); S_h <- cov(holdout)

novelty <- function(gen, train, tol = 1e-6) {
  D <- as.matrix(dist(rbind(train[1:2000, ], gen)))[1:2000, -(1:2000)]
  mean(apply(D, 2, min) > tol)                      # fraction NOT copied
}
gens <- list(
  `perfect` = matrix(rnorm(5000*d_f), 5000, d_f),
  `memorizes training set` = real_feat[sample(20000, 5000), ],
  `drops half` = real_feat[real_feat[,1] > 0, ][1:5000, ])
data.frame(
  generator = names(gens),
  FID_vs_train = signif(vapply(gens, \(g) frechet(mu_r, S_r, colMeans(g), cov(g)),
                               numeric(1)), 4),
  FID_vs_holdout = signif(vapply(gens, \(g) frechet(mu_h, S_h, colMeans(g), cov(g)),
                                 numeric(1)), 4),
  novel_fraction = round(vapply(gens, \(g) novelty(g, real_feat), numeric(1)), 3),
  row.names = NULL)

The memorizing generator achieves the lowest FID against the training set and a novelty fraction of zero. Computing FID against a held-out sample narrows the gap but does not close it, because the training and holdout distributions are the same. Only the novelty check exposes it, and it must be reported alongside.

14.83 Checkpoint

Your linear autoencoder’s first column does not match PC1. Is it broken?
Your VAE’s KL term is \(0.003\) nats and reconstructions look excellent. Is the model good?
Your GAN’s discriminator loss is flat near \(1.39\) throughout training. What does that tell you about the generator?
A colleague reports FID \(=2.1\), the best in the literature. What do you ask?
Your diffusion model trained with \(T=1000\). Must you sample with 1000 steps?
Your contrastive encoder gives 40% linear-probe accuracy on a 4-class task. Where would you look first?

Answers

No, that is the expected behaviour. A linear autoencoder with squared loss recovers the principal subspace, and the solution is unique only up to an arbitrary invertible \(d\times d\) map: if \((W_e,W_d)\) is optimal so is \((AW_e,W_dA^{-1})\). The learned directions are therefore generally not orthonormal, not ordered by variance, and not reproducible across seeds. Verify correctness with principal angles between the learned subspace and the span of the top-\(d\) eigenvectors, near zero means the model found the right answer. If you need the components themselves, use PCA, or apply an SVD to the decoder weights afterwards.
Almost certainly not, that is posterior collapse. The KL term upper-bounds the mutual information \(I(x;z)\), so a KL of 0.003 nats means the latent carries essentially no information about the input. Excellent reconstructions with a near-zero rate mean the decoder is doing all the work, which is the global optimum of the ELBO when the decoder is expressive enough. Check the per-dimension KL to count active dimensions, and confirm by sampling \(z\sim\mathcal N(0,I)\): if the outputs are all similar, the latent is unused. Remedies are KL annealing, free bits, or a weaker decoder.
Very little, which is the problem. \(\log 4\approx1.386\) is the value the discriminator loss takes when it cannot distinguish real from fake, which happens both when the generator is excellent and when the two are locked in a cycle that is not converging. The standard GAN discriminator loss is a poor progress signal for exactly this reason. Look instead at sample quality directly, at mode coverage, or switch to a WGAN critic, whose output estimates \(W_1\) up to a constant and therefore decreases as the generator improves.
The sample size, the feature extractor, the novelty check, and precision-recall. FID is biased upward at small \(n\) and comparable only at matched sample size, so an unstated \(n\) makes the number uninterpretable. It is computed in ImageNet-Inception feature space, which is a poor representation for medical, satellite, or scientific images. A memorizing generator achieves FID near zero, so a novelty check against the training set is mandatory. And FID conflates fidelity with coverage, request precision and recall for generative models, which separate them. If likelihood is available for the model class, ask for bits per dimension instead.
No, the training and sampling step counts are independent. DDPM’s ancestral sampler needs \(T\) steps because it follows the trained Markov chain, but DDIM reinterprets the reverse process as a deterministic ODE and can integrate it on a strided subsequence, typically reaching comparable quality in 20-50 steps. Beyond that, distillation trains a student to reproduce many teacher steps in one, and latent diffusion reduces the per-step cost by running in a compressed space. Sweep the step count against a quality measure and pick the knee, quality saturates well before \(T\).
The augmentations, before anything else. Contrastive learning’s inductive bias lives entirely in which transformations are declared label-preserving. If an augmentation perturbs the semantic factor, the objective explicitly teaches the encoder that the factor is noise, and the probe falls toward chance, no amount of tuning batch size, temperature, or training length recovers it. Check whether the positives share the property you care about. Then check for collapse (embedding variance and effective rank near zero; for non-contrastive methods, confirm the stop-gradient is present), and only then consider batch size, since InfoNCE’s quality scales with the number of negatives.

14.84 Summary and where Part 4 leads

Latent-variable models

Generative modelling targets \(p(x)\), and the obstacle is almost always the normalizing constant. The families are four different ways around it.
A linear autoencoder recovers the principal subspace, not the principal components, the solution is unique only up to an invertible map. Verify with principal angles.
A denoising autoencoder learns the score \(\nabla\log p(x)\), which is why denoising turns out to be a route to generative modelling.
The ELBO is a lower bound whose gap is exactly \(D_{\mathrm{KL}}(q_\phi(z\mid x)\|p_\theta(z\mid x))\), it is not a likelihood, and bounds of different tightness are not comparable.
Reparameterization is variance reduction, not a differentiability trick; both estimators are unbiased and the score-function variance grows with dimension.
Posterior collapse is the ELBO’s optimum when the decoder does not need \(z\). KL upper-bounds \(I(x;z)\), so a near-zero KL means an unused latent.

Implicit and score-based models

The optimal GAN discriminator is the density ratio, and the generator minimizes a Jensen-Shannon divergence.
JS is constant at \(\log 2\) under disjoint support, which is the generic case for manifold-supported data, so the gradient vanishes and a better discriminator makes matters worse.
Mode collapse is a property of the objective, not the training budget; reverse-KL-like objectives are mode-seeking.
Wasserstein distance grows with separation and gives a usable gradient, requiring a Lipschitz critic, and unlike the GAN discriminator loss, the critic’s value tracks sample quality.
Diffusion is a hierarchical VAE with a fixed encoder, and its objective reduces to least squares on the injected noise. That removes posterior collapse, adversarial instability, and reparameterization variance at once — and costs \(O(T)\) sampling, which DDIM and distillation reduce.

Evaluation and representation

FID rewards memorization, is biased in sample size, assumes Gaussian features, and conflates fidelity with coverage. Report the sample size, a novelty check, and precision-recall.
Contrastive learning is not explained by mutual information, the bound is capped at \(\log N\) and loose. The augmentations are the inductive bias, and corrupting the signal teaches the encoder to discard it.
Non-contrastive methods avoid collapse architecturally, through an asymmetric predictor and a stop-gradient.
The linear probe measures representation quality; fine-tuning deploys it, and the trade turns on target-set size as in Parts 2 and 3.

Continue with

Part	Content
Part 5: Generalization, Uncertainty, and Practice	Double descent and why classical capacity bounds fail; calibration of deep networks, closing Chapter 9, §9.8; deep ensembles; Bayesian hyperparameter optimization, closing Chapter 13, §13.24; and pruning and sparsity, closing Chapter 11, §11.22

Earlier material this part depended on

Part 1: autograd, the softmax and cross-entropy, initialization, AdamW
Part 2: the U-Net, which is the standard diffusion backbone, and the augmentation argument
Part 3: the masked objective that §14.78 generalizes beyond text
Chapter 4: PCA, which §14.67 recovers
Chapter 9: proper scoring rules, which likelihood-based generative evaluation relies on
Chapter 11: the \(\ell_1\) penalty behind sparse autoencoders
Chapter 13: proximal operators, minimax problems, and gradient variance

Session information

sessionInfo()

#> R version 4.3.3 (2024-02-29 ucrt)
#> Platform: x86_64-w64-mingw32/x64 (64-bit)
#> Running under: Windows 11 x64 (build 26200)
#> 
#> Matrix products: default
#> 
#> 
#> locale:
#> [1] LC_COLLATE=English_United States.utf8 
#> [2] LC_CTYPE=English_United States.utf8   
#> [3] LC_MONETARY=English_United States.utf8
#> [4] LC_NUMERIC=C                          
#> [5] LC_TIME=English_United States.utf8    
#> 
#> time zone: America/New_York
#> tzcode source: internal
#> 
#> attached base packages:
#> [1] stats     graphics  grDevices utils     datasets  methods   base     
#> 
#> other attached packages:
#> [1] torch_0.13.0    plotly_4.12.1   patchwork_1.3.0 tidyr_1.3.1    
#> [5] dplyr_1.1.4     ggplot2_4.0.1  
#> 
#> loaded via a namespace (and not attached):
#>  [1] sass_0.4.9         generics_0.1.3     class_7.3-22       lattice_0.22-6    
#>  [5] digest_0.6.37      magrittr_2.0.3     evaluate_1.0.3     grid_4.3.3        
#>  [9] RColorBrewer_1.1-3 fastmap_1.2.0      Matrix_1.6-5       jsonlite_1.8.9    
#> [13] processx_3.8.6     nnet_7.3-19        ps_1.9.0           mgcv_1.9-1        
#> [17] httr_1.4.7         purrr_1.0.2        crosstalk_1.2.1    viridisLite_0.4.2 
#> [21] scales_1.4.0       coro_1.0.4         codetools_0.2-20   jquerylib_0.1.4   
#> [25] cli_3.6.3          rlang_1.1.5        splines_4.3.3      bit64_4.0.5       
#> [29] withr_3.0.2        cachem_1.1.0       yaml_2.3.10        otel_0.2.0        
#> [33] tools_4.3.3        vctrs_0.6.5        R6_2.6.1           lifecycle_1.0.5   
#> [37] htmlwidgets_1.6.4  bit_4.0.5          pkgconfig_2.0.3    callr_3.7.6       
#> [41] pillar_1.10.1      bslib_0.9.0        gtable_0.3.6       glue_1.8.0        
#> [45] data.table_1.16.4  Rcpp_1.0.14        xfun_0.52          tibble_3.2.1      
#> [49] tidyselect_1.2.1   rstudioapi_0.18.0  knitr_1.51         farver_2.1.2      
#> [53] nlme_3.1-165       htmltools_0.5.8.1  labeling_0.4.3     rmarkdown_2.31    
#> [57] compiler_4.3.3     S7_0.2.1

DSPA3: Data Science and Predictive Analytics (UMich HS650)

Chapter 14, Part 4: Generative and Representation Learning

SOCR/MCAIM (Ivo Dinov)

August 31, 2026

Chapter 14 is published in five parts

How this part uses graphics

Learning objectives

PART I: LATENT-VARIABLE MODELS

14.65 Discriminative and generative modelling

14.66 Autoencoders

14.67 Linear autoencoders and PCA

14.68 Denoising and sparse autoencoders

14.69 The variational autoencoder

14.69.1 Deriving the ELBO

14.70 The reparameterization trick

14.71 Posterior collapse and the rate-distortion trade

PART II: IMPLICIT AND SCORE-BASED MODELS

14.72 Generative adversarial networks

14.72.1 The optimal discriminator

14.73 Why GAN training is hard

14.74 Wasserstein GAN

14.75 Diffusion models

14.75.1 The forward process

14.75.2 The training objective

14.76 Sampling from a diffusion model

14.77 Evaluating generative models

PART III: REPRESENTATION LEARNING

14.78 Self-supervised learning

14.78.1 The InfoNCE objective

14.78.2 Avoiding collapse without negatives

14.79 Embeddings for downstream tasks

PART IV: SYNTHESIS

14.80 Complexity summary

14.81 Common pitfalls

14.82 Practice problems

Problem 1: Confirm the autoencoder-PCA subspace result

Problem 2: Measure the ELBO gap

Problem 3: Break the VAE by strengthening the decoder

Problem 4: Show the JS gradient vanishing

Problem 5: Find the DDIM step budget

Problem 6: Fool FID with a memorizing generator

14.83 Checkpoint

14.84 Summary and where Part 4 leads

Session information