Disentanglement, beta-VAE & Posterior Collapse

"Give me one knob for rotation, one for size, one for color, and I will redecorate the universe one dimension at a time. Tangle them together and I am just a very expensive lottery ticket."

A Latent Axis That Wants a Job Description

You ended Section 31.3 by adding a coefficient $\beta$ to the KL term and watching reconstruction trade against latent regularity. That experiment was the doorway to this section. The plain VAE gives you a usable latent, but you have no control over what its dimensions mean: dimension 3 might mix digit identity, slant, and thickness all at once. For many applications, the controllable generation and image editing of Chapter 35 chief among them, you want each latent dimension to be a clean, separate control. That is disentanglement, and the beta-VAE is the simplest method that pursues it.

1. What Disentanglement Means Intermediate

Real images are generated by a small number of independent factors of variation, the latent-variable picture introduced in Chapter 30. A rendered sprite is determined by its shape, size, rotation, and position; a face by pose, lighting, identity, and expression. A latent representation is disentangled if changing one latent dimension changes exactly one of these factors while leaving the others fixed. Disentanglement is desirable because it makes the latent interpretable and editable: to make a face look left you turn the pose knob, with no side effects on identity or lighting. It also tends to generalize better, since a model that has separated the factors can recombine them in ways it never saw during training.

A plain VAE rarely disentangles on its own. Nothing in the ELBO of Section 31.3 rewards aligning latent axes with factors; the encoder is free to entangle them however minimizes reconstruction. The standard-normal prior is factorized (its dimensions are independent), so there is a weak pull toward independence, but reconstruction usually overwhelms it. The beta-VAE strengthens that pull deliberately.

2. The beta-VAE Intermediate

The beta-VAE changes the ELBO by a single scalar. Multiply the KL term by $\beta > 1$:

Why does this encourage disentanglement? A larger $\beta$ makes latent capacity expensive: the model is penalized harder for every bit of information it stores in the code and for any deviation from the factorized prior. Forced to economize, the encoder uses its scarce latent budget on the most reconstruction-relevant factors and, because the prior is factorized, prefers to put independent factors on independent axes. The result on simple, controlled datasets is that individual dimensions come to control individual factors. The cost is explicit and unavoidable: a higher $\beta$ blurs reconstructions, because latent capacity that would have carried fine detail has been taxed away. Figure 31.4.1 illustrates the spectrum from entangled to disentangled to collapsed as $\beta$ grows.

The implementation is trivial given the VAE of Section 31.3: one multiplier. The next block shows the loss and the latent-traversal procedure that reveals what each axis learned, by holding all dimensions fixed and sweeping one.

# The beta-VAE is the Section 31.3 loss with one multiplier on the KL term,
# plus a latent-traversal diagnostic: sweep a single latent dimension while
# freezing the rest to reveal which visual factor that axis has captured.
import torch

def beta_vae_loss(x_hat, x, mu, logvar, beta=4.0):
    recon = torch.nn.functional.binary_cross_entropy(x_hat, x, reduction="sum")
    kl = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return recon + beta * kl                 # the only change from Section 31.3

# Latent traversal: see what each axis controls.
model.eval()
with torch.no_grad():
    z = torch.zeros(1, 10)                    # start at the prior mean
    for dim in range(10):                     # one row per latent dimension
        row = []
        for val in torch.linspace(-3, 3, 8):  # sweep this dim from -3 to +3
            z_t = z.clone(); z_t[0, dim] = val
            row.append(model.dec(z_t).view(28, 28))
        # On a disentangled model, each row varies ONE visual factor
        # (e.g. row 2 rotates, row 5 changes thickness) and others stay fixed.

3. Measuring and Doubting Disentanglement Advanced

"Each axis controls one factor" needs a number to be a claim. Because measuring it requires knowing the true factors, disentanglement metrics are evaluated on synthetic datasets with controlled ground truth, most famously dSprites, where every image is a known combination of shape, scale, rotation, and position. Given that ground truth, several metrics quantify the alignment: the Mutual Information Gap (MIG) measures, for each true factor, how much more the single most-informative latent dimension tells you about it than the second-most-informative one, so a high gap means one factor maps to one axis. Other metrics (the FactorVAE score, the DCI framework, the original beta-VAE metric) formalize related notions of how cleanly latents and factors line up. The next block computes MIG schematically.

# Mutual Information Gap (MIG): for each ground-truth factor, measure how
# much more the best latent dimension knows about it than the runner-up.
# A large average gap means factors map cleanly to single axes.
import numpy as np
from sklearn.metrics import mutual_info_score

def mutual_information_gap(codes, factors, n_bins=20):
    """codes: (N, D) latent means; factors: (N, K) ground-truth factors."""
    D, K = codes.shape[1], factors.shape[1]
    mi = np.zeros((K, D))
    for k in range(K):
        for d in range(D):
            cz = np.digitize(codes[:, d], np.histogram_bin_edges(codes[:, d], n_bins))
            mi[k, d] = mutual_info_score(factors[:, k], cz)
    gaps = []
    for k in range(K):
        order = np.sort(mi[k])[::-1]               # MI of each latent with factor k
        entropy_k = mutual_info_score(factors[:, k], factors[:, k])
        gaps.append((order[0] - order[1]) / (entropy_k + 1e-9))  # top minus runner-up
    return float(np.mean(gaps))                    # higher = more disentangled
# A beta-VAE at beta=4 typically scores higher MIG than a plain VAE (beta=1),
# confirming the KL pressure aligned axes with factors, at some reconstruction cost.

4. Posterior Collapse Advanced

The opposite failure is more insidious because it can happen even at $\beta = 1$. Posterior collapse occurs when the encoder's distribution $q_\phi(z \mid x)$ becomes identical to the prior for every input, so the KL term hits zero and the latent code carries no information about $x$. The decoder, meanwhile, has learned to produce reasonable outputs while ignoring the useless code entirely. You are left with a VAE whose latent input does nothing: change $z$ and the output barely moves. This is fatal for a generative model, because the latent is supposed to be the source of variation. The illustration below pictures collapse as a decoder that has stopped reading its mail.

Collapse arises from a race at the start of training. If the decoder is powerful enough to model the data on its own (a deep autoregressive or heavily-skip-connected decoder is the usual culprit), then early in training, when the code is still uninformative noise, the cheapest way to lower the loss is to drive the KL to zero (free, since it removes a penalty) and let the strong decoder carry the reconstruction. Once the code is ignored, no gradient pushes the encoder to make it informative again, and the model is stuck. The standard fixes all work by protecting the latent during this vulnerable early phase. KL annealing ramps the KL weight from zero up to one over the first epochs, so the model first learns to use the code and only later is asked to regularize it. The free-bits trick gives each latent dimension a small KL "allowance" that is not penalized, so the model is not rewarded for collapsing dimensions to exactly zero. The next block implements both.

# Two defenses against posterior collapse: a per-dimension free-bits floor
# that stops the optimizer being rewarded for zeroing latent dimensions,
# and a KL-annealing schedule that ramps the KL weight up over early epochs.
import torch

def collapse_resistant_loss(x_hat, x, mu, logvar, kl_weight, free_bits=0.5):
    recon = torch.nn.functional.binary_cross_entropy(x_hat, x, reduction="sum")
    # Per-dimension KL, then clamp each to a floor of `free_bits` nats.
    kl_per_dim = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp())  # (batch, dim)
    kl = torch.clamp(kl_per_dim, min=free_bits).sum()            # free bits
    return recon + kl_weight * kl                               # kl_weight is annealed

# KL annealing schedule: ramp the weight from 0 to 1 over the first 10 epochs.
for epoch in range(num_epochs):
    kl_weight = min(1.0, epoch / 10.0)        # 0, 0.1, 0.2, ... up to 1.0
    for x, _ in loader:
        x_hat, mu, logvar = model(x)
        loss = collapse_resistant_loss(x_hat, x, mu, logvar, kl_weight)
        opt.zero_grad(); loss.backward(); opt.step()
    active = (mu.var(0) > 1e-2).sum().item()  # count latent dims actually used
    print(f"epoch {epoch}: kl_w={kl_weight:.1f}, active latent dims={active}")
# Without these tricks a strong decoder drives active dims toward 0 (collapse);
# with annealing + free bits, the latent stays informative.