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 -

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