Articles

VAE Coding Interview: Implement the Loss and Debug a Collapsing Latent Space

Implement VAE loss with clear reductions, check KL and reparameterization, and diagnose unused latents with numerical fixtures and controlled interventions.

A VAE coding interview can start with a short loss function and turn into a debugging exercise. Reconstruction improves, KL approaches zero, and the interviewer asks whether the latent space has collapsed. That curve alone cannot tell you whether the model found an uninformative posterior, the loss was scaled incorrectly, or the reconstruction gradient stopped reaching the encoder.

Make the likelihood, tensor shapes, reduction units, and gradient path explicit before changing the architecture or KL weight. Then test whether the decoder actually uses input-dependent latent information. A plausible remedy needs a mechanism and a comparable evaluation, not just a higher KL number.

This guide uses original research papers, official PyTorch documentation, and original numerical fixtures. It does not report an employer's question frequency.

Define the model before writing its loss

For our original exercise, each input is a vector of D binary observations. The decoder returns D Bernoulli logits. The encoder returns a K-dimensional mean and log variance for a diagonal Gaussian approximate posterior, and the prior is a standard normal distribution. Batch size is B; all examples use the same observed dimensions without masking.

The original VAE paper develops a variational objective with an expected log-likelihood term and a divergence between approximate posterior and prior. Maximizing that ELBO corresponds to minimizing reconstruction negative log-likelihood plus KL(q(z|x) || p(z)); the direction of the KL is part of the definition.

Name shapes before implementation: input and logits are [B, D]; mean and log variance are [B, K]; a sampled latent is [B, K]. Sum likelihood contributions over observation dimensions and KL contributions over latent dimensions, then average examples. Reconstruction usually uses a Monte Carlo sample; this Gaussian KL can be evaluated analytically.

Do not silently apply the binary exercise to arbitrary real-valued features. A Gaussian observation model has different likelihood terms and variance assumptions. Likewise, treating grayscale values as Bernoulli targets is a modeling choice that deserves explanation. Our fixture uses actual zeros and ones to keep the contract unambiguous.

Implement the loss with visible reduction units

This original NumPy reference ran locally. It checks the loss arithmetic independently of an encoder or decoder training loop.

import numpy as np

def vae_terms(logits, x, mu, logvar):
    logits, x, mu, logvar = [np.asarray(a, dtype=float)
                            for a in (logits, x, mu, logvar)]
    assert logits.ndim == x.ndim == mu.ndim == logvar.ndim == 2
    assert logits.shape == x.shape and mu.shape == logvar.shape
    assert x.shape[0] == mu.shape[0] and x.shape[0] > 0
    assert np.isin(x, [0, 1]).all()
    bce = (np.maximum(logits, 0) - logits*x
           + np.log1p(np.exp(-np.abs(logits))))
    reconstruction = bce.sum(axis=1)
    kl = 0.5 * (mu**2 + np.exp(logvar) - 1 - logvar).sum(axis=1)
    return (reconstruction + kl).mean(), reconstruction.mean(), kl.mean()

The stable binary-logit expression avoids explicitly taking the logarithm of a probability rounded to zero. It does not solve overflow from an extreme log variance: that exponential still needs monitoring. The assertions express our exercise contract; they are not a complete production validation layer.

The official BCEWithLogitsLoss documentation combines sigmoid and binary cross-entropy for numerical stability and defines reduction behavior. Its default mean averages output elements. That default is different from summing each example's D contributions and then averaging B examples.

Log the two components separately, along with their reduction conventions. A single scalar can hide one term overwhelming the other, and a graph labeled simply “KL” may already have a coefficient or batch normalization applied to it.

Check a two-example fixture before any training run

Our original fixture has B=2, D=2, and K=2. Both examples have zero logits and binary targets [0,1]. The first posterior has mean [0,0] and log variance [0,0]. The second has mean [1,0] and log variance [0,log(4)].

Each reconstruction term is 2*log(2), approximately 1.3863. The first KL is zero because its posterior equals the standard-normal prior. The second KL is 0.5*(1 + 4 - 1 - log(4)), approximately 1.3069. The batch-mean KL is approximately 0.6534, so the total batch-mean loss is approximately 2.0397.

x = np.array([[0., 1.], [0., 1.]])
logits = np.zeros_like(x)
mu = np.array([[0., 0.], [1., 0.]])
logvar = np.array([[0., 0.], [0., np.log(4)]])
loss, reconstruction, kl = vae_terms(logits, x, mu, logvar)
assert np.isclose(reconstruction, 2*np.log(2))
assert np.isclose(kl, (4 - np.log(4))/4)
assert np.isclose(loss, reconstruction + kl)

These checks ran locally. We also duplicated the complete batch and confirmed the mean loss was unchanged. That check detects accidental batch summation; it does not establish correct gradients or validate a trained model.

For a sign check, move a prior-matching posterior mean away from zero while holding logits fixed. Its KL should increase. For a scale check, distinguish variance four from standard deviation two. Neither a negative KL in this ordinary fixture nor an unexplained dependence on repeated batch size should be dismissed as a training phenomenon.

Why can a harmless-looking mean change the objective?

With D observation dimensions, averaging reconstruction over all B times D elements makes it one D-th of our per-example reconstruction. If KL remains summed over K and averaged over B, the relative KL weight has increased by a factor of D compared with the agreed objective.

In our two-feature fixture, that mistaken reconstruction is approximately 0.6931, while mean KL remains 0.6534. The resulting total is approximately 1.3466, not 2.0397. With hundreds of observation dimensions, the imbalance can be much larger.

This is not equivalent to scaling the entire loss by one constant: only one component changed. A learning-rate adjustment cannot restore the original relative weighting. Similarly, dividing KL by K while leaving reconstruction unchanged introduces another objective choice.

If you intentionally use reconstruction + beta*KL, record beta and both units. Beta one under our conventions gives the standard negative-ELBO estimator. Another beta modifies the objective; avoid labeling every weighted sum as the same ELBO. Compare runs using consistent reporting before interpreting their curves.

Keep the reparameterization path differentiable

For diagonal Gaussian parameters, standard deviation is exp(0.5*logvar). Draw parameter-independent standard-normal noise and construct z = mu + std*epsilon. Using exp(logvar) as the standard deviation would square the intended variance relationship.

Official PyTorch distributions documentation distinguishes ordinary sampling from the differentiable pathwise construction available through rsample. It also defines Normal's scale as standard deviation. Here are original PyTorch functions reviewed against those APIs; they were not executed in this environment.

import torch
import torch.nn.functional as F

def reparameterize(mu, logvar):
    std = torch.exp(0.5 * logvar)
    return mu + std * torch.randn_like(std)

def vae_loss(logits, x, mu, logvar, beta=1.0):
    assert logits.ndim == x.ndim == mu.ndim == logvar.ndim == 2
    assert logits.shape == x.shape and mu.shape == logvar.shape
    assert x.shape[0] == mu.shape[0] and x.shape[0] > 0
    rec = F.binary_cross_entropy_with_logits(
        logits, x, reduction="none").sum(dim=1)
    kl = 0.5 * (mu.square() + logvar.exp() - 1 - logvar).sum(dim=1)
    return (rec + beta*kl).mean(), rec.mean(), kl.mean()

These functions assume finite floating-point tensors on compatible devices and the binary-target contract above. The caller must connect encoder outputs through reparameterize, then pass z into the decoder producing logits. A correct loss helper alone does not demonstrate that connection.

The official Tensor.detach documentation says the detached tensor leaves the current gradient graph. Detaching z before decoding breaks the reconstruction path to the encoder. Converting it through an unrelated array or scalar can do the same.

Encoder gradients may still be nonzero because the KL term depends on encoder outputs. Checking only the total encoder gradient can therefore miss a broken reconstruction path. Inspect the two terms' contributions separately with a deliberately nondegenerate fixture.

Use a controlled derivative check rather than a random curve

For our original scalar diagnostic, fix the noise at zero, set decoder logit equal to the latent, and use target zero. At mu=1, reconstruction is log(1+exp(mu)), whose derivative is approximately 0.7311. The unit-variance posterior KL derivative with respect to mu is 1.

Central finite differences in our NumPy check reproduce both values. If the decoder's logit is held fixed while mu changes, reconstruction sensitivity becomes zero while KL sensitivity stays one. This is a controlled illustration of a disconnected reconstruction dependency, not an executed PyTorch autograd experiment.

Hold the noise fixed for such a comparison. Drawing new noise on each side of a tiny perturbation can overwhelm the derivative estimate. Also avoid a test point where the expected gradient is zero; at mu=0 with unit variance, KL legitimately has a stationary point.

Once using a training framework, inspect whether the optimizer contains encoder parameters, whether the reconstruction graph reaches them, and whether gradients are finite. Check before and after the intended optimizer step. These are our diagnostic recommendations, not evidence that a particular training run failed for those reasons.

Decide whether the latent is being ignored

Research on posterior collapse describes approximate posteriors approaching the prior and latent information being unused. Lucas and colleagues' linear-VAE analysis also shows why collapse should not be attributed universally to one simplistic explanation about the ELBO. Their analyzed setting does not prove one cause for every nonlinear architecture.

Our recommended investigation follows this order:

Observation	Next discriminating check
KL unexpectedly small	Verify unweighted per-example and per-dimension KL
Encoder receives some gradients	Isolate reconstruction from KL contributions
Reconstructions look plausible	Swap latents between different inputs
Swapping latents changes little	Compare logits and held-out reconstruction losses
Some dimensions are inactive	Inspect active dimensions before declaring total collapse

Keep the decoder's other inputs fixed when changing z. An autoregressive decoder receiving true previous tokens has information routes that a simple binary-image decoder lacks. Shuffling z tests dependence, but it can also move combinations outside the training distribution; interpret the intervention alongside other evidence.

Inspect whether posterior means vary across diverse inputs and whether variances approach the prior. A small average KL can hide a few active dimensions or examples. A larger KL can reflect a poorly matched prior without guaranteeing useful representations. Combine these measurements with a stated task and held-out evaluation.

Choose a remedy that addresses the observed mechanism

Fix a sign, reduction, missing optimizer parameter, or detached path before changing model capacity. If the implementation checks pass, test a training intervention while holding the dataset split, likelihood, reporting units, and evaluation contract fixed.

Bowman and colleagues study sentence VAEs with autoregressive decoders and discuss KL-cost annealing and weakening decoder conditioning. Those are findings in that setting, not a guarantee that any warm-up schedule repairs every VAE.

For a proposed KL warm-up, specify the coefficient schedule and evaluate after it reaches the intended objective. For a weaker decoder or restricted side information, compare the cost in reconstruction quality with improved latent use. Do not select a model solely because its KL increased.

During the interview, propose one controlled change and the evidence that would support it. “I will check reconstruction gradients before tuning beta” is a concrete debugging step. “I will add every anti-collapse technique” leaves the original cause unresolved.

Practise the implementation and its follow-ups

Question	Useful follow-up
Explain VAEs, ELBO, KL, and Reparameterization	State shapes, KL direction and sample path
Implement and Debug Backprop in NumPy	Check derivatives with controlled inputs
Implement Softmax Cross-Entropy Forward and Backward Passes and a Training Loop	Explain numerical stability and reduction units
Implement PyTorch training loop	Verify parameter, gradient and step wiring
Compare PCA and Autoencoders Under the Curse of Dimensionality	Separate representation utility from reconstruction

For deterministic bottlenecks and anomaly thresholds, use the separate autoencoder guide. For this coding exercise, return to the VAE question and explain the two-example loss before discussing collapse.

Sources and Further Reading

Research checked October 4, 2026. Original NumPy arithmetic and finite-difference checks ran locally. PyTorch snippets are documentation-reviewed, not executed training or autograd results; no candidate-process claims are made.