Variational Autoencoder (VAE) — Visual Experiments¶
Goal: Build intuition for VAE through experiments — from the intractable integral to the ELBO objective.
What we'll explore¶
| # | Experiment | Question answered |
|---|---|---|
| 1 | Point vs Distribution Encoding | What changes when we encode to a distribution instead of a point? |
| 2 | KL Divergence Geometry | What does the KL regularization term look like as a force on the encoder? |
| 3 | ELBO Decomposition | How do the reconstruction loss and KL penalty trade off to form the ELBO? |
| 4 | Latent Space Coverage | How does KL regularization fill the latent space gaps left by AE? |
| 5 | β-VAE: Information Rate | How does β control the reconstruction–compression Pareto frontier? |
Linked theory: research/03-vae.md · Chạy: uv run jupyter lab từ thư mục latent-anything-theory/ (cần các package trong pyproject.toml).
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib.patches import Ellipse
from sklearn.decomposition import PCA
from scipy.stats import gaussian_kde
import warnings
warnings.filterwarnings('ignore')
np.random.seed(42)
plt.rcParams.update({
'figure.dpi': 110,
'font.size': 11,
'axes.titlesize': 12,
'axes.labelsize': 11,
'figure.facecolor': 'white',
'axes.facecolor': '#f8f9fa',
'axes.grid': True,
'grid.alpha': 0.35,
})
print('Setup complete. NumPy:', np.__version__)
Setup complete. NumPy: 2.4.6
Experiment 1: Point Encoding (AE) vs Distribution Encoding (VAE)¶
Question: What fundamentally changes when we replace $z = g_\phi(x)$ (a single point) with $z \sim q_\phi(z|x) = \mathcal{N}(\mu(x),\, \sigma^2(x) \cdot I)$ (a distribution)?
A vanilla AE is deterministic: every input $x$ maps to exactly one latent code $z$. If the training data forms 3 clusters, the latent space has 3 isolated islands — any point sampled outside those islands has never been seen by the decoder. A VAE encoder instead outputs the parameters $(\mu, \sigma)$ of a Gaussian, then samples $z$ from it. This stochasticity spreads each cluster into an overlapping cloud, filling the gaps.
np.random.seed(1)
n_per = 200
cluster_centers = [(-2.5, 0.5), (2.0, 1.0), (0.2, -2.2)]
cluster_std = 0.35
X_parts = [
np.random.randn(n_per, 2) * cluster_std + np.array(c)
for c in cluster_centers
]
X_data = np.vstack(X_parts)
y_data = np.repeat([0, 1, 2], n_per)
# AE: deterministic PCA encoding — each x maps to one fixed z
pca = PCA(n_components=2)
Z_ae = pca.fit_transform(X_data)
ae_cluster_centers = np.array([Z_ae[y_data == i].mean(0) for i in range(3)])
# VAE simulation: encoder outputs (mu, sigma); z is sampled from N(mu, sigma^2 I)
sigma_vae = 0.75 # typical learned encoder std
Z_vae_mu = Z_ae.copy()
Z_vae = Z_vae_mu + sigma_vae * np.random.randn(*Z_ae.shape)
print('AE cluster centers (latent):')
for i, c in enumerate(ae_cluster_centers):
print(f' Class {i}: ({c[0]:+.2f}, {c[1]:+.2f})')
print(f'VAE encoder sigma: {sigma_vae}')
AE cluster centers (latent): Class 0: (-2.33, +0.87) Class 1: (+2.15, +1.11) Class 2: (+0.18, -1.98) VAE encoder sigma: 0.75
palette = ['#e74c3c', '#2ecc71', '#3498db']
fig, axes = plt.subplots(1, 2, figsize=(14, 6))
fig.suptitle(
'Experiment 1: Point Encoding (AE) vs Probabilistic Encoding (VAE)',
fontsize=13, fontweight='bold'
)
# --- Left: AE latent space (isolated point clusters) ---
ax1 = axes[0]
for cls, col in enumerate(palette):
mask = y_data == cls
ax1.scatter(Z_ae[mask, 0], Z_ae[mask, 1], c=col, alpha=0.55, s=14,
label=f'Class {cls}')
# Highlight gap regions with an annotation
for i, j in [(0, 1), (1, 2), (0, 2)]:
mid = (ae_cluster_centers[i] + ae_cluster_centers[j]) / 2
ax1.annotate(
'gap', xy=mid, fontsize=8, color='gray', ha='center',
bbox=dict(boxstyle='round,pad=0.2', facecolor='white', edgecolor='gray', alpha=0.7)
)
ax1.set_xlabel('z\u2080'); ax1.set_ylabel('z\u2081')
ax1.set_title('AE: x \u2192 z (one fixed point)\nIsolated islands \u2014 gaps are empty')
ax1.legend(fontsize=9)
# --- Right: VAE latent space (overlapping Gaussians) ---
ax2 = axes[1]
# Draw Gaussian ellipses for each cluster
for cls, col in enumerate(palette):
mu_c = ae_cluster_centers[cls]
for n_std, alpha in [(2, 0.10), (1, 0.22)]:
ell = Ellipse(
xy=mu_c,
width=2 * n_std * sigma_vae,
height=2 * n_std * sigma_vae,
facecolor=col, edgecolor=col, lw=1.5, alpha=alpha
)
ax2.add_patch(ell)
# Sampled z values
mask = y_data == cls
ax2.scatter(Z_vae[mask, 0], Z_vae[mask, 1], c=col, alpha=0.25, s=10)
# Cluster mean star
ax2.scatter(*mu_c, marker='*', s=220, color=col, edgecolors='black', zorder=10)
handles = [mpatches.Patch(color=c, label=f'Class {i}') for i, c in enumerate(palette)]
ax2.legend(handles=handles, fontsize=9)
ax2.set_xlabel('z\u2080'); ax2.set_ylabel('z\u2081')
ax2.set_title(
'VAE: x \u2192 (\u03bc(x), \u03c3(x)), then z ~ N(\u03bc, \u03c3\u00b2I)\n'
'Gaussians overlap \u2014 gaps are filled'
)
ax2.set_xlim(ax1.get_xlim())
ax2.set_ylim(ax1.get_ylim())
plt.tight_layout()
plt.show()
What you see¶
What the plot shows: The AE latent space (left) has three tight, well-separated clusters. The regions labeled "gap" have zero training support — the decoder has never learned to decode any point there. The VAE latent space (right) shows the same cluster means (★), but each input maps to a Gaussian (shaded ellipses show 1σ and 2σ). The sampled $z$ values spread across the ellipses, overlapping between clusters.
Why it looks this way: The VAE encoder is stochastic by design. For each input $x$, it outputs $(\mu(x), \sigma(x))$ and then $z$ is drawn from $\mathcal{N}(\mu, \sigma^2 I)$. The KL term forces $\sigma$ to remain non-negligible — it cannot shrink to zero (which would collapse back to the AE). This spread of probability mass is what fills the gaps between clusters.
Key takeaway: The switch from deterministic $z = g_\phi(x)$ to stochastic $z \sim q_\phi(z|x)$ is not just a technical detail. It is the mechanism that transforms the latent space from a sparse collection of points into a continuous probability landscape, making random sampling and interpolation meaningful.
Experiment 2: KL Divergence Geometry — The Regularization Force¶
Question: What does the KL penalty $D_{KL}(\mathcal{N}(\mu, \sigma^2) \| \mathcal{N}(0,1))$ look like as a geometric landscape, and what forces does it exert on the encoder?
For the Gaussian case, the KL has a clean analytical form: $$D_{KL} = \frac{1}{2}\left(\mu^2 + \sigma^2 - \ln \sigma^2 - 1\right)$$
This is always $\geq 0$, with a unique minimum of $0$ at $(\mu=0,\, \sigma=1)$ — exactly the prior $\mathcal{N}(0,1)$. The KL term acts as a restoring force in two directions: it pulls the encoder mean $\mu$ toward 0, and it pulls the encoder variance $\sigma^2$ toward 1.
np.random.seed(2)
def kl_gaussian(mu: np.ndarray, sigma: np.ndarray) -> np.ndarray:
"""KL(N(mu, sigma^2) || N(0,1)) — closed-form for diagonal Gaussians."""
return 0.5 * (mu**2 + sigma**2 - np.log(sigma**2 + 1e-12) - 1)
mu_range = np.linspace(-3.5, 3.5, 250)
sigma_range = np.linspace(0.08, 3.2, 250)
MU, SIGMA = np.meshgrid(mu_range, sigma_range)
KL_surface = kl_gaussian(MU, SIGMA)
# Cross-sections
kl_vs_mu = kl_gaussian(mu_range, 1.0) # varying mu, sigma=1
kl_vs_sigma = kl_gaussian(0.0, sigma_range) # varying sigma, mu=0
print(f'KL at (mu=0, sigma=1): {kl_gaussian(0.0, 1.0):.6f} (theoretical min = 0)')
print(f'KL at (mu=2, sigma=1): {kl_gaussian(2.0, 1.0):.4f}')
print(f'KL at (mu=0, sigma=0.1): {kl_gaussian(0.0, 0.1):.4f} (near-deterministic)')
print(f'KL at (mu=0, sigma=3.0): {kl_gaussian(0.0, 3.0):.4f} (too wide)')
KL at (mu=0, sigma=1): -0.000000 (theoretical min = 0) KL at (mu=2, sigma=1): 2.0000 KL at (mu=0, sigma=0.1): 1.8076 (near-deterministic) KL at (mu=0, sigma=3.0): 2.9014 (too wide)
fig, axes = plt.subplots(1, 3, figsize=(17, 5))
fig.suptitle(
'Experiment 2: KL Divergence Geometry \u2014 The Regularization Force',
fontsize=13, fontweight='bold'
)
# --- Left: 2D heatmap over (mu, sigma) ---
ax1 = axes[0]
cf = ax1.contourf(
MU, SIGMA, np.clip(KL_surface, 0, 8),
levels=25, cmap='YlOrRd'
)
plt.colorbar(cf, ax=ax1, label='KL divergence')
ax1.contour(
MU, SIGMA, KL_surface,
levels=[0.1, 0.5, 1.0, 2.0, 4.0],
colors='black', linewidths=0.7, alpha=0.4
)
ax1.scatter(0, 1, marker='*', s=300, color='#3498db',
edgecolors='black', zorder=10,
label='Minimum: KL=0 at (\u03bc=0, \u03c3=1)')
# Arrows showing the "restoring force" direction
arrow_kw = dict(arrowstyle='->', color='white', lw=2)
for mu_pt, sigma_pt in [(2.5, 0.4), (-2.5, 0.4), (2.5, 2.5), (-2.5, 2.5)]:
dx = -mu_pt * 0.25
ds = (1.0 - sigma_pt) * 0.3
ax1.annotate('', xy=(mu_pt + dx, sigma_pt + ds), xytext=(mu_pt, sigma_pt),
arrowprops=arrow_kw)
ax1.set_xlabel('\u03bc (encoder mean)')
ax1.set_ylabel('\u03c3 (encoder std)')
ax1.set_title(
'KL(N(\u03bc,\u03c3\u00b2) \u2016 N(0,1)) heatmap\n'
'Arrows show gradient force toward minimum'
)
ax1.legend(fontsize=8.5, loc='upper right')
# --- Middle: KL vs mu (sigma=1) ---
ax2 = axes[1]
for sigma_val, col, lbl in [
(1.0, '#e74c3c', '\u03c3 = 1.0 (prior width)'),
(0.5, '#9b59b6', '\u03c3 = 0.5 (narrow)'),
(2.0, '#27ae60', '\u03c3 = 2.0 (wide)'),
]:
ax2.plot(mu_range, kl_gaussian(mu_range, sigma_val),
color=col, lw=2.5, label=lbl)
ax2.axvline(0, color='gray', ls=':', lw=1.2, alpha=0.6)
ax2.set_xlabel('\u03bc (encoder mean)')
ax2.set_ylabel('KL divergence')
ax2.set_title('KL vs Mean \u03bc\nParabolic penalty: pulls \u03bc \u2192 0')
ax2.legend(fontsize=9)
ax2.set_ylim(bottom=0)
# --- Right: KL vs sigma (mu=0) ---
ax3 = axes[2]
ax3.plot(sigma_range, kl_vs_sigma, color='#3498db', lw=2.5, label='\u03bc = 0')
ax3.axvline(1.0, color='gray', ls='--', lw=1.5, label='\u03c3 = 1 (minimum)')
ax3.fill_between(
sigma_range, kl_vs_sigma,
where=(sigma_range < 0.8), alpha=0.18, color='#e74c3c',
label='Too narrow (\u03c3 < 1) \u2192 near-deterministic'
)
ax3.fill_between(
sigma_range, kl_vs_sigma,
where=(sigma_range > 1.6), alpha=0.10, color='#e67e22',
label='Too wide (\u03c3 > 1) \u2192 uninformative'
)
ax3.scatter(1.0, kl_gaussian(0.0, 1.0), s=120, color='#3498db',
edgecolors='black', zorder=10)
ax3.set_xlabel('\u03c3 (encoder std)')
ax3.set_ylabel('KL divergence')
ax3.set_title('KL vs Std \u03c3\nPulls \u03c3 \u2192 1: prevents collapse AND over-spread')
ax3.legend(fontsize=8.5)
ax3.set_ylim(bottom=0)
plt.tight_layout()
plt.show()
What you see¶
What the plot shows: The 2D heatmap (left) reveals that the KL surface is a bowl with its minimum at $(\mu=0, \sigma=1)$ — the prior itself. The white arrows show the gradient force pulling every encoder configuration toward this minimum. The cross-sections confirm the two independent penalties: KL is a parabola in $\mu$ (center), punishing off-center means; and a U-shaped curve in $\sigma$ (right), punishing both too-narrow and too-wide encoders.
Why it looks this way: Expanding the formula $D_{KL} = \tfrac{1}{2}(\mu^2 + \sigma^2 - \ln\sigma^2 - 1)$ shows two separable contributions. The $\mu^2$ term is a quadratic penalty on the mean — the farther the encoder cluster is from the origin, the higher the cost. The $\sigma^2 - \ln\sigma^2 - 1$ term is a convex function with minimum at $\sigma=1$: it prevents the encoder from going fully deterministic ($\sigma \to 0$, which would recreate the AE gap problem) and from becoming too diffuse ($\sigma \gg 1$, which destroys information).
Key takeaway: The KL term is a two-sided regularizer that continuously exerts a force pulling $\mu \to 0$ and $\sigma \to 1$. It prevents latent collapse (the AE degenerate case) and prevents the encoder from becoming uninformative — this is what makes the VAE latent space smooth and continuously decodable.
Experiment 3: ELBO = Reconstruction + KL — A Balancing Act¶
Question: How do the reconstruction term and the KL term of the ELBO trade off, and what is the optimal encoder variance?
The ELBO decomposes as: $$\text{ELBO} = \underbrace{\mathbb{E}_{q(z|x)}[\log p(x|z)]}_{\text{reconstruction}} - \underbrace{D_{KL}(q(z|x) \| p(z))}_{\text{KL penalty}}$$
For a 1D Gaussian model with fixed $\mu_z = \mu_x$ and decoder variance $\tau^2 = 1$, the optimal encoder variance has a clean closed form: $\sigma_z^\star = \frac{1}{\sqrt{2}}$. This experiment visualizes both loss terms and their sum as functions of the encoder standard deviation $\sigma_z$.
np.random.seed(3)
# Model: x = 2.0, encoder q(z|x) = N(mu_z, sigma_z^2)
# mu_z = 2.0 (encoder mean = data), tau^2 = 1 (decoder variance)
# Reconstruction loss: E_q[-log p(x|z)] = 0.5/tau^2 * (sigma_z^2 + (x-mu_z)^2)
# KL: 0.5 * (mu_z^2 + sigma_z^2 - log(sigma_z^2) - 1)
# ELBO: -recon_loss - kl_loss (maximize)
x0 = 2.0
mu_z = 2.0 # encoder mean equals data value
tau2 = 1.0 # decoder variance
sigma_z = np.linspace(0.05, 2.8, 500)
recon_loss = 0.5 / tau2 * (sigma_z**2 + (x0 - mu_z)**2)
kl_loss = 0.5 * (mu_z**2 + sigma_z**2 - np.log(sigma_z**2 + 1e-12) - 1)
elbo = -(recon_loss + kl_loss)
# Analytical optimum: sigma_z^2 = tau^2 / (1 + tau^2) = 0.5
sigma_opt = np.sqrt(tau2 / (1 + tau2))
elbo_opt = float(elbo[np.argmin(np.abs(sigma_z - sigma_opt))])
print(f'Analytical optimal sigma_z* = sqrt(1/2) = {sigma_opt:.4f}')
print(f'ELBO at sigma_z*: {elbo_opt:.4f}')
print(f'Recon loss at sigma_z*: {float(recon_loss[np.argmin(np.abs(sigma_z - sigma_opt))]):.4f}')
print(f'KL loss at sigma_z*: {float(kl_loss[np.argmin(np.abs(sigma_z - sigma_opt))]):.4f}')
Analytical optimal sigma_z* = sqrt(1/2) = 0.7071 ELBO at sigma_z*: -2.3466 Recon loss at sigma_z*: 0.2491 KL loss at sigma_z*: 2.0975
fig, axes = plt.subplots(1, 3, figsize=(17, 5))
fig.suptitle(
'Experiment 3: ELBO Decomposition \u2014 Reconstruction vs KL Penalty',
fontsize=13, fontweight='bold'
)
# --- Left: both loss terms vs sigma_z ---
ax1 = axes[0]
ax1.plot(sigma_z, recon_loss, color='#e74c3c', lw=2.5, label='Reconstruction loss')
ax1.plot(sigma_z, kl_loss, color='#3498db', lw=2.5, label='KL penalty')
ax1.axvline(sigma_opt, color='gray', ls='--', lw=1.8, alpha=0.8,
label=f'Optimal \u03c3* = {sigma_opt:.3f}')
ax1.set_xlabel('Encoder \u03c3\u1d63 (std of q(z|x))')
ax1.set_ylabel('Loss value')
ax1.set_title(
'Reconstruction and KL terms vs \u03c3\u1d63\n'
'Each pulls in the opposite direction'
)
ax1.legend(fontsize=9)
ax1.set_ylim(0, 8)
ax1.set_xlim(0, 2.8)
# Annotate the two extreme regimes
ax1.annotate(
'AE regime:\nhigh KL, low recon',
xy=(0.15, kl_loss[5]), xytext=(0.6, 6.5),
arrowprops=dict(arrowstyle='->', color='gray'),
fontsize=8, color='gray', ha='center'
)
ax1.annotate(
'Noisy regime:\nhigh recon, low KL',
xy=(2.4, recon_loss[-30]), xytext=(1.6, 5.5),
arrowprops=dict(arrowstyle='->', color='gray'),
fontsize=8, color='gray', ha='center'
)
# --- Middle: ELBO vs sigma_z ---
ax2 = axes[1]
ax2.plot(sigma_z, elbo, color='#27ae60', lw=2.5, label='ELBO (maximize \u2191)')
ax2.axvline(sigma_opt, color='gray', ls='--', lw=1.8, alpha=0.8)
ax2.scatter(sigma_opt, elbo_opt, s=150, color='#e74c3c',
edgecolors='black', zorder=10,
label=f'Maximum at \u03c3* = {sigma_opt:.3f}')
ax2.set_xlabel('Encoder \u03c3\u1d63')
ax2.set_ylabel('ELBO (= \u2212recon \u2212 KL)')
ax2.set_title('ELBO has a unique maximum\nwhere reconstruction and KL balance')
ax2.legend(fontsize=9)
ax2.set_xlim(0, 2.8)
# --- Right: the encoder distributions at 3 sigma values vs the prior ---
ax3 = axes[2]
z_range = np.linspace(-2.5, 6.0, 500)
configs = [
(0.12, '#e74c3c', 'Too narrow: \u03c3=0.12\n(AE-like, high KL)'),
(sigma_opt,'#27ae60', f'Optimal: \u03c3={sigma_opt:.3f}\n(max ELBO)'),
(2.2, '#3498db', 'Too wide: \u03c3=2.2\n(uninformative)'),
]
for sigma_val, col, lbl in configs:
q = np.exp(-0.5 * ((z_range - mu_z) / sigma_val)**2) / (sigma_val * np.sqrt(2 * np.pi))
ax3.plot(z_range, q, lw=2.5, color=col, label=lbl)
prior = np.exp(-0.5 * z_range**2) / np.sqrt(2 * np.pi)
ax3.plot(z_range, prior, 'k--', lw=1.8, alpha=0.7, label='Prior p(z) = N(0,1)')
ax3.axvline(mu_z, color='gray', ls=':', lw=1, alpha=0.7)
ax3.axvline(0, color='black', ls=':', lw=0.8, alpha=0.5)
ax3.set_xlabel('z')
ax3.set_ylabel('Probability density')
ax3.set_title(
'q(z|x) at three \u03c3 values vs prior p(z)\n'
'Optimal \u03c3 balances width vs prior alignment'
)
ax3.legend(fontsize=8.5)
plt.tight_layout()
plt.show()
What you see¶
What the plot shows: The left panel shows the two loss terms moving in opposite directions as $\sigma_z$ varies — the reconstruction loss (red) rises with $\sigma_z$ because a wider encoder introduces more noise, while the KL penalty (blue) rises as $\sigma_z \to 0$ because the encoder becomes near-deterministic (far from the prior). The middle panel shows their sum (ELBO) has a single maximum at $\sigma_z^\star = 1/\sqrt{2} \approx 0.707$. The right panel shows the three encoder distributions: the too-narrow one is a spike far from the prior; the optimal one has width matching the prior; the too-wide one blurs over both.
Why it looks this way: For a Gaussian model, maximizing ELBO over $\sigma_z$ yields the closed-form optimum $\sigma_z^\star = \sqrt{\tau^2 / (1 + \tau^2)}$ where $\tau^2$ is the decoder variance. With $\tau^2 = 1$, this gives $\sigma_z^\star = 1/\sqrt{2}$. The formula shows that a more precise decoder (smaller $\tau^2$) pushes the optimal encoder variance lower — the encoder can afford to be tighter when the decoder is more accurate.
Key takeaway: The ELBO is not an ad hoc loss — it is a rigorously derived lower bound on $\log P(X)$, and maximizing it simultaneously trains the decoder to reconstruct accurately and trains the encoder to approximate the prior. The two-term decomposition reveals the fundamental tension between compression (KL) and fidelity (reconstruction) that defines all useful latent representations.
Experiment 4: Latent Space Coverage — AE Gaps vs VAE Regularity¶
Question: How does increasing the KL regularization weight $\beta$ transform the shape of the latent space, from isolated AE clusters toward a smooth prior?
The $\beta$-VAE objective $\text{ELBO}_\beta = \text{recon} - \beta \cdot D_{KL}$ generalizes the VAE: $\beta=0$ recovers the AE (no regularization), $\beta=1$ is the standard VAE, and $\beta > 1$ applies stronger compression pressure. This experiment simulates the latent space for four values of $\beta$ on the same 3-cluster dataset, showing how the encoder distributions shift from isolated to overlapping to merged.
np.random.seed(4)
n_per = 250
cluster_specs = [(-2.2, 0.8), (2.1, 0.6), (0.0, -2.0)]
cluster_std = 0.32
X_parts = [
np.random.randn(n_per, 2) * cluster_std + np.array(c)
for c in cluster_specs
]
X_data = np.vstack(X_parts)
y_data = np.repeat([0, 1, 2], n_per)
# Base: AE encoding (PCA = deterministic point mapping)
pca = PCA(n_components=2)
Z_ae = pca.fit_transform(X_data)
Z_ae_norm = Z_ae / Z_ae.std() # normalize to unit scale for comparability
# Simulate VAE latent for 4 beta values:
# As beta increases, the KL term pulls mu toward 0 and sigma toward 1.
# We model: mu_eff = mu_ae * shrink(beta), sigma_eff = sigma(beta)
def simulate_vae_latent(Z_ae_norm, beta, n_samples):
shrink = 1.0 / (1.0 + beta * 1.8)
sigma = np.clip(0.15 + beta * 0.4, 0.15, 1.05)
Z_mu = Z_ae_norm * shrink
Z_samp = Z_mu + sigma * np.random.randn(*Z_ae_norm.shape)
return Z_samp, shrink, sigma
beta_vals = [0.0, 0.5, 1.0, 4.0]
latent_by_beta = {}
for beta in beta_vals:
if beta == 0.0:
latent_by_beta[beta] = (Z_ae_norm, 1.0, 0.0)
else:
latent_by_beta[beta] = simulate_vae_latent(Z_ae_norm, beta, n_per)
print('Beta | Mean shrink | Sigma')
print('-' * 34)
for beta, (Z, shrink, sigma) in latent_by_beta.items():
print(f'{beta:4.1f} | {shrink:.3f} | {sigma:.3f}')
Beta | Mean shrink | Sigma ---------------------------------- 0.0 | 1.000 | 0.000 0.5 | 0.526 | 0.350 1.0 | 0.357 | 0.550 4.0 | 0.122 | 1.050
palette = ['#e74c3c', '#2ecc71', '#3498db']
titles = [
'\u03b2 = 0\n(AE: no KL \u2192 isolated clusters)',
'\u03b2 = 0.5\n(mild KL \u2192 clusters start merging)',
'\u03b2 = 1.0\n(standard VAE \u2192 overlapping)',
'\u03b2 = 4.0\n(strong KL \u2192 near N(0,I))',
]
fig, axes = plt.subplots(2, 4, figsize=(18, 8))
fig.suptitle(
'Experiment 4: Effect of KL Regularization \u03b2 on Latent Space Coverage',
fontsize=13, fontweight='bold'
)
lim = 3.2
grid_pts = np.linspace(-lim, lim, 80)
GX, GY = np.meshgrid(grid_pts, grid_pts)
grid_pos = np.vstack([GX.ravel(), GY.ravel()])
# Prior density for reference
prior_density = np.exp(-0.5 * (GX**2 + GY**2)) / (2 * np.pi)
for col_idx, (beta, title) in enumerate(zip(beta_vals, titles)):
Z, shrink, sigma = latent_by_beta[beta]
ax_top = axes[0, col_idx]
ax_bot = axes[1, col_idx]
# --- Row 1: scatter plot of latent codes ---
for cls, col in enumerate(palette):
mask = y_data == cls
ax_top.scatter(Z[mask, 0], Z[mask, 1], c=col, alpha=0.4, s=10)
ax_top.set_xlim(-lim, lim); ax_top.set_ylim(-lim, lim)
ax_top.set_xlabel('z\u2080'); ax_top.set_ylabel('z\u2081')
ax_top.set_title(title, fontsize=10)
ax_top.set_aspect('equal')
# --- Row 2: latent density (KDE) vs prior ---
kde = gaussian_kde(Z.T, bw_method=0.25)
density = kde(grid_pos).reshape(80, 80)
cf = ax_bot.contourf(GX, GY, density, levels=15, cmap='YlOrRd', alpha=0.85)
ax_bot.contour(GX, GY, prior_density,
levels=[0.04, 0.10, 0.15],
colors=['#3498db'], linewidths=1.5, linestyles='--', alpha=0.8)
ax_bot.set_xlim(-lim, lim); ax_bot.set_ylim(-lim, lim)
ax_bot.set_xlabel('z\u2080'); ax_bot.set_ylabel('z\u2081')
ax_bot.set_title(
'KDE density (orange) vs prior N(0,I) (blue dashed)',
fontsize=9
)
ax_bot.set_aspect('equal')
plt.tight_layout()
plt.show()
What you see¶
What the plot shows: The top row shows the scatter of latent codes at each $\beta$. At $\beta=0$ (AE), the three clusters sit at isolated fixed positions with visible white gaps. As $\beta$ increases to 0.5 and 1.0, the clusters shrink toward the origin and grow wider, starting to overlap. At $\beta=4$, the latent space looks close to $\mathcal{N}(0, I)$ — a single smooth Gaussian cloud. The bottom row shows the kernel density estimate (orange) overlaid with the prior density contours (blue dashed), confirming how the marginal distribution of $z$ approaches the prior as $\beta$ grows.
Why it looks this way: The KL term $D_{KL}(q(z|x) \| p(z))$ acts as a regularizer that penalizes two things: (a) encoder means that are far from zero (the $\mu^2$ term), and (b) encoder variances far from 1 (the $\sigma^2 - \ln\sigma^2$ term). As $\beta$ grows, both penalties dominate the training objective and the encoder is forced to "collapse" the posterior toward the prior. The trade-off is that with very large $\beta$, the encoder loses discriminative power and reconstruction quality degrades.
Key takeaway: The KL regularization is what transforms the latent space from a deterministic codebook (AE) into a continuous probability density (VAE). The standard VAE ($\beta=1$) strikes the canonical balance: enough regularity to support random sampling and interpolation, while still preserving enough structure for the decoder to reconstruct faithfully. This is the core contribution of the VAE framework.
Experiment 5: β-VAE — Controlling the Information Rate¶
Question: How does the hyperparameter $\beta$ in the β-VAE objective control the reconstruction–compression trade-off, and what is its connection to the Information Bottleneck?
The β-VAE (Higgins et al., 2017) generalises the standard VAE by placing a weight $\beta \geq 0$ in front of the KL term:
$$\text{ELBO}_\beta = \underbrace{\mathbb{E}_{q(z|x)}[\log p(x|z)]}_{\text{reconstruction}} - \beta \cdot \underbrace{D_{KL}(q(z|x) \| p(z))}_{\approx\, I(Z;X)}$$
- $\beta = 0$: no regularisation → degenerates to an AE.
- $\beta = 1$: standard VAE.
- $\beta > 1$: stronger regularisation → more compression and potentially more disentangled latent factors, at the cost of reconstruction quality.
This objective is exactly the IB Lagrangian from notebook 01: $\beta$ is the same Lagrange multiplier that controls the compression–prediction trade-off on the Information Plane.
np.random.seed(5)
# Analytical beta-VAE: 1D Gaussian model, x = mu_z = 2.0, decoder variance tau^2 = 1
# Optimal encoder variance under beta-ELBO:
# d/d(sigma^2)[-recon - beta*KL] = 0 => sigma^2_opt(beta) = beta*tau^2 / (1 + beta*tau^2)
mu_z = 2.0
tau2 = 1.0
beta_range = np.logspace(-1.5, 1.5, 400) # beta from ~0.03 to ~32
sigma_opt = np.sqrt(beta_range / (1.0 + beta_range)) # optimal sigma*(beta)
recon_opt = 0.5 / tau2 * sigma_opt**2 # recon loss at optimum
kl_opt = 0.5 * (mu_z**2 + sigma_opt**2
- np.log(sigma_opt**2 + 1e-12) - 1) # KL at optimum
# Information rate: I(Z;X) ~= 0.5 * log(1 + mu^2/sigma^2) (SNR approximation)
snr = mu_z**2 / (sigma_opt**2 + 1e-12)
info_rate = 0.5 * np.log1p(snr)
mark_betas = [0.1, 0.5, 1.0, 2.0, 5.0, 10.0]
mark_sigmas = [float(np.interp(b, beta_range, sigma_opt)) for b in mark_betas]
mark_recons = [float(np.interp(b, beta_range, recon_opt)) for b in mark_betas]
mark_infos = [float(np.interp(b, beta_range, info_rate)) for b in mark_betas]
print(f'{"beta":>6} {"sigma_opt":>9} {"recon":>7} {"I(Z;X)":>7}')
print('-' * 38)
for b, s, r, i in zip(mark_betas, mark_sigmas, mark_recons, mark_infos):
print(f'{b:6.1f} {s:9.4f} {r:7.4f} {i:7.4f}')
beta sigma_opt recon I(Z;X) -------------------------------------- 0.1 0.3015 0.0455 1.9033 0.5 0.5773 0.1667 1.2825 1.0 0.7071 0.2500 1.0986 2.0 0.8165 0.3333 0.9730 5.0 0.9129 0.4167 0.8789 10.0 0.9535 0.4545 0.8432
mcs = plt.cm.plasma(np.linspace(0.1, 0.9, len(mark_betas)))
fig, axes = plt.subplots(1, 3, figsize=(17, 5))
fig.suptitle(
'Experiment 5: β-VAE — Controlling the Information Rate',
fontsize=13, fontweight='bold'
)
# --- Left: optimal sigma vs beta ---
ax1 = axes[0]
ax1.semilogx(beta_range, sigma_opt, color='#9b59b6', lw=2.5, label='σ*(β)')
ax1.axvline(1.0, color='gray', ls='--', lw=1.5, alpha=0.8, label='β=1 (standard VAE)')
ax1.axhline(1.0, color='#e74c3c', ls=':', lw=1.5, alpha=0.7, label='σ=1 (prior std)')
for b, s, col in zip(mark_betas, mark_sigmas, mcs):
ax1.scatter(b, s, s=110, color=col, edgecolors='black', zorder=10)
ax1.annotate(f'β={b}', (b, s), textcoords='offset points', xytext=(5, 4), fontsize=7.5)
ax1.set_xlabel('β (log scale)')
ax1.set_ylabel('Optimal encoder σ*')
ax1.set_title(
'Optimal encoder std grows with β\n'
'Higher β → wider encoder → less information in z'
)
ax1.legend(fontsize=8.5)
ax1.set_xlim(beta_range[0], beta_range[-1])
ax1.set_ylim(0, 1.08)
# --- Middle: reconstruction loss and KL at optimum vs beta ---
ax2 = axes[1]
ax2.semilogx(beta_range, recon_opt, color='#e74c3c', lw=2.5, label='Reconstruction loss')
ax2.semilogx(beta_range, kl_opt, color='#3498db', lw=2.5, label='KL penalty')
ax2.axvline(1.0, color='gray', ls='--', lw=1.5, alpha=0.8, label='β=1')
for b, r, col in zip(mark_betas, mark_recons, mcs):
ax2.scatter(b, r, s=90, color=col, edgecolors='black', zorder=10)
ax2.set_xlabel('β (log scale)')
ax2.set_ylabel('Loss value at optimum')
ax2.set_title(
'Reconstruction degrades as β increases\n'
'KL stabilises as encoder approaches prior N(0,1)'
)
ax2.legend(fontsize=8.5)
ax2.set_ylim(bottom=0)
ax2.set_xlim(beta_range[0], beta_range[-1])
# --- Right: Information Plane --- I(Z;X) vs reconstruction quality ---
ax3 = axes[2]
ax3.plot(info_rate, -recon_opt, color='#27ae60', lw=3, label='β-VAE Pareto frontier', zorder=5)
mid = len(info_rate) // 2
ax3.annotate(
'', xy=(info_rate[mid + 20], -recon_opt[mid + 20]),
xytext=(info_rate[mid - 20], -recon_opt[mid - 20]),
arrowprops=dict(arrowstyle='->', color='#27ae60', lw=2.5)
)
ax3.text(
info_rate[mid + 25], -recon_opt[mid + 25] - 0.008,
'β increases →', fontsize=8.5, color='#27ae60'
)
for b, i_m, r_m, col in zip(mark_betas, mark_infos, mark_recons, mcs):
ax3.scatter(i_m, -r_m, s=140, color=col, edgecolors='black', zorder=10)
ax3.annotate(f'β={b}', (i_m, -r_m), textcoords='offset points', xytext=(4, 5), fontsize=8)
ax3.set_xlabel('I(Z;X) — Information rate (nats)')
ax3.set_ylabel('Reconstruction quality (−recon loss)')
ax3.set_title(
'Information Plane: β traces the Pareto frontier\n'
'Mirrors the IB information plane (notebook 01)'
)
ax3.legend(fontsize=9)
plt.tight_layout()
plt.show()
What you see¶
What the plot shows: The left panel shows the optimal encoder std $\sigma^*(\beta)$ rising monotonically toward 1 as $\beta$ increases — the encoder is forced to spread more probability mass to satisfy the KL constraint. The middle panel shows how both losses evolve at the optimum: reconstruction loss climbs steadily while the KL term eventually stabilises (once $\sigma \to 1$, the only remaining KL is from the off-centre mean $\mu \neq 0$). The right panel is the Information Plane: each coloured dot is one $\beta$ value, tracing a Pareto frontier from high-information/good-reconstruction ($\beta \to 0$, top-right) to low-information/poor-reconstruction ($\beta \to \infty$, bottom-left).
Why it looks this way: The closed-form solution $\sigma^2_{\text{opt}}(\beta) = \beta\tau^2 / (1 + \beta\tau^2)$ shows that $\beta$ directly controls the "noise budget" allocated to the latent code. This reduces the mutual information $I(Z;X) \approx \tfrac{1}{2}\ln(1 + \mu^2/\sigma^2)$. The information plane here is structurally identical to the IB information plane in notebook 01 — the VAE/β-VAE family implements the IB objective with $\beta$ playing the role of the Lagrange multiplier.
Key takeaway: β-VAE is not a separate model — it is a tunable instantiation of the Information Bottleneck principle. $\beta$ selects an operating point on the Pareto frontier: $\beta=1$ is the standard VAE, $\beta > 1$ prioritises a compact structured latent space (encouraging disentangled factors at the cost of reconstruction), and $\beta < 1$ relaxes regularisation to preserve more detail. The choice is application-driven: generative modelling favours $\beta \approx 1$, disentanglement research typically uses $\beta \geq 4$.
Summary¶
| Experiment | Key insight |
|---|---|
| 1. Point vs Distribution | Stochastic encoding fills latent gaps by spreading each cluster into a probability cloud |
| 2. KL Geometry | KL acts as a two-sided force: pulls $\mu \to 0$ (parabolic) and $\sigma \to 1$ (U-shaped) |
| 3. ELBO Decomposition | Optimal encoder variance balances reconstruction fidelity against prior alignment |
| 4. Latent Coverage | Increasing $\beta$ continuously morphs the latent distribution from isolated clusters toward $\mathcal{N}(0, I)$ |
| 5. β-VAE Information Rate | $\beta$ traces the IB Pareto frontier: higher $\beta$ = more compression, lower $I(Z;X)$, worse reconstruction |
Connection to the broader theory¶
The VAE resolves the two core limitations of the vanilla AE established in the previous notebook: the intractable marginal $\log P(X)$ is replaced by a tractable lower bound (ELBO), and the latent gap problem is eliminated by imposing a continuous prior through the KL penalty. The ELBO objective is precisely the Information Bottleneck Lagrangian from notebook 01, with reconstruction playing the role of $I(Z;Y)$ and the KL term playing the role of $\beta \cdot I(Z;X)$ — the VAE is therefore a concrete probabilistic implementation of the IB principle. The β-VAE objective makes this IB connection explicit: $\beta$ is the Lagrange multiplier that positions the model on the Pareto frontier between compression and fidelity, and larger $\beta$ encourages disentangled latent factors by imposing stronger independence pressure via the KL term. The next natural step is to examine what happens when the decoder is also stochastic and the latent space has explicit geometric structure (Riemannian VAEs, FlatVI).