The SGVB Estimator and AEVB Algorithm

This section covers Section 2.3 of the paper: the practical estimator that makes VAE training possible.

The Goal

We need a practical way to:

1. Estimate the ELBO $\mathcal{L}(\theta, \phi; \mathbf{x}^{(i)})$
2. Compute gradients $\nabla_{\theta, \phi} \mathcal{L}$
3. Use standard stochastic gradient methods (SGD, Adam, etc.)

The Reparameterization Insight

Under certain conditions (detailed in the next section), we can reparameterize the random variable $\tilde{\mathbf{z}} \sim q_\phi(\mathbf{z}|\mathbf{x})$ using a differentiable transformation:

Equation (4):

$$\tilde{\mathbf{z}} = g_\phi(\boldsymbol{\epsilon}, \mathbf{x}) \quad \text{with} \quad \boldsymbol{\epsilon} \sim p(\boldsymbol{\epsilon})$$

Where:

• $\boldsymbol{\epsilon}$ is an auxiliary noise variable with a fixed distribution $p(\boldsymbol{\epsilon})$
• $g_\phi(\boldsymbol{\epsilon}, \mathbf{x})$ is a differentiable function parameterized by $\phi$

This allows us to rewrite expectations:

Equation (5):

$$\mathbb{E}{q\phi(\mathbf{z}|\mathbf{x}^{(i)})}[f(\mathbf{z})] = \mathbb{E}{p(\boldsymbol{\epsilon})}\left[f(g\phi(\boldsymbol{\epsilon}, \mathbf{x}^{(i)}))\right] \simeq \frac{1}{L} \sum_{l=1}^{L} f(g_\phi(\boldsymbol{\epsilon}^{(l)}, \mathbf{x}^{(i)}))$$

Where $\boldsymbol{\epsilon}^{(l)} \sim p(\boldsymbol{\epsilon})$.

Why this matters: The expectation is now over $p(\boldsymbol{\epsilon})$, which doesn’t depend on $\phi$. So we can differentiate through the samples!

SGVB Estimator A — Equation (6)

Applying this to the ELBO (Equation 2 form), we get the generic SGVB estimator:

$$\tilde{\mathcal{L}}^A(\theta, \phi; \mathbf{x}^{(i)}) = \frac{1}{L} \sum_{l=1}^{L} \log p_\theta(\mathbf{x}^{(i)}, \mathbf{z}^{(i,l)}) - \log q_\phi(\mathbf{z}^{(i,l)}|\mathbf{x}^{(i)})$$

Where: $$\mathbf{z}^{(i,l)} = g_\phi(\boldsymbol{\epsilon}^{(i,l)}, \mathbf{x}^{(i)}) \quad \text{and} \quad \boldsymbol{\epsilon}^{(l)} \sim p(\boldsymbol{\epsilon})$$

SGVB Estimator B — Equation (7) (Lower Variance)

When the KL divergence $D_{KL}(q_\phi(\mathbf{z}|\mathbf{x}^{(i)}) | p_\theta(\mathbf{z}))$ can be computed analytically (e.g., both are Gaussian), we get a better estimator:

$$\tilde{\mathcal{L}}^B(\theta, \phi; \mathbf{x}^{(i)}) = -D_{KL}(q_\phi(\mathbf{z}|\mathbf{x}^{(i)}) | p_\theta(\mathbf{z})) + \frac{1}{L} \sum_{l=1}^{L} \log p_\theta(\mathbf{x}^{(i)}|\mathbf{z}^{(i,l)})$$

Where: $$\mathbf{z}^{(i,l)} = g_\phi(\boldsymbol{\epsilon}^{(i,l)}, \mathbf{x}^{(i)}) \quad \text{and} \quad \boldsymbol{\epsilon}^{(l)} \sim p(\boldsymbol{\epsilon})$$

Why lower variance? The KL term is computed exactly (no sampling), only the reconstruction term needs Monte Carlo estimation.

Minibatch Estimator — Equation (8)

For a dataset with $N$ datapoints, we estimate the full ELBO using minibatches:

$$\mathcal{L}(\theta, \phi; \mathbf{X}) \simeq \tilde{\mathcal{L}}^M(\theta, \phi; \mathbf{X}^M) = \frac{N}{M} \sum_{i=1}^{M} \tilde{\mathcal{L}}(\theta, \phi; \mathbf{x}^{(i)})$$

Where $\mathbf{X}^M = {\mathbf{x}^{(i)}}_{i=1}^{M}$ is a randomly drawn minibatch of size $M$.

Key finding from the paper: With a large enough minibatch (e.g., $M = 100$), we can set $L = 1$ (one sample per datapoint).

The AEVB Algorithm

Algorithm 1 from the paper:

Input: θ, φ (initial parameters)

repeat:
    X^M ← Random minibatch of M datapoints from full dataset
    ε   ← Random samples from noise distribution p(ε)
    g   ← ∇_{θ,φ} L̃^M(θ, φ; X^M, ε)  # Gradients of minibatch estimator
    θ, φ ← Update parameters using gradients g (e.g., SGD or Adagrad)
until convergence of parameters (θ, φ)

return θ, φ

Settings used in experiments: $M = 100$, $L = 1$

Connection to Auto-Encoders

Looking at Estimator B (Equation 7):

$$\tilde{\mathcal{L}}^B = \underbrace{-D_{KL}(q_\phi(\mathbf{z}|\mathbf{x}) | p_\theta(\mathbf{z}))}{\text{Regularizer}} + \underbrace{\frac{1}{L} \sum{l=1}^{L} \log p_\theta(\mathbf{x}|\mathbf{z}^{(l)})}_{\text{Expected negative reconstruction error}}$$

This reveals the auto-encoder connection:

The function $g_\phi(\cdot)$ maps a datapoint $\mathbf{x}^{(i)}$ and noise $\boldsymbol{\epsilon}^{(l)}$ to a sample from the approximate posterior:

$$\mathbf{z}^{(i,l)} = g_\phi(\boldsymbol{\epsilon}^{(l)}, \mathbf{x}^{(i)}) \quad \text{where} \quad \mathbf{z}^{(i,l)} \sim q_\phi(\mathbf{z}|\mathbf{x}^{(i)})$$

Then $\log p_\theta(\mathbf{x}^{(i)}|\mathbf{z}^{(i,l)})$ is the log-probability of $\mathbf{x}^{(i)}$ under the decoder — this is the reconstruction term.

Why AEVB is Efficient

Traditional variational inference problems:

• Need to run optimization per datapoint
• Expensive iterative inference (e.g., MCMC) for each $\mathbf{x}$

AEVB solution:

• Amortized inference: The encoder $q_\phi(\mathbf{z}|\mathbf{x})$ is a neural network that directly maps any $\mathbf{x}$ to its approximate posterior
• No per-datapoint optimization: Just a forward pass through the encoder
• Stochastic gradient descent: Works with minibatches, scales to large datasets

Computational Complexity

Per minibatch update:

1. Forward pass through encoder: $O(M \cdot \text{encoder complexity})$
2. Sample $\boldsymbol{\epsilon}$: $O(M \cdot \text{latent dim})$
3. Compute $\mathbf{z} = g_\phi(\boldsymbol{\epsilon}, \mathbf{x})$: $O(M \cdot \text{latent dim})$
4. Forward pass through decoder: $O(M \cdot \text{decoder complexity})$
5. Compute loss and gradients: $O(M \cdot \text{output dim})$
6. Backward pass: Same complexity as forward

Total: Linear in minibatch size, linear in network size — same as training a standard neural network!

Summary