import random
import matplotlib.pyplot as plt
import numpy as np
%load_ext autoreload
%autoreload 2
SEED = 42
random.seed(SEED)
rng = np.random.default_rng(SEED)
plt.rcParams["figure.dpi"] = 90
Exercise 1: KL Divergence of a Diagonal Gaussian (Hand Calculation)¶
Task: Let $q = \mathcal{N}(\mu, \Sigma)$ with $\mu = [1.0, -1.0]$ and $\Sigma = \text{diag}(e^{0.5}, e^{-0.2})$ (so the log-variances are $[0.5, -0.2]$), and let $p = \mathcal{N}(0, I)$.
- Hand-compute $D_{\text{KL}}(q \| p)$ using the closed form
$D_{\text{KL}} = -\frac{1}{2} \sum_j \left(1 + \log\sigma_j^2 - \mu_j^2 - \sigma_j^2\right)$.
This is exactly the KL term of the VAE ELBO (the term
vae_elbo_lossinml_first_principlesadds to the reconstruction loss — the full decomposition is verified in first_principles.ipynb). - Verify the closed form in code.
- Verify it a second, independent way with a Monte-Carlo estimate $D_{\text{KL}} \approx \frac{1}{N}\sum_{i=1}^N \left[\log q(z_i) - \log p(z_i)\right]$ with $z_i \sim q$.
Solution 1 (Hand Calculation)¶
$D_{\text{KL}} = -\frac{1}{2} \left[ (1 + 0.5 - 1.0^2 - e^{0.5}) + (1 - 0.2 - (-1.0)^2 - e^{-0.2}) \right]$
With $e^{0.5} \approx 1.6487$ and $e^{-0.2} \approx 0.8187$:
- Component 1: $1 + 0.5 - 1 - 1.6487 = -1.1487$
- Component 2: $1 - 0.2 - 1 - 0.8187 = -1.0187$
Result: $D_{\text{KL}} = -\frac{1}{2}(-1.1487 - 1.0187) \approx 1.0837$.
mu = np.array([1.0, -1.0])
logvar = np.array([0.5, -0.2])
def kl_closed_form(mu, logvar):
"""KL( N(mu, diag(exp(logvar))) || N(0, I) ) — the ELBO KL term."""
return -0.5 * np.sum(1 + logvar - mu**2 - np.exp(logvar))
kl_cf = kl_closed_form(mu, logvar)
print(f"Closed-form KL: {kl_cf:.4f}")
assert np.allclose(kl_cf, 1.0837, atol=1e-3)
# Monte-Carlo estimate: KL = E_q[log q(z) - log p(z)], z ~ q
n = 200_000
std = np.exp(0.5 * logvar)
z = mu + std * rng.standard_normal((n, 2))
def log_gaussian(z, mu, logvar):
return -0.5 * np.sum(np.log(2 * np.pi) + logvar + (z - mu) ** 2 / np.exp(logvar), axis=1)
kl_mc = np.mean(log_gaussian(z, mu, logvar) - log_gaussian(z, np.zeros(2), np.zeros(2)))
print(f"Monte-Carlo KL: {kl_mc:.4f} ({n} samples)")
assert np.allclose(kl_mc, kl_cf, atol=0.02)
Closed-form KL: 1.0837 Monte-Carlo KL: 1.0792 (200000 samples)
Exercise 2: Reparameterization Trick Gradients (Implementation)¶
Task: Implement the forward pass of $z = \mu + \sigma \odot \epsilon$ with $\sigma = e^{\frac{1}{2}\log\sigma^2}$. Then compute the analytic gradients $\frac{\partial z}{\partial \mu}$ and $\frac{\partial z}{\partial \log\sigma^2}$ and verify them against finite differences with an explicit tolerance.
mu = np.array([0.5, 0.1])
logvar = np.array([-0.1, 0.2])
eps = np.array([0.4, -0.6])
# 1. Forward
std = np.exp(0.5 * logvar)
z = mu + std * eps
# 2. Analytic gradients
dz_dmu = np.ones_like(mu)
dz_dlogvar = 0.5 * std * eps
# 3. Finite differences
h = 1e-5
dz_dmu_fd = ((mu + h + std * eps) - z) / h
std_h = np.exp(0.5 * (logvar + h))
dz_dlogvar_fd = ((mu + std_h * eps) - z) / h
print(f"Analytic dz_dlogvar: {dz_dlogvar}")
print(f"FD dz_dlogvar: {dz_dlogvar_fd}")
assert np.allclose(dz_dmu, dz_dmu_fd, atol=1e-4)
assert np.allclose(dz_dlogvar, dz_dlogvar_fd, atol=1e-4)
Analytic dz_dlogvar: [ 0.19024588 -0.33155128] FD dz_dlogvar: [ 0.19024636 -0.3315521 ]
Exercise 3: Optimal Discriminator and the Value of the Game (Hand Calculation)¶
Task: Consider distributions supported on two points $\lbrace a, b \rbrace$: $p_{\text{data}} = [0.8, 0.2]$ and $p_g = [0.3, 0.7]$.
- Hand-compute the optimal discriminator $D^\ast(x) = \dfrac{p_{\text{data}}(x)}{p_{\text{data}}(x) + p_g(x)}$ at both points.
- Hand-compute the value of the game $V(G, D^\ast) = \mathbb{E}_{p_{\text{data}}}[\log D^\ast(x)] + \mathbb{E}_{p_g}[\log(1 - D^\ast(x))]$ and confirm the theory.md identity $V(G, D^\ast) = -\log 4 + 2 \cdot \text{JSD}(p_{\text{data}} \Vert p_g)$.
- Verify both in code, including a check against
gan_discriminator_lossfromml_first_principles(which computes $-V$ as a mean over samples).
Solution 3 (Hand Calculation)¶
Optimal discriminator:
- $D^\ast(a) = \frac{0.8}{0.8 + 0.3} = \frac{8}{11} \approx 0.7273$
- $D^\ast(b) = \frac{0.2}{0.2 + 0.7} = \frac{2}{9} \approx 0.2222$
Value of the game (natural log):
$V = 0.8 \log\tfrac{8}{11} + 0.2 \log\tfrac{2}{9} + 0.3 \log\tfrac{3}{11} + 0.7 \log\tfrac{7}{9}$ $\approx -0.2548 - 0.3008 - 0.3898 - 0.1759 = -1.1213$
JSD check with mixture $m = \frac{1}{2}(p_{\text{data}} + p_g) = [0.55, 0.45]$:
- $D_{\text{KL}}(p_{\text{data}} \Vert m) \approx 0.1376$
- $D_{\text{KL}}(p_g \Vert m) \approx 0.1274$
- $\text{JSD} = \frac{1}{2}(0.1376 + 0.1274) \approx 0.1325$
Result: $-\log 4 + 2 \cdot 0.1325 = -1.3863 + 0.2650 \approx -1.1213 = V$. The identity holds.
from ml_first_principles.generative_models import gan_discriminator_loss
p_data = np.array([0.8, 0.2])
p_g = np.array([0.3, 0.7])
d_star = p_data / (p_data + p_g)
print(f"D*(a) = {d_star[0]:.4f}, D*(b) = {d_star[1]:.4f}")
V = np.sum(p_data * np.log(d_star)) + np.sum(p_g * np.log(1 - d_star))
print(f"V(G, D*) = {V:.4f}")
assert np.allclose(V, -1.1213, atol=1e-4)
# Identity: V = -log 4 + 2 * JSD(p_data || p_g)
m = 0.5 * (p_data + p_g)
kl = lambda p, q: np.sum(p * np.log(p / q))
jsd = 0.5 * (kl(p_data, m) + kl(p_g, m))
print(f"JSD = {jsd:.4f}, -log4 + 2*JSD = {-np.log(4) + 2 * jsd:.4f}")
assert np.allclose(V, -np.log(4) + 2 * jsd, atol=1e-12)
# Library check: gan_discriminator_loss averages over samples, so encode the
# probabilities as sample counts (real: 8 a's + 2 b's; fake: 3 a's + 7 b's).
d_real = np.repeat(d_star, [8, 2])
d_fake = np.repeat(d_star, [3, 7])
assert np.allclose(-gan_discriminator_loss(d_real, d_fake), V, atol=1e-9)
print("Matches -gan_discriminator_loss at the optimal discriminator.")
D*(a) = 0.7273, D*(b) = 0.2222 V(G, D*) = -1.1213 JSD = 0.1325, -log4 + 2*JSD = -1.1213 Matches -gan_discriminator_loss at the optimal discriminator.
Exercise 4: DDPM Forward-Noising Marginal (Implementation)¶
Task: theory.md derives the closed-form marginal of the DDPM forward process: $q(x_t \mid x_0) = \mathcal{N}\left(x_t;\ \sqrt{\bar{\alpha}_t}\, x_0,\ (1 - \bar{\alpha}_t) I\right)$ with $\alpha_t = 1 - \beta_t$ and $\bar{\alpha}_t = \prod_{s=1}^t \alpha_s$.
For the schedule $\beta = [0.1, 0.2, 0.3]$:
- Hand-compute $\bar{\alpha}_t$ for $t = 1, 2, 3$.
- Implement
ddpm_forward_marginal(x0, t, alpha_bar, eps)returning $x_t = \sqrt{\bar{\alpha}_t}\, x_0 + \sqrt{1 - \bar{\alpha}_t}\, \epsilon$. - Verify deterministically (fixed seed, explicit
atol) that the sample mean and variance of $x_t$ match the closed form, and that composing two single noising steps $q(x_t \mid x_{t-1})$ gives the same mean/variance as the $t=2$ marginal. - Check the limiting behavior: for a standard long schedule, $\bar{\alpha}_T \to 0$, so $x_T$ is (almost) pure noise.
Solution 4 (Hand Calculation)¶
$\alpha = 1 - \beta = [0.9, 0.8, 0.7]$, so:
- $\bar{\alpha}_1 = 0.9$
- $\bar{\alpha}_2 = 0.9 \times 0.8 = 0.72$
- $\bar{\alpha}_3 = 0.72 \times 0.7 = 0.504$
At $t = 3$ the marginal is $\mathcal{N}\left(\sqrt{0.504}\, x_0,\ 0.496\, I\right)$: about half of the signal variance has been replaced by noise.
betas = np.array([0.1, 0.2, 0.3])
alphas = 1.0 - betas
alpha_bar = np.cumprod(alphas)
assert np.allclose(alpha_bar, [0.9, 0.72, 0.504], atol=1e-12)
def ddpm_forward_marginal(x0, t, alpha_bar, eps):
"""Sample x_t ~ q(x_t | x_0) = N(sqrt(abar_t) x0, (1 - abar_t) I); t is 1-based."""
ab = alpha_bar[t - 1]
return np.sqrt(ab) * x0 + np.sqrt(1.0 - ab) * eps
x0 = np.array([2.0, -1.0])
n = 200_000
for t in (1, 2, 3):
eps = rng.standard_normal((n, 2))
xt = ddpm_forward_marginal(x0, t, alpha_bar, eps)
mean_err = np.max(np.abs(xt.mean(axis=0) - np.sqrt(alpha_bar[t - 1]) * x0))
assert np.allclose(xt.mean(axis=0), np.sqrt(alpha_bar[t - 1]) * x0, atol=0.01)
assert np.allclose(xt.var(axis=0), (1 - alpha_bar[t - 1]) * np.ones(2), atol=0.01)
print(f"t={t}: abar={alpha_bar[t - 1]:.3f} mean matches (max err {mean_err:.4f}), var matches")
# Composing two single steps q(x_1|x_0), q(x_2|x_1) reproduces the t=2 marginal
x1 = np.sqrt(alphas[0]) * x0 + np.sqrt(betas[0]) * rng.standard_normal((n, 2))
x2 = np.sqrt(alphas[1]) * x1 + np.sqrt(betas[1]) * rng.standard_normal((n, 2))
assert np.allclose(x2.mean(axis=0), np.sqrt(alpha_bar[1]) * x0, atol=0.01)
assert np.allclose(x2.var(axis=0), (1 - alpha_bar[1]) * np.ones(2), atol=0.01)
print("Two composed single steps match the t=2 closed-form marginal.")
# Limiting behavior: standard linear schedule drives abar_T -> 0 (pure noise)
T = 1000
abar_long = np.cumprod(1 - np.linspace(1e-4, 0.02, T))
print(f"abar_T for linear schedule (T={T}): {abar_long[-1]:.2e}")
assert abar_long[-1] < 1e-4
t=1: abar=0.900 mean matches (max err 0.0009), var matches t=2: abar=0.720 mean matches (max err 0.0014), var matches t=3: abar=0.504 mean matches (max err 0.0013), var matches Two composed single steps match the t=2 closed-form marginal. abar_T for linear schedule (T=1000): 4.04e-05
Exercise 5: Vanishing Generator Gradients (Failure Analysis)¶
Task: The minimax generator loss is $\mathcal{L}_{\text{sat}} = \log(1 - D(G(z)))$; the non-saturating alternative (implemented by gan_generator_loss in ml_first_principles) is $\mathcal{L}_{\text{ns}} = -\log D(G(z))$. Write $d = \sigma(s)$ where $s$ is the discriminator's logit on a fake sample.
- Derive $\frac{\partial \mathcal{L}_{\text{sat}}}{\partial s}$ and $\frac{\partial \mathcal{L}_{\text{ns}}}{\partial s}$ by hand.
- Evaluate both at $s \in \lbrace -6, -4, -2, 0 \rbrace$ and verify against central finite differences.
- Explain which loss vanishes when the discriminator confidently rejects fakes ($d \to 0$, i.e. $s \to -\infty$) — exactly the regime of early training — and why this is a failure mode.
Solution 5 (Derivation and Analysis)¶
With $d = \sigma(s)$ and $\frac{d\sigma}{ds} = d(1 - d)$ (chain rule):
- $\dfrac{\partial \mathcal{L}_{\text{sat}}}{\partial s} = \dfrac{-1}{1 - d} \cdot d(1 - d) = -d$
- $\dfrac{\partial \mathcal{L}_{\text{ns}}}{\partial s} = -\dfrac{1}{d} \cdot d(1 - d) = d - 1$
Result: as $d \to 0$ the saturating gradient $-d \to 0$ while the non-saturating gradient $d - 1 \to -1$.
Why this is a failure mode: early in training the generator's samples are obviously fake, so a decent discriminator assigns them $d \approx 0$ (large negative logits). Under $\mathcal{L}_{\text{sat}}$ the generator then receives gradients of magnitude $\approx d \approx 0$ — it cannot learn precisely when it most needs to. The non-saturating loss keeps the gradient near $-1$ in that regime, which is why practical GANs use it. WGAN removes the sigmoid/JSD saturation altogether: its critic is unbounded and 1-Lipschitz, so gradients stay informative even when $p_{\text{data}}$ and $p_g$ have disjoint support (see theory.md §4).
from ml_first_principles.generative_models import gan_generator_loss
def sigmoid(s):
return 1.0 / (1.0 + np.exp(-s))
logits = np.array([-6.0, -4.0, -2.0, 0.0])
d = sigmoid(logits)
grad_sat = -d # d/ds log(1 - sigmoid(s))
grad_ns = d - 1.0 # d/ds -log(sigmoid(s))
# Central finite-difference check of both analytic gradients
h = 1e-6
fd_sat = (np.log(1 - sigmoid(logits + h)) - np.log(1 - sigmoid(logits - h))) / (2 * h)
fd_ns = (-np.log(sigmoid(logits + h)) + np.log(sigmoid(logits - h))) / (2 * h)
assert np.allclose(grad_sat, fd_sat, atol=1e-6)
assert np.allclose(grad_ns, fd_ns, atol=1e-6)
# The library generator loss is the non-saturating objective -E[log D(G(z))]
assert np.allclose(gan_generator_loss(d), -np.mean(np.log(d)), atol=1e-9)
print(f"{'logit':>6} {'D(G(z))':>9} {'saturating grad':>16} {'non-saturating grad':>20}")
for s, di, gs, gn in zip(logits, d, grad_sat, grad_ns):
print(f"{s:6.1f} {di:9.5f} {gs:16.5f} {gn:20.5f}")
logit D(G(z)) saturating grad non-saturating grad -6.0 0.00247 -0.00247 -0.99753 -4.0 0.01799 -0.01799 -0.98201 -2.0 0.11920 -0.11920 -0.88080 0.0 0.50000 -0.50000 -0.50000
Exercise 6: VAE vs GAN vs Diffusion (Conceptual)¶
Task: Explain the fundamental differences between VAE, GAN, and Diffusion in terms of:
- Training objective
- Mode coverage (how well they cover the entire data distribution)
- Sample quality (visual fidelity of outputs)
Solution 6¶
Training Objective:
- VAE: Maximizes the Evidence Lower Bound (ELBO), an explicit approximate likelihood.
- GAN: Solves a minimax game between a Generator and Discriminator (implicit likelihood).
- Diffusion: Matches the score function by learning to denoise at various noise scales (often optimizing a variational bound similar to ELBO).
Mode Coverage:
- VAE: Excellent mode coverage. The KL divergence penalty encourages spreading mass over the entire data distribution, often leading to blurry outputs.
- GAN: Poor mode coverage. Prone to mode collapse because the generator only needs to produce a few highly realistic modes to fool the discriminator (Exercise 5 shows the closely related vanishing-gradient failure).
- Diffusion: Excellent mode coverage, similar to VAEs but often better due to the flexible diffusion steps.
Sample Quality:
- VAE: Generally blurry and lower quality due to pixel-wise independent assumptions.
- GAN: Very high fidelity and sharp samples.
- Diffusion: High fidelity, often matching or exceeding GANs, but at the cost of slow iterative sampling.