Autoencoders and VAEs

A sequence model compresses an ordered history into a hidden state. A screenshot crop poses the same question in pixels. Can two numbers preserve a bright 2 by 2 alert badge on a 5 by 5 grayscale screenshot? Suppose you hand the decoder a random pair of numbers. What can it draw if training only showed it codes produced by an encoder?

An autoencoder learns to reconstruct its input from a latent code, which we make compact here.^{[1]Reference 1Reducing the Dimensionality of Data with Neural Networkshttps://doi.org/10.1126/science.1127647} A variational autoencoder (VAE) learns a probabilistic latent model with a prior we can sample for generation. Its objective connects reconstruction and prior sampling, but useful samples still depend on the model, data, and training.^{[2]Reference 2Auto-Encoding Variational Bayes.https://arxiv.org/abs/1312.6114}

Compress one screenshot crop

Start with a tiny crop. Each grayscale number is between 0 (dark) and 1 (bright). A real screenshot has many more pixels, but a 5 by 5 crop lets us inspect every single number. The bright 2 by 2 block is the highlighted region. Before looking at encoder maps, predict: which pixels should one latent coordinate watch, and which should it ignore?

View larger

An autoencoder has two learned functions:

The encoder $f_\phi$ maps the input $x$ to the latent code $z$. Meanwhile, the decoder $g_\theta$ maps that code to a reconstruction $\hat{x}$. Here $\phi$ and $\theta$ denote their learned weights. A narrow $z$ creates the hourglass bottleneck: the decoder receives input-specific information through those two coordinates. The network must discover compact summaries. A large network can still memorize a tiny dataset.

For this crop, the first encoder map averages the four bright center pixels and the second averages two lower pixels:

We chose these weights by hand, so the bottleneck arithmetic is visible before asking optimization to discover a useful map.

For normalized pixels, mean-squared error (MSE) measures reconstruction directly:

$N$ is the number of pixels, here 25. Each term is one pixel's squared error. The next snippet runs that same encoder, then a matching crafted decoder, so you can see the 25-to-2 shrink and the reconstruction error together.

import numpy as np

x = np.array([
    0.08, 0.12, 0.09, 0.11, 0.07,
    0.15, 0.85, 0.88, 0.14, 0.10,
    0.11, 0.90, 0.92, 0.13, 0.09,
    0.07, 0.18, 0.16, 0.12, 0.08,
    0.06, 0.09, 0.10, 0.07, 0.05,
], dtype=np.float32)

center = [6, 7, 11, 12]
lower = [16, 17]
W_encoder = np.zeros((25, 2), dtype=np.float32)
W_encoder[center, 0] = 0.25
W_encoder[lower, 1] = 0.50

W_decoder = np.zeros((2, 25), dtype=np.float32)
W_decoder[0, center] = 1.0
W_decoder[0, [5, 8, 10, 13]] = 0.12
W_decoder[1, lower] = 1.0
W_decoder[1, [15, 18, 21, 22]] = 0.35

z = x @ W_encoder
x_hat = np.clip(z @ W_decoder + 0.08, 0.0, 1.0)
mse = np.mean((x - x_hat) ** 2)

print("input numbers:", x.size)
print("latent numbers:", z.size, np.round(z, 3))
print("reconstruction MSE:", round(float(mse), 4))

input numbers: 25
latent numbers: 2 [0.888 0.17 ]
reconstruction MSE: 0.0027

The input shrinks from 25 numbers to 2. This crafted decoder does well on this one pattern because we built its weights around it. Notice what it loses: rearranging the four selected center pixels while preserving their average leaves $z_1$ unchanged. The decoder then produces the same reconstruction for different inputs. Learned weights can capture variation across a family of crops.

Learn the bottleneck

The next program makes noisy variations of the same crop and trains an autoencoder in PyTorch. The training loop uses gradients to update weights with an optimizer, as in the backpropagation lesson. Compare its MSE with a deliberately uninformative baseline that returns the average crop for every input.

import torch
from torch import nn
import torch.nn.functional as F

torch.manual_seed(3)
base = torch.tensor([
    0.08, 0.12, 0.09, 0.11, 0.07,
    0.15, 0.85, 0.88, 0.14, 0.10,
    0.11, 0.90, 0.92, 0.13, 0.09,
    0.07, 0.18, 0.16, 0.12, 0.08,
    0.06, 0.09, 0.10, 0.07, 0.05,
])
patches = torch.clamp(base + 0.025 * torch.randn(48, 25), 0.0, 1.0)

encoder = nn.Sequential(nn.Linear(25, 8), nn.Tanh(), nn.Linear(8, 2))
decoder = nn.Sequential(nn.Linear(2, 8), nn.Tanh(), nn.Linear(8, 25), nn.Sigmoid())
optimizer = torch.optim.Adam(
    list(encoder.parameters()) + list(decoder.parameters()), lr=0.03
)

with torch.no_grad():
    before = F.mse_loss(decoder(encoder(patches)), patches).item()

for _ in range(250):
    reconstruction = decoder(encoder(patches))
    loss = F.mse_loss(reconstruction, patches)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

with torch.no_grad():
    after = F.mse_loss(decoder(encoder(patches)), patches).item()
    latent = encoder(patches[:1])
    mean_crop = patches.mean(dim=0, keepdim=True).expand_as(patches)
    baseline = F.mse_loss(mean_crop, patches).item()

print(f"MSE before training: {before:.4f}")
print(f"MSE after training:  {after:.4f}")
print(f"constant-crop MSE:   {baseline:.4f}")
print("latent shape:", tuple(latent.shape))

MSE before training: 0.1537
MSE after training:  0.0006
constant-crop MSE:   0.0006
latent shape: (1, 2)

The trained MSE is close to the constant-crop baseline. These nearly identical inputs don't require an informative code: a decoder can do well by remembering their shared pattern. Optimization works as expected. To test whether the latent captures useful variation, train and evaluate on distinct patterns, hold out crops, and check whether perturbing the latent changes the reconstruction.

An autoencoder's reconstruction loss scores codes produced by its encoder for training inputs. A random pair can fall far from those codes, where the decoder has received little training signal:

The encoder can arrange codes without matching a chosen prior. It might place center-highlight crops near $(-1.8, 1.2)$ and lower-strip crops near $(1.7, -1.3)$. A standard-normal draw can fall in a region rarely reached by encoded training crops, where reconstruction loss offers little evidence about decoder quality. Outputs there may be useful, blurry, or implausible. The encoder may produce separated or overlapping clusters, and random samples may succeed or fail. A separate model could also learn a distribution over deterministic codes.

View larger

A VAE learns distributions over codes

A VAE starts with a generative story: draw a latent $z$ from a prior $p(z)$, then draw an image from a decoder distribution $p_\theta(x\mid z)$. For an observed crop, inference runs that story backward: which latent values could have produced it? The encoder predicts an approximate posterior, written $q_\phi(z \mid x)$, over those possible values.

The word approximate matters: computing the exact posterior $p_\theta(z \mid x)$ is generally impractical, so the encoder learns a tractable stand-in. This approximation method is variational inference. Here it's amortized: one encoder predicts posterior parameters for every crop, sharing its learned inference procedure across inputs.

For the common diagonal-Gaussian case, it predicts a mean vector $\mu$ and a log-variance vector $\log \sigma^2$:

The model also declares a simple prior, usually:

The posterior describes codes for one observed crop. The prior is the distribution we sample from when no input crop is provided. In $\mathcal{N}(0,I)$, each coordinate has mean zero and variance one, and the coordinates are independent. A diagonal posterior similarly has independent coordinates conditional on this crop, though its means and variances depend on the input. Mixing the posteriors of different crops needn't preserve that independence in the aggregate codes.

Why predict logvar instead of a raw standard deviation? An unconstrained linear layer can output negative values, while standard deviation must stay positive. Exponentiating half a finite log variance gives a positive scale. Floating-point exponentiation can still underflow to zero or overflow to infinity. Softplus offers another positive parameterization; log variance is one convenient choice.

import numpy as np

logvar = np.array([-0.70, 0.10])
sigma = np.exp(0.5 * logvar)
variance = np.exp(logvar)

print("log variance:", logvar)
print("variance:", np.round(variance, 3))
print("standard deviation:", np.round(sigma, 3))
print("all sigma positive:", bool(np.all(sigma > 0)))

log variance: [-0.7  0.1]
variance: [0.497 1.105]
standard deviation: [0.705 1.051]
all sigma positive: True

Those two vectors define a Gaussian for one crop. The encoder can place that Gaussian far from the numbers we later draw. What penalty pulls it toward a sampleable prior while preserving crop information?

Keep posteriors near a sampleable prior

If the encoder may put each crop's distribution anywhere it likes, prior samples still won't reliably land where the decoder has trained. Before writing the penalty, predict what it should do: it should be zero when the posterior matches the prior, and grow when means or variances move away.

VAEs measure that mismatch with Kullback-Leibler (KL) divergence. KL isn't a symmetric distance: direction matters. For a diagonal Gaussian posterior and standard-normal prior, the closed form is:^{[2]Reference 2Auto-Encoding Variational Bayes.https://arxiv.org/abs/1312.6114}

The penalty is zero when $\mu=0$ and $\sigma=1$. Moving means away from zero or making variances very different from one increases it.

Use the same two-dimensional posterior we'll sample from next: $\mu = [0.40, -0.20]$ and $\log\sigma^2 = [-0.70, 0.10]$. One coordinate at a time:

The next snippet checks that closed form on three posteriors, including this one.

import numpy as np

def kl_to_standard_normal(mu, logvar):
    # expm1(logvar) computes exp(logvar) - 1 accurately near zero.
    return float(0.5 * np.sum(mu**2 + np.expm1(logvar) - logvar))

posteriors = {
    "matches prior": (np.array([0.0, 0.0]), np.array([0.0, 0.0])),
    "small offset": (np.array([0.4, -0.2]), np.array([-0.7, 0.1])),
    "far offset": (np.array([2.0, -1.5]), np.array([-1.5, 1.2])),
}

for name, (mu, logvar) in posteriors.items():
    print(f"{name:13} KL = {kl_to_standard_normal(mu, logvar):.4f}")

matches prior KL = 0.0000
small offset  KL = 0.2009
far offset    KL = 4.0466

KL pressure has a cost. If it dominates reconstruction, all inputs can be encoded nearly like the prior and the decoder may ignore $z$. We'll check for that behavior after the training loop. First we need a way to sample from $q_\phi(z \mid x)$ while keeping gradients connected to $\mu$ and $\sigma$.

Sample while gradients still flow

Kingma and Welling reparameterize a Gaussian draw so reconstruction gradients can reach encoder parameters.^{[2]Reference 2Auto-Encoding Variational Bayes.https://arxiv.org/abs/1312.6114} Writing $z\sim\mathcal{N}(\mu,\sigma^2)$ specifies a distribution. The sampling API determines how the draw connects to autograd. PyTorch's ordinary Normal.sample() supplies no pathwise gradient to its mean or scale. Reparameterization makes that path explicit. Other estimators, including score-function gradients, can train stochastic models without differentiating the sampled value directly.^{[3]Reference 3Probability Distributionshttps://docs.pytorch.org/docs/2.14/distributions.html}

The reparameterization trick moves all stochasticity into an independent noise variable $\epsilon$ drawn from a standard normal distribution, then applies an algebraic transformation:

During the backward pass, $\epsilon$ is treated as a fixed constant tensor. The symbol $\odot$ means elementwise multiplication. For each coordinate $i$, the partial derivatives are:

Across all coordinates, the corresponding Jacobians are $\partial z/\partial\mu = I$ and $\partial z/\partial\sigma = \operatorname{diag}(\epsilon)$; derivatives between different coordinates are zero.

Reconstruction gradients $\partial L_{\text{recon}}/\partial z$ can then reach $\mu$, $\sigma$, and the encoder weights $\phi$. PyTorch's torch.distributions.Normal.rsample() implements this draw. Using sample() removes this reconstruction path to the encoder, but the decoder can still train, and an analytic KL term can still update encoder parameters.^{[3]Reference 3Probability Distributionshttps://docs.pytorch.org/docs/2.14/distributions.html}

View larger

import numpy as np

mu = np.array([0.40, -0.20])
logvar = np.array([-0.70, 0.10])
epsilon = np.array([0.50, -1.00])

sigma = np.exp(0.5 * logvar)
z = mu + sigma * epsilon

print("mu:", np.round(mu, 3))
print("sigma:", np.round(sigma, 3))
print("epsilon:", np.round(epsilon, 3))
print("sampled z:", np.round(z, 3))

mu: [ 0.4 -0.2]
sigma: [0.705 1.051]
epsilon: [ 0.5 -1. ]
sampled z: [ 0.752 -1.251]

Check the framework contract directly. Predict which draw keeps an autograd path, then verify the derivative through logvar: $\partial z_j/\partial\log\sigma_j^2=\tfrac12\sigma_j\epsilon_j=\tfrac12(z_j-\mu_j)$.

import torch
from torch.distributions import Normal

torch.manual_seed(11)
mu = torch.tensor([0.40, -0.20], requires_grad=True)
logvar = torch.tensor([-0.70, 0.10], requires_grad=True)
distribution = Normal(mu, torch.exp(0.5 * logvar))
ordinary = distribution.sample()
pathwise = distribution.rsample()
mu_gradient, logvar_gradient = torch.autograd.grad(
    pathwise.sum(), (mu, logvar), retain_graph=True
)
kl = 0.5 * (mu.square() + torch.expm1(logvar) - logvar).sum()
kl_mu_gradient = torch.autograd.grad(kl, mu)[0]

print("sample has an autograd path:", ordinary.requires_grad)
print("rsample has an autograd path:", pathwise.requires_grad)
print("rsample mean gradient:", mu_gradient.tolist())
print("logvar derivative matches:", torch.allclose(
    logvar_gradient, 0.5 * (pathwise.detach() - mu.detach())
))
print("analytic KL mean gradient:", [round(x, 3) for x in kl_mu_gradient.tolist()])

sample has an autograd path: False
rsample has an autograd path: True
rsample mean gradient: [1.0, 1.0]
logvar derivative matches: True
analytic KL mean gradient: [0.4, -0.2]

That sampled $z$ is what the decoder sees during training. The full path is now:

View larger

Generation later skips the encoder. It draws $z$ from the prior $p(z)$ and runs only the decoder. The path is differentiable. Those gradients still need a quantity to improve.

The VAE objective

The VAE performs variational inference by optimizing an evidence lower bound (ELBO) on the data log-likelihood.^{[2]Reference 2Auto-Encoding Variational Bayes.https://arxiv.org/abs/1312.6114} Written as a quantity to maximize:

The first term rewards a decoder that makes the observed crop likely after sampling its latent. The second discourages a posterior that moves too far from the prior. The bound reaches $\log p_\theta(x)$ when $q_\phi(z \mid x)$ matches the true posterior $p_\theta(z \mid x)$. The remaining gap is $D_{\mathrm{KL}}(q_\phi(z \mid x)\,\|\,p_\theta(z \mid x))$.

In code we minimize the negative ELBO, or a weighted variant:

For $\beta=1$, this is the negative ELBO when L_recon estimates the selected negative log-likelihood. To connect that term to the pixel errors you've already calculated, let the decoder output a mean image $\hat{x}$ and assume independent Gaussian pixel noise with fixed variance $\tau^2$:

For fixed $\tau$, the last term is constant with respect to the weights. The first term is scaled summed squared error. Mean MSE needs its own reduction factor. A smaller observation variance penalizes the same pixel errors more strongly. Changing the reconstruction reduction also changes the optimization problem unless you adjust the KL coefficient.

The reconstruction term averages over posterior samples. Our code uses one reparameterized draw per example, giving a Monte Carlo estimate of that expectation. For this fixed model and example, averaging $n$ independent draws divides the estimator variance by $n$, at the cost of more decoder work. The Gaussian KL term is evaluated analytically.^{[2]Reference 2Auto-Encoding Variational Bayes.https://arxiv.org/abs/1312.6114}

This ledger computes both terms for one crop using a fixed hypothetical reconstruction:

import numpy as np

x = np.array([0.10, 0.90, 0.85, 0.12])
x_hat = np.array([0.12, 0.84, 0.80, 0.14])
mu = np.array([0.40, -0.20])
logvar = np.array([-0.70, 0.10])

reconstruction_mse = np.mean((x - x_hat) ** 2)
kl = -0.5 * np.sum(1 + logvar - mu**2 - np.exp(logvar))

for beta in [0.0, 0.1, 1.0]:
    total = reconstruction_mse + beta * kl
    print(f"beta={beta:.1f}: recon={reconstruction_mse:.4f} KL={kl:.4f} total={total:.4f}")

beta=0.0: recon=0.0017 KL=0.2009 total=0.0017
beta=0.1: recon=0.0017 KL=0.2009 total=0.0218
beta=1.0: recon=0.0017 KL=0.2009 total=0.2026

View larger

Higgins et al. studied $\beta$-VAE, increasing the KL weight to restrict latent information and encourage disentangled representations.^{[4]Reference 4beta-VAE: Learning Basic Visual Concepts with a Constrained Variational Frameworkhttps://openreview.net/forum?id=Sy2fzU9gl} A factorized prior alone leaves learned coordinates free to mix real-world factors. Unsupervised disentanglement requires assumptions about the model and data.^{[5]Reference 5Challenging Common Assumptions in the Unsupervised Learning of Disentangled Representationshttps://proceedings.mlr.press/v97/locatello19a.html} Measure reconstruction, latent usefulness, and prior samples on your task; choose the KL weight from those results.

reconstruction_losses = [0.004, 0.010, 0.028]
kl_losses = [1.20, 0.42, 0.08]
betas = [0.01, 0.10, 1.00]

for beta, recon, kl in zip(betas, reconstruction_losses, kl_losses):
    total = recon + beta * kl
    print(f"beta={beta:>4.2f}: reconstruction={recon:.3f}, KL={kl:.2f}, objective={total:.3f}")

beta=0.01: reconstruction=0.004, KL=1.20, objective=0.016
beta=0.10: reconstruction=0.010, KL=0.42, objective=0.052
beta=1.00: reconstruction=0.028, KL=0.08, objective=0.108

These rows account for the objective terms. Each uses a different objective, so compare the totals as objective definitions. Log reconstruction and KL separately, then inspect samples at each setting to assess model quality.

Train a tiny VAE

Now train a two-dimensional VAE on two families of synthetic crops: one with a bright center highlight and one with a bright lower strip. A constant image can't fit both well. After training, compare reconstruction from each posterior mean with reconstruction from the other family's means. Swapping codes should hurt if they carry pattern information. Finally, decode three prior draws and inspect whether the center or lower strip is bright.

The decoder returns Gaussian mean images in this MSE-based example. We display those means without adding observation noise. The loss below is mean pixel error plus a chosen KL weight. This coefficient is an experimental setting with the mean-MSE scale defined above.

import torch
from torch import nn
import torch.nn.functional as F

torch.manual_seed(8)
center_mark = torch.tensor([
    0.05, 0.10, 0.08, 0.10, 0.06,
    0.12, 0.88, 0.86, 0.13, 0.08,
    0.10, 0.91, 0.90, 0.12, 0.07,
    0.06, 0.15, 0.14, 0.10, 0.06,
    0.05, 0.08, 0.08, 0.05, 0.04,
])
lower_mark = torch.tensor([
    0.04, 0.07, 0.06, 0.08, 0.05,
    0.10, 0.15, 0.12, 0.10, 0.06,
    0.08, 0.18, 0.16, 0.09, 0.06,
    0.06, 0.84, 0.88, 0.11, 0.07,
    0.04, 0.07, 0.08, 0.05, 0.04,
])
patches = torch.cat([
    torch.clamp(center_mark + 0.02 * torch.randn(32, 25), 0.0, 1.0),
    torch.clamp(lower_mark + 0.02 * torch.randn(32, 25), 0.0, 1.0),
])

class TinyVAE(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = nn.Sequential(nn.Linear(25, 12), nn.Tanh())
        self.mu = nn.Linear(12, 2)
        self.logvar = nn.Linear(12, 2)
        self.decoder = nn.Sequential(
            nn.Linear(2, 12), nn.Tanh(), nn.Linear(12, 25), nn.Sigmoid()
        )

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

model = TinyVAE()
optimizer = torch.optim.Adam(model.parameters(), lr=0.03)
beta = 0.02

def loss_terms(x_hat, x, mu, logvar):
    recon = F.mse_loss(x_hat, x)
    kl = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp()).sum(1).mean()
    return recon + beta * kl, recon, kl

for _ in range(400):
    x_hat, mu, logvar = model(patches)
    loss, recon, kl = loss_terms(x_hat, patches, mu, logvar)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

torch.manual_seed(99)
with torch.no_grad():
    x_hat, mu, logvar = model(patches)
    loss, recon, kl = loss_terms(x_hat, patches, mu, logvar)
    generated = model.decoder(torch.randn(3, 2))
    mean_mse = F.mse_loss(model.decoder(mu), patches)
    swapped_mse = F.mse_loss(model.decoder(mu.roll(32, dims=0)), patches)
    kl_per_dim = -0.5 * (1 + logvar - mu.square() - logvar.exp()).mean(0)

print(f"reconstruction={recon.item():.4f} KL={kl.item():.4f} total={loss.item():.4f}")
print("generated batch shape:", tuple(generated.shape))
print(f"posterior-mean MSE={mean_mse.item():.4f}; swapped-code MSE={swapped_mse.item():.4f}")
print("KL per dimension:", [round(value, 4) for value in kl_per_dim.tolist()])
for index, crop in enumerate(generated):
    center = crop[[6, 7, 11, 12]].mean().item()
    lower = crop[[16, 17]].mean().item()
    print(f"prior draw {index}: center={center:.3f}, lower={lower:.3f}")

reconstruction=0.0075 KL=0.7382 total=0.0222
generated batch shape: (3, 25)
posterior-mean MSE=0.0005; swapped-code MSE=0.1276
KL per dimension: [0.0005, 0.7378]
prior draw 0: center=0.142, lower=0.875
prior draw 1: center=0.877, lower=0.150
prior draw 2: center=0.800, lower=0.219

Each object in the program now has a job: mu and logvar define each approximate posterior, torch.randn_like supplies $\epsilon$, decoder consumes sampled $z$, and the logged losses reveal the compromise.

The VAE's shape contract is small enough to write down. Here, B is batch size:

The reductions matter too. F.mse_loss averages pixel errors over the batch, while the KL term sums its two latent dimensions before averaging over the batch. beta = 0.02 is specific to this setup. If you change latent width or loss reduction, retune it and keep logging both terms.

Swapping the families' codes raises MSE from 0.0005 to 0.1276, showing that the decoder uses input-dependent information. Prior draw 0 has a bright lower strip; draw 1 has a bright center; draw 2 is less distinct. Three draws can't establish distributional quality, and the reconstruction checks use training crops. Add held-out patterns and inspect full decoded grids before making a generalization claim.

Watch for posterior collapse

Low KL can look like success because each posterior matches the prior. But if reconstruction stays acceptable while changing $z$ barely changes the output, what happened? The decoder found a shortcut and stopped using its latent input. In information-theoretic terms, the mutual information $I(X;Z)$ between observations and latent codes can fall toward zero: knowing the sampled code no longer tells you much about which crop produced it. This failure is called posterior collapse. Van den Oord et al. describe it as a common issue for VAEs paired with high-capacity autoregressive decoders.^{[6]Reference 6Neural Discrete Representation Learning.https://arxiv.org/abs/1711.00937}

An expressive autoregressive decoder can model data through previously observed pixels or tokens while ignoring the global latent. Training can settle near $q_\phi(z\mid x)=p(z)$, with little input-specific information in the code.^{[7]Reference 7Generating Sentences from a Continuous Spacehttps://aclanthology.org/K16-1002/} That decoder can still generate varied, plausible outputs. An unconditional deterministic mean-image decoder optimized with MSE favors the dataset's mean image when it ignores $z$. Averaging follows from that decoder and loss; posterior collapse describes the broader failure to use latent information.

There are two pressures hidden in the average KL. Let $q(z)$ be the mixture of encoder posteriors over the data and let $I_q(X;Z)$ measure information under that joint data/encoder distribution. Then:^{[8]Reference 8ELBO Surgery: Yet Another Way to Carve Up the Variational Evidence Lower Boundhttps://www.cs.columbia.edu/~blei/fogm/2022F/readings/HoffmanJohnson2016.pdf}

Aggregate prior matching and per-crop posterior matching measure separate properties. If every posterior equals the prior, $I_q(X;Z)=0$: a sampled code cannot tell us which crop it came from. A large average KL can reflect prior mismatch, useful information, or both.

The training code already reports KL per latent dimension. One near-zero dimension can be unused capacity when the other distinguishes the two patterns. This illustrative monitor raises a warning only when KL is tiny in every dimension; it doesn't diagnose collapse by itself:

import numpy as np

logged_kl_by_dimension = {
    "using latent": np.array([0.31, 0.18]),
    "suspected collapse": np.array([0.0003, 0.0001]),
}
threshold = 0.001

for run_name, kl_dims in logged_kl_by_dimension.items():
    low_kl_warning = bool(np.all(kl_dims < threshold))
    print(f"{run_name:18} total_KL={kl_dims.sum():.4f} low_KL_warning={low_kl_warning}")

using latent       total_KL=0.4900 low_KL_warning=False
suspected collapse total_KL=0.0004 low_KL_warning=True

Interpret a low KL value alongside reconstruction quality, generated samples, and whether predictions respond when $z$ changes. A rapid collapse toward zero still warrants investigation.

Three interventions worth testing are:

1. KL annealing: Start with weak KL pressure and increase it toward the intended final weight. Bowman et al. anneal from zero to one for their likelihood-scaled objective.^{[7]Reference 7Generating Sentences from a Continuous Spacehttps://aclanthology.org/K16-1002/} Our mean-MSE example uses a different scale, so choose the endpoint to match its loss scale. The nearly constant-crop experiment showed that zero KL weight can leave the latent unused.
2. Free bits: The original formulation sums $\max(\tau,\mathbb{E}_X D_{\mathrm{KL},j})$ over latent groups, with each group's KL averaged over the minibatch before applying the floor.^{[9]Reference 9Improved Variational Inference with Inverse Autoregressive Flowhttps://arxiv.org/abs/1606.04934} Below $\tau$, that KL contribution has zero gradient; reconstruction can reward additional information without extra KL cost. Mutual information can still remain low. With natural logarithms, the threshold is measured in nats despite the name.
3. Decoder rebalancing: Reduce decoder capacity or weaken autoregressive conditioning. Bowman's word dropout removes some teacher-forced context so latent information becomes more useful.^{[7]Reference 7Generating Sentences from a Continuous Spacehttps://aclanthology.org/K16-1002/} Test whether the decoder uses $z$ after the change, and monitor sample quality because reducing decoder capacity can hurt it.

The order in free bits is observable. Suppose one latent group has per-example KL values 0.2 and 2.2 nats, and its floor is 1.0. Its batch average is 1.2, already above the floor, so the original formulation keeps the KL gradient for both examples. Flooring each example first would suppress the first example's gradient and charge a different total:

import torch

kl_values = torch.tensor([0.2, 2.2], dtype=torch.float64, requires_grad=True)
floor = 1.0
group_floor = kl_values.mean().clamp_min(floor)
example_floor = kl_values.clamp_min(floor).mean()
group_grad = torch.autograd.grad(group_floor, kl_values, retain_graph=True)[0]
example_grad = torch.autograd.grad(example_floor, kl_values)[0]
print(f"floor after group average: {group_floor.item():.1f}")
print("derivatives:", group_grad.tolist())
print(f"average after example floors: {example_floor.item():.1f}")
print("derivatives:", example_grad.tolist())

floor after group average: 1.2
derivatives: [0.5, 0.5]
average after example floors: 1.6
derivatives: [0.0, 0.5]

These are derivatives with respect to the KL values. Autograd then multiplies them by each value's derivative with respect to the encoder parameters. The example avoids the exact threshold, where the floor has a corner; below the group threshold, that group's KL contribution has zero derivative for every example.

Inspect interpolation to see whether it is meaningful. A decoder built from continuous layers changes continuously with $z$ in a deterministic autoencoder too; mathematical continuity alone says little about image plausibility. The path $z(\alpha)=(1-\alpha)z_A+\alpha z_B$ can cross regions poorly represented during training in either model. Gaussian posterior support alone doesn't tell us where most probability mass lies or where decoding works well.

VQ-VAE: replace continuous samples with codes

A Vector Quantised Variational Autoencoder (VQ-VAE) takes a different route. Its encoder output is mapped to the nearest vector in a learned codebook, producing a discrete latent representation. The original VQ-VAE paper also learns a prior over those discrete codes.^{[6]Reference 6Neural Discrete Representation Learning.https://arxiv.org/abs/1711.00937}

For one encoder vector, quantization is a nearest-neighbor lookup:

import numpy as np

encoder_output = np.array([0.72, -0.10])
codebook = np.array([
    [-0.80, 0.20],
    [0.65, -0.05],
    [0.10, 0.90],
])

squared_distances = ((codebook - encoder_output) ** 2).sum(axis=1)
index = int(squared_distances.argmin())
quantized = codebook[index]

print("squared distances:", np.round(squared_distances, 3))
print("selected code ID:", index)
print("decoder receives:", quantized)

squared distances: [2.4   0.007 1.384]
selected code ID: 1
decoder receives: [ 0.65 -0.05]

Nearest-code assignments are piecewise constant, so their ordinary derivatives don't supply a useful encoder training signal. VQ-VAE uses a straight-through gradient estimator for the encoder path, plus a codebook loss and a commitment loss so encoder outputs stay close to useful entries:^{[6]Reference 6Neural Discrete Representation Learning.https://arxiv.org/abs/1711.00937}

Here $\operatorname{sg}[\cdot]$ stops gradients, $e$ is the selected codebook vector, and this $\beta$ weights the VQ-VAE commitment term. In a $\beta$-VAE, $\beta$ weights the KL term. In the original straight-through setup, reconstruction trains the decoder and encoder; it supplies no codebook gradient. The codebook loss updates selected entries, while commitment updates the encoder. A moving-average codebook update is another option.^{[6]Reference 6Neural Discrete Representation Learning.https://arxiv.org/abs/1711.00937} The important distinction is the interface:

The nearest code doesn't change during this backward pass. Which tensors receive gradients from reconstruction alone, and which need the auxiliary losses? This two-coordinate example makes the routing visible:

import torch

encoded = torch.tensor([0.72, -0.10], requires_grad=True)
codebook = torch.tensor([
    [-0.80, 0.20], [0.65, -0.05], [0.10, 0.90]
], requires_grad=True)
index = ((codebook - encoded) ** 2).sum(1).argmin()
quantized = codebook[index]
straight_through = encoded + (quantized - encoded).detach()
target = torch.tensor([0.50, 0.10])
reconstruction = 0.5 * (straight_through - target).square().sum()
encoder_recon, codebook_recon = torch.autograd.grad(
    reconstruction, (encoded, codebook), retain_graph=True, allow_unused=True
)
codebook_loss = (encoded.detach() - quantized).square().sum()
commitment = 0.25 * (encoded - quantized.detach()).square().sum()
(reconstruction + codebook_loss + commitment).backward()

rounded = lambda vector: [round(x, 3) for x in vector.tolist()]
print("forward code:", rounded(straight_through.detach()))
print("reconstruction encoder gradient:", rounded(encoder_recon))
print("reconstruction codebook gradient:", codebook_recon)
print("total encoder gradient:", rounded(encoded.grad))
print("selected code gradient:", rounded(codebook.grad[index]))
assert torch.count_nonzero(codebook.grad[[0, 2]]).item() == 0

forward code: [0.65, -0.05]
reconstruction encoder gradient: [0.15, -0.15]
reconstruction codebook gradient: None
total encoder gradient: [0.185, -0.175]
selected code gradient: [-0.14, 0.1]

A hand-picked VQ-VAE code ID can be decoded. Plausible sampling depends on drawing code-ID combinations from the learned prior, which captures combinations used by the generative model during training.

Continuous VAEs and discrete VQ-VAEs can both turn a large input into a smaller latent another model can consume. A later generator can model those latents directly, reducing the work of modeling every input pixel.

Latent models reduce the expensive space

The same bottleneck contract scales beyond this tiny crop. A later generator can work in latent space, then hand its final representation to the decoder that turns it back into pixels. Rombach et al. train diffusion models in the latent space of pretrained autoencoders, applying denoising operations to the compressed representation.^{[10]Reference 10High-Resolution Image Synthesis with Latent Diffusion Models.https://arxiv.org/abs/2112.10752} They write the spatial downsampling factor as $f = H/h = W/w$ and evaluate the efficiency and quality trade-off across values from 2 to 32, including $f=4$ and $f=8$. The autoencoder's job is precise: compress before the expensive generation steps, then decode a final latent back to pixels.

Their compression stage uses perceptual and adversarial losses, with weak KL or vector-quantization regularization. A separate diffusion model learns the latent distribution.^{[10]Reference 10High-Resolution Image Synthesis with Latent Diffusion Models.https://arxiv.org/abs/2112.10752} Sampling Gaussian noise and decoding it directly isn't equivalent to running that generator: diffusion must first transform noise into a latent that fits the learned distribution.

You can check the size reduction by hand. A downsampling factor of 8 reduces each spatial axis by 8, so it reduces spatial positions by 64. On a 512 by 512 grid that is 262,144 positions becoming 4,096:

height, width = 512, 512
downsample = 8
latent_height = height // downsample
latent_width = width // downsample

pixel_positions = height * width
latent_positions = latent_height * latent_width

print("pixel grid:", (height, width), "positions:", pixel_positions)
print("latent grid:", (latent_height, latent_width), "positions:", latent_positions)
print("spatial reduction:", pixel_positions // latent_positions, "x")

pixel grid: (512, 512) positions: 262144
latent grid: (64, 64) positions: 4096
spatial reduction: 64 x

Channels and decoder quality still matter. For a hypothetical RGB image and four-channel latent at $f=8$, the spatial area shrinks by 64 while the channel count grows from 3 to 4, for a 48-fold reduction in scalar values. Those element counts alone do not determine runtime or compressed-file size. Rombach et al. find that aggressive compression hurts their downstream results; detail discarded by the encoder can't be uniquely recovered from its code. A generator may synthesize plausible detail without reproducing the original pixels.

Debug the latent interface

Start by logging reconstruction and KL losses separately. A low total can hide a reconstruction term that dominates, or a KL term that has collapsed toward zero.

If sampling or KL returns NaN or Inf, check the scale path first: compute standard deviation as sigma = exp(0.5 * logvar). The expression exp(logvar) gives variance. Inspect the range of logvar before adding an explicit clamp, since a clamp can hide unstable optimization.

Keep the two distributions distinct. The encoder supplies $q_\phi(z \mid x)$ for an observed crop; generation samples from the prior $p(z)$. If decoded outputs barely change when $z$ changes, the decoder may be ignoring its latent input even when reconstruction looks fine.

For a VQ-VAE, log code usage as well. If only a few code IDs are selected, the codebook's capacity is going unused and the discrete bottleneck needs a closer look.

One loss number mixes three checks: whether the crop was rebuilt, whether a random draw lands where the decoder was trained, and whether changing the code changes the picture.

Test what the code preserves