Optimizing the VAE Objective in Practice

From a VAE Objective to an Optimization Problem

A count VAE is trained by adjusting encoder and decoder parameters so that its latent-variable model assigns useful probability to observed count vectors while keeping the approximate posterior compatible with the prior. The evidence lower bound (ELBO) gives a per-cell objective for this trade-off. In practice, the data set may contain many cells, so training uses repeated minibatches to estimate how parameter changes affect the average objective.

Let \(x_i\) be the count vector for cell \(i\), let \(q_\phi(z\mid x_i)\) be the encoder distribution, and let \(p_\theta(x_i\mid z)\) and \(p(z)\) be the decoder likelihood and latent prior. The parameters are \(\theta\) and \(\phi\), collectively denoted \(\psi=(\theta,\phi)\). For the standard objective, the per-cell ELBO is:

The first term is the expected log likelihood, often called the reconstruction term. The second is the Kullback–Leibler divergence, which is nonnegative and penalizes a variational distribution that differs from the prior. With \(N\) cells, one common empirical target is the average ELBO, \(\mathcal{J}(\psi)=N^{-1}\sum_{i=1}^N\mathcal{L}_i(\psi)\). Maximizing this target is equivalent to minimizing its negative, the average negative-ELBO loss. Software may report the negative log likelihood and KL term as positive quantities, so always check which sign convention is being plotted.

The optimization target is an empirical average over the observed cells. The cells are treated as fixed for this target; randomness enters through minibatch selection and, when used, Monte Carlo draws from the encoder. The standard minibatch argument assumes uniform cell sampling, a correctly weighted minibatch average, and gradients with finite expectations. Neural-network objectives are generally non-concave, so these steps do not guarantee a global maximum.

Why the Minibatch Gradient Works

Suppose a minibatch \(B\) contains \(m\) cells sampled uniformly from the \(N\) observed cells. An estimator of the full-data average objective is the average ELBO within the minibatch. The same expression applies to its gradient, provided differentiation and expectation can be interchanged and the gradient estimates are well behaved.

This is an unbiased gradient estimate for the average objective. If optimizing the sum rather than the average, the corresponding estimator is \(N/m\) times the minibatch sum. Mixing these conventions changes gradient scale by a factor of \(N\), which can make a learning rate inappropriate even though the underlying objective has not changed.

The encoder expectation is usually estimated with the reparameterization trick for Gaussian latents: sample a noise variable and express \(z\) as a differentiable function of that noise and encoder parameters. Conditional on the minibatch, a Monte Carlo gradient based on this construction estimates the gradient of the minibatch expected objective. A single latent draw adds sampling noise; averaging more draws can reduce that noise but costs computation. A learning rate is the scale of a parameter update, not a correction for a biased or incorrectly scaled objective.

KL Warm-Up Changes the Target During Training

The KL term can be given a time-varying weight \(\beta_t\), producing the objective \(\mathcal{L}_{i,\beta_t}= \mathbb{E}_{q_\phi}[\log p_\theta(x_i\mid z)]-\beta_t\mathrm{KL}(q_\phi\|p)\). KL warm-up starts with a smaller weight and increases it toward one according to a schedule. The intent is to let the decoder and encoder begin fitting the observations before applying the full prior-matching penalty. It changes the objective being optimized; it is not merely a change to the optimizer’s step size.

When \(\beta_t=1\), this is the usual ELBO. At other values it is generally a weighted objective, not the standard ELBO lower bound. Therefore, a curve plotted during warm-up can rise partly because the objective’s KL weight is changing. For a fair comparison across training epochs or models, evaluate the standard ELBO with \(\beta=1\) separately, even if the training objective uses a schedule.

Learning Rates and Training-Curve Reading

For gradient ascent with a fixed learning rate \(\eta>0\), a basic update is \(\psi_{t+1}=\psi_t+\eta\widehat{g}_t\), where \(\widehat{g}_t\) is a minibatch estimate of the ELBO gradient. For gradient descent on the negative ELBO, the update has the corresponding negative-gradient direction. Adaptive methods such as Adam rescale updates using running estimates of gradient means and squared gradients, but still require a learning-rate choice. The scale of the ELBO, batch size, parameterization, and optimizer all affect useful numerical values; there is no universally optimal rate.

A rate that is too large may cause objective values to oscillate, jump sharply, or become non-finite. A rate that is too small may produce slow, nearly flat progress. A practical approach is to begin with a conservative rate, examine the scale and stability of updates, and adjust using both training and validation curves. A schedule that reduces the rate after progress plateaus can make later updates less erratic, but should not be used to conceal a persistently unstable objective or an incorrect implementation.

Track the reconstruction contribution and KL contribution separately as well as their combined objective. With the conventional ELBO signs, both log likelihood and ELBO are often negative, while KL is nonnegative. If software instead displays negative log likelihood, lower values are better for that component. Record the sign convention and whether values are sums or per-cell averages. Evaluate validation cells without parameter updates and with the same objective definition; when warm-up is used, a standard-\(\beta\) validation ELBO helps distinguish actual model fit from the changing training target.

Diagnostics, Alternatives, and Limitations

Optimization curves are diagnostics, not proof that the model is correctly specified or that a global optimum has been found. Several checks help distinguish expected stochastic variation from problems:

• Large oscillations or abrupt deterioration: lower the learning rate and check numerical stability, loss scaling, and gradient magnitudes.
• Flat training and validation curves: determine whether the rate is too small, gradients are very noisy, or the objective has reached a plateau. A flat curve by itself does not identify the cause.
• Training improves while validation worsens: this may indicate overfitting or a difference between training and validation processing. Compare the same fixed objective on both sets.
• Unexpected signs or scales: verify whether the software reports log likelihood or negative log likelihood, sums or means, and whether the KL weight is included.
• A changing KL weight: report the schedule and do not interpret the warm-up objective as a single fixed metric.

Uniform minibatch stochastic gradient ascent is the basic approach because it gives a simple unbiased estimate of the empirical gradient. Full-batch optimization uses all cells at once and removes cell-sampling variability, but each update is more costly. Deterministic optimization of a Monte Carlo approximation can still have latent-sampling noise; using more latent draws may reduce it. Adaptive optimizers can be more robust to differently scaled parameter gradients, while a plain fixed-rate update is simpler to interpret. None removes the need to define the objective consistently, monitor held-out performance, and check numerical behavior.

The minibatch estimator’s unbiasedness relies on the sampling scheme and weighting. If cells have unequal sampling probabilities, an unweighted average need not estimate the full-data average gradient; appropriate weights are needed. Also, minibatch gradients are noisy, and a favorable minibatch step does not guarantee that the full-data objective rises at every update. Because the neural-network objective is non-convex, results can depend on initialization, data order, and random draws. Repeat runs and report the optimization choices when differences between fitted models matter.

A frequent mistake is to say that a higher warm-up objective proves better standard-ELBO fit. That reasoning compares values computed with different KL weights. The correction is to calculate or evaluate every checkpoint using the same \(\beta=1\) objective, and to treat the scheduled objective as the optimization device rather than the final comparison metric.