73 - VAE: The Beta Weight, Posterior Collapse and Training Instabilities

A badly trained VAE does not give you an error: the loss goes down, the program ends, the images come out. The defect is inside, and you only see it if you know what to measure. In the previous lesson the VAE worked on the first try, but that is a lucky case: all its ingredients were in balance. Here we break the balance, one ingredient at a time, and watch what breaks: the weight β between reconstruction and KL, posterior collapse (code dimensions that stop carrying information) and numerical instabilities. For each one there is a measurement that reveals it in a few seconds. All the numbers come from real runs, repeated twice with identical outputs, on 8×8 handwritten digits (scikit-learn's load_digits) and not on MNIST: the run did not download any dataset, and the MNIST figures from books do not apply here.

The words of this lesson

We use the terms of the two previous lessons (ELBO, KL, prior, reparametrisation, active dimension) and need only a few more. A β-VAE is a VAE in which the KL is multiplied by a coefficient: loss = reconstruction + β · KL. With β = 1 it is the standard VAE, that is the true ELBO; the name comes from Higgins and colleagues (2017), who use β > 1 to obtain more "orderly" representations (attribution not verified on the text, see the sources). Posterior collapse happens when, for one or more latent dimensions, the encoder produces the same N(0, 1) whatever the image: the KL of that dimension is zero and the decoder ignores it, for it that is a dead dimension.

KL annealing (or warm-up) starts with β = 0 and gradually brings it to 1 in the first epochs, to give the model time to learn to use the code before penalising it. The aggregate posterior, q(z), is the average of q(z | x) over the whole dataset, that is the distribution of the codes the encoder really produces: to generate with z ~ N(0, I) it must resemble the prior. Finally the rate-distortion trade-off: the KL is the cost of the code in nats (rate), the reconstruction error is the distortion, and each value of β picks a point on that curve.

The analogy: the courier's budget

Imagine a courier who has to transmit an image to a colleague and pays per nat: the more details he sends, the more he spends. β is the price per nat. If the message costs nothing (β = 0) the courier sends everything, with insane precision, and the colleague reconstructs very well; but the messages are written in a dialect that only the two of them know, so nobody else can order "a random message" and get anything sensible. If it costs too much, the courier sends nothing any more and the colleague always answers with the same mean image. The standard VAE sits halfway, and the experiment that follows says where the boundaries lie.

The starting point, self-contained

The experiments reuse the model and the loss of the previous lesson, but rewritten here in full, so the block can be copied and run on its own. Two things change: the code dimension latent is a parameter, because we will vary it, and at the bottom there is the "judge", a small digit classifier trained separately. The judge does not measure perceptual quality: it is used to count, over 2000 generated images, which digits appear and how often, that is the diversity of the samples.

import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

torch.set_num_threads(2)
X, y = load_digits(return_X_y=True)                                # 1797 digits, 8x8 pixels, values 0..16
Xtr, Xte, ytr, yte = train_test_split(X / 16.0, y, test_size=0.3, random_state=0, stratify=y)
Xtr, Xte = torch.tensor(Xtr, dtype=torch.float32), torch.tensor(Xte, dtype=torch.float32)
ytr, yte = torch.tensor(ytr), torch.tensor(yte)
D = 64

class VAE(nn.Module):
    def __init__(self, latent=10, hidden=128):
        super().__init__()
        self.enc = nn.Sequential(nn.Linear(D, hidden), nn.ReLU())
        self.mu = nn.Linear(hidden, latent)
        self.logvar = nn.Linear(hidden, latent)
        self.dec = nn.Sequential(nn.Linear(latent, hidden), nn.ReLU(), nn.Linear(hidden, D))

    def encode(self, x):
        h = self.enc(x)
        return self.mu(h), self.logvar(h)

    def reparametrize(self, mu, logvar):
        return mu + torch.exp(0.5 * logvar) * torch.randn_like(mu)

    def decode(self, z):
        return self.dec(z)                                          # logits

    def forward(self, x):
        mu, logvar = self.encode(x)
        return self.decode(self.reparametrize(mu, logvar)), mu, logvar

def vae_loss(logits, x, mu, logvar):
    """Negative ELBO per image in nats: reconstruction (BCE summed over pixels) + KL (summed over dimensions)."""
    recon = F.binary_cross_entropy_with_logits(logits, x, reduction="sum") / x.size(0)
    kl = (0.5 * (mu.pow(2) + logvar.exp() - logvar - 1)).sum() / x.size(0)
    return recon + kl, recon, kl

# The "judge": a small classifier that tells us which digit it sees in a generated image
torch.manual_seed(0)
clf = nn.Sequential(nn.Linear(D, 64), nn.ReLU(), nn.Linear(64, 10))
opt_c = torch.optim.Adam(clf.parameters(), lr=1e-3)
for _ in range(60):
    perm = torch.randperm(len(Xtr))
    for i in range(0, len(Xtr), 64):
        idx = perm[i:i + 64]
        loss = F.cross_entropy(clf(Xtr[idx]), ytr[idx]); opt_c.zero_grad(); loss.backward(); opt_c.step()
with torch.no_grad():
    print("judge accuracy on real test digits:", round((clf(Xte).argmax(-1) == yte).float().mean().item(), 4))

judge accuracy on real test digits: 0.9778

The vae_loss loss sums the BCE over the 64 pixels and the KL over the dimensions, then divides by the number of images: the result is in nats per image, and it is the negative of the ELBO. The judge reaches 97.8% on the test set, enough to trust the counts we derive from it; on the blurry images of the samples, however, it will assign a digit anyway, so a balanced count says that no class is missing, not that the images are sharp.

Experiment 1: sweeping β

We train the same VAE for five values of β (0, 0.25, 1, 2, 4), with a 10-dimension code, 60 epochs, one seed. For each we measure reconstruction and KL on the test set, always with the true loss (β = 1), the MSE decoding the mean, the number of active dimensions (variance of μ on the test set > 0.01) and the entropy of the digits the judge sees in 2000 samples from N(0, I): the maximum is ln 10 = 2.30, all digits equiprobable.

def train_vae(beta=1.0, latent=10, epochs=60, seed=0, warmup_epochs=0):
    """beta weights the KL term: loss = reconstruction + beta * KL. With warmup_epochs, beta grows linearly from 0 to its value."""
    torch.manual_seed(seed)
    model = VAE(latent=latent)
    opt = torch.optim.Adam(model.parameters(), lr=1e-3)
    n_batches = (len(Xtr) + 63) // 64
    step = 0
    for epoch in range(epochs):
        perm = torch.randperm(len(Xtr))
        for i in range(0, len(Xtr), 64):
            x = Xtr[perm[i:i + 64]]
            logits, mu, logvar = model(x)
            _, recon, kl = vae_loss(logits, x, mu, logvar)
            frac = min(1.0, step / (warmup_epochs * n_batches)) if warmup_epochs > 0 else 1.0
            loss = recon + beta * frac * kl                       # the weight of the KL, possibly ramped up
            opt.zero_grad(); loss.backward(); opt.step()
            step += 1
    return model

@torch.no_grad()
def report(model, latent):
    logits, mu, logvar = model(Xte)
    _, recon, kl = vae_loss(logits, Xte, mu, logvar)            # always the true ELBO (beta = 1)
    mse = F.mse_loss(torch.sigmoid(model.decode(mu)), Xte).item()
    active = int((mu.var(dim=0) > 0.01).sum())
    samples = torch.sigmoid(model.decode(torch.randn(2000, latent)))
    hist = torch.bincount(clf(samples).argmax(-1), minlength=10).float() / 2000
    entropy = max(0.0, -(hist[hist > 0] * hist[hist > 0].log()).sum().item())
    return recon.item(), kl.item(), mse, active, entropy

The train_vae function is the training of the previous lesson with three differences. The beta coefficient multiplies the KL in the loss. The latent parameter varies the code dimension. The warmup_epochs parameter, if positive, makes the factor frac grow linearly from 0 to 1 (step / (warmup_epochs · n_batches)) and uses it in place of 1. The report function, instead, does not know β: it always evaluates with vae_loss, that is with the true ELBO, so the rows of the table can be compared with each other. If you evaluated each model with the loss it was trained with, you would be comparing different functions.

models = {}
print(f"{'beta':>5s} {'recon':>7s} {'KL':>6s} {'MSE':>7s} {'active':>7s} {'class entropy (max 2.30)':>26s}")
for beta in (0.0, 0.25, 1.0, 2.0, 4.0):
    models[beta] = train_vae(beta=beta)
    recon, kl, mse, active, entropy = report(models[beta], 10)
    print(f"{beta:5.2f} {recon:7.1f} {kl:6.1f} {mse:7.4f} {active:4d}/10 {entropy:14.2f}")

 beta   recon     KL     MSE  active   class entropy (max 2.30)
 0.00    16.7  137.2  0.0127   10/10           1.64
 0.25    18.2    9.7  0.0148   10/10           2.29
 1.00    21.9    2.9  0.0302   10/10           2.28
 2.00    25.3    0.9  0.0491    4/10           2.11
 4.00    27.2    0.0  0.0743    0/10           0.00

The chart relates the two central columns: each point is a value of β, the horizontal axis is the KL on a logarithmic scale, the vertical one the reconstruction. The β = 4 point is missing, because it has KL zero and cannot be drawn on a logarithmic scale.

Let us read the table from the top. With β = 0 no penalty touches the code: the KL explodes to 137 nats per image (against 2.9 at β = 1), the reconstruction is the best (16.7 nats, MSE 0.0127) and the model is, in effect, an autoencoder. But the class entropy drops to 1.64: generating from N(0, I) does not work, the same defect seen with the autoencoder of lesson 70. With β = 0.25 the KL falls to 9.7, the reconstruction stays close (18.2) and the entropy is already 2.29, the highest in the table: diversified samples with almost all of the reconstruction quality. With β = 1 (the true ELBO) the values are KL 2.9, reconstruction 21.9, MSE 0.0302, entropy 2.28.

Then the behaviour changes abruptly, not gradually. At β = 2 only 4 dimensions out of 10 remain active, the reconstruction rises to 25.3 and the MSE to 0.0491. At β = 4 the active dimensions are zero: KL 0.0, reconstruction 27.2, entropy 0.00. This is a total posterior collapse: the encoder gives every image the same N(0, I), the decoder receives only noise and learns to always emit the same image. We verify it, because the coincidence is suspicious: if the β = 4 model emits the mean training image, its reconstruction must coincide with that of the mean image.

mean = Xtr.mean(0).clamp(1e-6, 1 - 1e-6)                        # the average training image
bce = -(Xte * mean.log() + (1 - Xte) * (1 - mean).log()).sum(1).mean().item()
mse = F.mse_loss(mean.expand_as(Xte), Xte).item()
print(f"average image used for every digit: reconstruction {bce:.1f} nat, MSE {mse:.4f}")
with torch.no_grad():
    s = torch.sigmoid(models[4.0].decode(torch.randn(2000, 10)))
    print("digits the judge sees in 2000 samples of the beta=4 model:", torch.bincount(clf(s).argmax(-1), minlength=10).tolist())

average image used for every digit: reconstruction 27.2 nat, MSE 0.0739
digits the judge sees in 2000 samples of the beta=4 model: [0, 0, 0, 0, 0, 0, 0, 0, 2000, 0]

It matches: 27.2 nats and MSE 0.0739 (against 0.0743 for the β = 4 model), and the judge sees the same digit in all 2000 samples. The code at β = 4 carries nothing, and the decoder has become a constant. The point on the scale where this happens depends on data, architecture and training length: on MNIST with larger decoders the values of β will be different. Here it is one seed and one configuration: the threshold between "4 active dimensions" and "none" is not a constant, the shape of the curve is.

No β is right in absolute terms. To compress well you pick a small β; for an orderly, samplable space you stay near 1; beyond a certain value you buy simplicity of the code with a useless model. The authors of the β-VAE claim that β > 1 favours "disentangled" factors: we did not measure that here, we only measured the loss of active dimensions.

Experiment 2: posterior collapse

Collapse does not happen only with a large β. We give the VAE a code of 40 dimensions instead of 10 and count how many it uses, with β = 1; then we try the KL warm-up (β grows from 0 to 1 in 10 epochs) and finally β = 3.

print(f"{'setting':38s} {'recon':>7s} {'KL':>6s} {'MSE':>7s} {'active':>9s}")
for name, kwargs in [("latent 10, beta 1",                     dict(latent=10)),
                     ("latent 40, beta 1",                     dict(latent=40)),
                     ("latent 40, beta 1, KL warm-up 10 ep",   dict(latent=40, warmup_epochs=10)),
                     ("latent 40, beta 3",                     dict(latent=40, beta=3.0))]:
    m = train_vae(**kwargs)
    recon, kl, mse, active, entropy = report(m, kwargs["latent"])
    print(f"{name:38s} {recon:7.1f} {kl:6.1f} {mse:7.4f} {active:5d}/{kwargs['latent']}")

setting                                  recon     KL     MSE    active
latent 10, beta 1                         21.9    2.9  0.0302    10/10
latent 40, beta 1                         23.0    2.4  0.0367    38/40
latent 40, beta 1, KL warm-up 10 ep       22.6    2.6  0.0343    33/40
latent 40, beta 3                         27.1    0.0  0.0789     1/40

With 40 dimensions available the VAE uses 38 of them, and reconstruction and KL stay almost the same as with 10 (23.0 and 2.4 nats against 21.9 and 2.9): the information needed for the digits sits in a few dimensions, and the others are switched off by the KL because they are not needed. This is a benign collapse, a regularisation that works. The warm-up left 33 active dimensions and −ELBO 25.2 (22.6 + 2.6) against 25.4 (23.0 + 2.4) without warm-up: a difference of 0.2 nats, with one seed, which cannot be told apart from chance. With β = 3 the collapse is almost total: 1 dimension out of 40, reconstruction 27.1.

Why spend attention on something benign? Because in other contexts collapse is fatal. Bowman and colleagues (2016) train a VAE for sentences with an LSTM decoder and observe that the decoder learns to explain the data without looking at z: at that point no signal flows any more between decoder and encoder, and the model settles in an equilibrium with KL zero, equivalent to an ordinary language model. In their paper standard training, with no remedies, indeed produces models with performance equal to the baseline language model and KL zero. That is why they propose KL cost annealing (KL weight at zero, then growing) and word dropout. Our decoder is made of dense layers and cannot model an image without z, so the serious case does not arise. The general rule stands: count the active dimensions every time you train a VAE.

Diagnostics: does the aggregate posterior resemble the prior?

To generate with z ~ N(0, I) the codes the encoder produces on the whole dataset must be distributed more or less like N(0, I). The quick check: for each dimension, the mean and variance of the z on the test set, which should be about 0 and about 1. We reuse the models trained in experiment 1.

@torch.no_grad()
def aggregate_posterior(model, X):
    mu, logvar = model.encode(X)
    z = model.reparametrize(mu, logvar)                       # one z per image, as the decoder sees it in training
    return z.mean(dim=0), z.var(dim=0)

for beta in (0.0, 1.0, 4.0):                                  # the models trained in the sweep above
    mean_z, var_z = aggregate_posterior(models[beta], Xte)
    print(f"beta={beta:3.1f}   avg |mean| of z per dim {mean_z.abs().mean():.2f}   variance of z per dim: min {var_z.min():.2f}  mean {var_z.mean():.2f}  max {var_z.max():.2f}")

beta=0.0   avg |mean| of z per dim 2.82   variance of z per dim: min 4.13  mean 8.75  max 17.75
beta=1.0   avg |mean| of z per dim 0.05   variance of z per dim: min 0.88  mean 1.00  max 1.11
beta=4.0   avg |mean| of z per dim 0.05   variance of z per dim: min 0.95  mean 1.00  max 1.10

The aggregate_posterior function takes the test images, computes a code for each with the reparametrisation (as the decoder would see it in training) and returns the mean and variance per dimension. With β = 0 the mean variance of the codes is 8.75, up to 17.75, and the absolute mean 2.82: the decoder works on a scale very different from N(0, 1), so sampling from the prior lands in a zone it has never seen. With β = 1 the mean variance is 1.00 and the absolute mean 0.05: an aggregate posterior aligned with the prior. Here, though, there is a trap, visible in the third row: even with β = 4 the mean variance is 1.00 and the mean 0.05. The collapsed model passes the check perfectly, because its posterior is identical to the prior by construction. The aggregate posterior diagnostic is a necessary condition, not a sufficient one: it must be read together with the count of active dimensions and the diversity of the samples.

Training instabilities

VAEs have numerically delicate points: exp(logvar) can overflow and the gradients through σ can become large. We try five learning rates, with and without two countermeasures: limiting logvar to [−10, 10] and clipping the gradient norm at 5 (torch.nn.utils.clip_grad_norm_).

@torch.no_grad()
def evaluate_loss(model):
    logits, mu, logvar = model(Xte)
    return vae_loss(logits, Xte, mu, logvar)[0].item()

def train_with_lr(lr, stabilise=False, epochs=30, seed=0):
    """Returns the epoch at which the loss became non-finite (or None) and the final test loss."""
    torch.manual_seed(seed)
    model = VAE(latent=10)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    for epoch in range(1, epochs + 1):
        perm = torch.randperm(len(Xtr))
        for i in range(0, len(Xtr), 64):
            x = Xtr[perm[i:i + 64]]
            logits, mu, logvar = model(x)
            if stabilise:
                logvar = logvar.clamp(-10, 10)                     # keep exp(logvar) in a safe range
            loss, _, _ = vae_loss(logits, x, mu, logvar)
            opt.zero_grad(); loss.backward()
            if stabilise:
                nn.utils.clip_grad_norm_(model.parameters(), 5.0)  # tame huge gradients
            opt.step()
            if not torch.isfinite(loss):
                return epoch, float("nan")
    return None, evaluate_loss(model)

for lr in (1e-3, 1e-2, 5e-2, 2e-1, 5e-1):
    nan_epoch, final = train_with_lr(lr)
    stab_epoch, stab_final = train_with_lr(lr, stabilise=True)
    print(f"lr={lr:g}: plain -> non-finite at epoch {nan_epoch}, final {final:.1f}   |   with clamp+clipping -> non-finite at epoch {stab_epoch}, final {stab_final:.1f}")

lr=0.001: plain -> non-finite at epoch None, final 25.4   |   with clamp+clipping -> non-finite at epoch None, final 25.3
lr=0.01: plain -> non-finite at epoch None, final 24.9   |   with clamp+clipping -> non-finite at epoch None, final 24.7
lr=0.05: plain -> non-finite at epoch None, final 26.8   |   with clamp+clipping -> non-finite at epoch None, final 27.3
lr=0.2: plain -> non-finite at epoch None, final 14224.7   |   with clamp+clipping -> non-finite at epoch None, final 27.8
lr=0.5: plain -> non-finite at epoch 1, final nan   |   with clamp+clipping -> non-finite at epoch 1, final nan

The function returns the epoch at which the loss becomes non-finite (check with torch.isfinite), or None, plus the final test loss. Up to lr 0.05 no configuration gives NaN, and the countermeasures hardly matter (25.4 against 25.3 at lr 0.001; 26.8 against 27.3 at 0.05, where they are even slightly worse). At lr 0.2 the instructive thing happens: without countermeasures training does not produce NaN but the final loss is 14224.7; with clamp and clipping it is 27.8. The failure is silent, an enormous loss with no error at all. At lr 0.5 both versions produce non-finite values already in the first epoch: the countermeasures do not rescue an absurd learning rate, and I did not investigate why.

A note on the code, because it is easy to copy badly: in the block the clamp acts on the logvar passed to the loss, so it protects the KL term, but the sampling inside model(x) still uses the unlimited value. To protect the exp(0.5 · logvar) of the sampling too, the clamp must go inside reparametrize. On larger models (deep convolutional decoders, unnormalised data) NaNs become more likely: the two measures remain good practice, but this experiment does not prove it, it only shows the lr 0.2 case on a small model.

Other typical failures and where to look for them

• Blurry images. Not a bug: it comes from the likelihood over independent pixels, which favours averages of possible futures. For sharp images you need other decoders or other losses.
• Excellent reconstruction, poor generation. Symptom of a β that is too small: aggregate posterior far from the prior. Check the variance of the z.
• All samples identical, poor reconstruction. β too large: check the active dimensions and whether the judge always sees the same digit.
• Very few active dimensions. Collapse: reduce β, reduce the latent dimension, try the warm-up or make the decoder less powerful.
• ELBO getting worse with more data. Probably a sum-or-mean error in the two components of the loss.

The list is a table of symptoms and causes: each item links an observation to a check, which is why the order does not matter. The most useful line is the third, because the experiment showed that the symptom "all samples identical" and the symptom "perfect aggregate posterior" can appear together.

Trade-off

A small β gives better reconstruction at the expense of the regularity of the code, a large β the opposite up to collapse (table of experiment 1). A large latent dimension gives more room and more dimensions to switch off: it is often worth overestimating and letting the KL do its work, watching the active dimensions. The warm-up can help with powerful decoders like Bowman's; here it gave no distinguishable benefit and costs one more hyperparameter.

Real pitfalls

1. Comparing models with different β using their own loss. A loss with β = 4 is not comparable with one with β = 1; always evaluate with β = 1.
2. Judging a VAE by reconstruction alone. With β = 0 the reconstruction is the best (16.7 nats) and the generator the worst among the non-collapsed ones.
3. Not looking at the active dimensions. A VAE with 40 dimensions and 38 active is normal; one with 10 dimensions and 0 active is broken.
4. Trusting the aggregate posterior alone. The collapsed model has it perfect.
5. Thinking that the absence of NaN means healthy training. At lr 0.2 there is no NaN and the loss is 14224.7.
6. One seed, one configuration. All the numbers here are single runs: the trends (KL falls, reconstruction rises with β) are clear, the small differences (33 against 38 active dimensions) are not reliable.

The six pitfalls have a common cause: what the program prints does not coincide with what you want to measure. A single figure (the loss, the variance of the codes, the absence of NaN) looks like a diagnosis and is not. The checklist that follows is simple: record reconstruction and KL separately at every epoch, count the active dimensions, check the mean and variance of the aggregate posterior, sample a few thousand images from the prior and measure the diversity, compare with an autoencoder of the same size, and if you use β ≠ 1 report the loss at β = 1 anyway.

A worked example by hand: why β controls the "precision" of the code

A one-dimensional case in which the result can be derived with a pencil. The image is x = 0 and the decoder is p(x | z) = N(z, s²) with s = 0.3. The encoder produces a Gaussian with mean μ = 0 and standard deviation σ to be chosen. The loss is "reconstruction + β·KL"; with μ = 0 the expected reconstruction is σ²/(2s²) (the larger σ is, the more z departs from 0 and the more the decoder errs, up to constants) and the KL is ½(σ² − ln σ² − 1) (the smaller σ is, the more precise and expensive the code). Differentiating with respect to σ² and setting to zero:

1/(2s²) + (β/2)·(1 − 1/σ²) = 0     →     σ² = β / ( β + 1/s² )

With 1/s² = 11.11:

β = 0.1:  σ² = 0.0089  (σ = 0.094)   KL = 1.86 nat
β = 1:    σ² = 0.0826  (σ = 0.287)   KL = 0.79 nat
β = 10:   σ² = 0.474   (σ = 0.688)   KL = 0.11 nat

The result has the same shape as the measured effect: a high β makes the code less precise (larger σ) and lowers its information cost (smaller KL); a low β makes it precise and expensive. As β → 0, σ² → 0: the code becomes an exact point, as in an autoencoder. In the real VAE the same logic acts on every component, with the difference that the precision needed depends on the image. These calculations are done by hand, not with the program, and rechecked with a separate calculation.

Frequently asked questions

How do I choose β in practice?

Start from 1, which is the true ELBO. If you need more faithful reconstructions try β < 1 (here 0.25 gave a reconstruction of 18.2 nats with samples still diversified); if you want a more regular or compact space, β > 1, but check the active dimensions: at β = 2 only 4 out of 10 remained. Always evaluate with the loss at β = 1.

Is it true that a high β "disentangles" the factors?

It is the idea of the β-VAE authors; the results depend on the data and we did not verify it. What we measured is the reduction of active dimensions, not interpretability.

Can I change β during training?

Yes: that is the KL warm-up, you start with β = 0 and rise to 1 in the first epochs. With powerful decoders it serves to avoid the collapse described by Bowman and colleagues; in our experiment it gave no clear benefit.

If collapse is benign, why worry about it?

Because inactive dimensions can signal that the model is not using the code as it should, for example with a decoder that is too powerful. Counting them costs one line.

Does an aggregate posterior close to the prior guarantee good samples?

No: it is necessary, not sufficient. The β = 4 model has mean variance 1.00 and all samples identical.

Exercises with solutions

Before moving on, try to solve them with pen and paper; the solutions are right below each one.

1. In the toy case with s = 0.5 and β = 2, what is the optimal encoder variance σ² = β/(β + 1/s²)?
   Solution: 1/s² = 4, so σ² = 2/(2 + 4) = 0.333.
2. What happens to σ² when β → 0 and when β → ∞?
   Solution: β → 0: σ² → 0 (exact code, like an autoencoder); β → ∞: σ² → 1 (the code becomes the prior: no information).
3. A dimension has a variance of μ equal to 0.004 and mean KL 0.002. Is it active?
   Solution: No (variance < 0.01 and KL almost null): it has collapsed.
4. From the table of experiment 1, which of the β = 0.25 and β = 1 models has the better −ELBO on the test set, and why?
   Solution: β = 1: 21.9 + 2.9 = 24.8 nats against 18.2 + 9.7 = 27.9. The β = 1 model optimises exactly that quantity; the β = 0.25 one optimises another.

The first two exercises use the formula derived by hand, the third applies the diagnosis rule, the fourth forces you to read the table with the right loss. If you get the fourth wrong, you compared the reconstruction and not the −ELBO.

What you can do now

You can use β as a knob and recognise where it stops working: at β = 0 the KL is 137 nats and the codes have a mean variance of 8.75, off scale with respect to the prior; at β = 4 the code is empty and the decoder emits the mean image. You know that with a code that is too large the VAE switches off the useless dimensions by itself (38 active out of 40), and that checking the aggregate posterior alone does not uncover a collapsed model. You know that instability shows up as an enormous loss (14224.7 at lr 0.2) more than as a NaN. We close the module on generative models here; in the next one we change paradigm with reinforcement learning, the agent that learns by trial and error.