Importance Weighted Autoencoders (IWAEs)

As we've discussed, the standard Evidence Lower Bound (ELBO) is central to training VAEs. However, the ELBO is, by definition, a lower bound on the log marginal likelihood $\log p_\theta(x)$. The gap between the ELBO and the true log-likelihood is precisely the KL divergence $KL(q_\phi(z|x) || p_\theta(z|x))$. A simpler approximate posterior $q_\phi(z|x)$ often leads to a larger gap, meaning the bound is looser. While making $q_\phi(z|x)$ more expressive (e.g., using normalizing flows, which we'll cover) is one way to tighten this bound, Importance Weighted Autoencoders (IWAEs) offer an alternative: achieve a tighter bound by using multiple samples from the same $q_\phi(z|x)$.

The IWAE Objective: A Tighter Bound with Multiple Samples

The core idea behind IWAEs, introduced by Burda, Grosse, and Salakhutdinov, is to leverage multiple samples from the approximate posterior $q_\phi(z|x)$ to form a better Monte Carlo estimate of $p_\theta(x)$. Recall that the marginal likelihood can be written as $p_\theta(x) = \mathbb{E}_{z \sim q_\phi(z|x)} \left[ \frac{p_\theta(x,z)}{q_\phi(z|x)} \right]$. The standard ELBO is $\mathbb{E}_{z \sim q_\phi(z|x)} \left[ \log \frac{p_\theta(x,z)}{q_\phi(z|x)} \right]$. Jensen's inequality tells us that $\log \mathbb{E}[Y] \ge \mathbb{E}[\log Y]$, which is why the ELBO is a lower bound.

The IWAE objective, denoted $L_K(x)$, uses $K$ samples $z_1, \dots, z_K$ drawn independently from $q_\phi(z|x)$ for a given input $x$:

$$L_K(x) = \mathbb{E}_{z_1, \dots, z_K \sim q_\phi(z|x)} \left[ \log \left( \frac{1}{K} \sum_{k=1}^K \frac{p_\theta(x, z_k)}{q_\phi(z_k|x)} \right) \right]$$

Each term $w_k = \frac{p_\theta(x, z_k)}{q_\phi(z_k|x)}$ is an importance weight. The IWAE objective averages these weights before taking the logarithm. This is a subtle but significant difference from the standard ELBO.

The $L_K$ objective has several desirable properties:

1. Monotonically Improving Bound: For any $K \ge 1$, we have: $$\log p_\theta(x) \ge L_K(x) \ge L_{K-1}(x)$$ This means that as you increase the number of samples $K$, the IWAE objective provides a progressively tighter lower bound on the true log marginal likelihood.
2. Connection to ELBO: For $K=1$, the IWAE objective $L_1(x)$ is exactly the standard ELBO: $$L_1(x) = \mathbb{E}_{z_1 \sim q_\phi(z|x)} \left[ \log \left( \frac{p_\theta(x, z_1)}{q_\phi(z_1|x)} \right) \right]$$
3. Convergence to True Log-Likelihood: As $K \to \infty$, the IWAE objective $L_K(x)$ converges to the true log marginal likelihood $\log p_\theta(x)$. This is because $\frac{1}{K} \sum_{k=1}^K w_k$ becomes a better estimate of $p_\theta(x)$.

The diagram below illustrates the general process of calculating the IWAE objective using $K$ samples.

Multiple samples $z_1, \dots, z_K$ are drawn from the inference network $q_\phi(z|x)$. Each sample contributes an importance weight $w_k = p_\theta(x|z_k)p(z_k)/q_\phi(z_k|x)$. These weights are averaged, and the logarithm of this average forms the IWAE objective $L_K$.

Intuition: Why Averaging Weights First Matters

The standard ELBO ($L_1$) can be thought of as averaging $\log(w_k)$. The IWAE objective ($L_K$) averages $w_k$ first, then takes the logarithm: $\log(\text{avg}(w_k))$. Due to Jensen's inequality ($\log \mathbb{E}[Y] \ge \mathbb{E}[\log Y]$), this change directly leads to a tighter bound.

Intuitively, $q_\phi(z|x)$ might be a poor approximation to the true posterior $p_\theta(z|x)$. It might assign low probability density to regions where $p_\theta(z|x)$ is high. If we only take one sample $z_1$ (as in the standard ELBO computation), and it lands in such a poorly estimated region, we get a bad estimate for that data point. With IWAE, if even one of the $K$ samples $z_k$ happens to fall into a region of high true posterior probability (and thus gets a large importance weight $w_k$), it can significantly lift the average $\frac{1}{K}\sum w_j$. This makes the overall estimate less susceptible to any single "unlucky" sample from $q_\phi(z|x)$. The inference network $q_\phi(z|x)$ is still "simple", but the multi-sample estimation of the bound is more robust.

Benefits of Using IWAE

Training a VAE by maximizing $L_K$ (for $K>1$) instead of the standard ELBO ($L_1$) offers several advantages:

1. Tighter Log-Likelihood Bounds: This is the most direct benefit. IWAEs typically report better (i.e., higher) log-likelihood estimates on test data compared to VAEs trained with the ELBO, especially when using the same $K$ for evaluation.
2. Improved Model Learning: Optimizing a tighter bound can lead to better parameter updates for both the generative model $p_\theta(x|z)$ (decoder) and the inference network $q_\phi(z|x)$ (encoder). It has been observed that models trained with IWAE can learn richer, more informative latent representations and generate higher quality samples. The encoder $q_\phi(z|x)$ might learn to propose samples $z_k$ that are more diverse or better cover the modes of the true posterior, as these will contribute more effectively to a high $L_K$ value.
3. No Change to Model Architecture: IWAE achieves these benefits without requiring a more complex inference network architecture (like structured VI or normalizing flows would). The expressiveness of $q_\phi(z|x)$ remains the same; it's the objective function that changes.

Computational Cost and Considerations

The primary drawback of IWAE is its increased computational cost during training and evaluation. For each data point and each gradient step:

• You need to sample $K$ latent variables $z_k$ from $q_\phi(z|x)$.
• You need to perform $K$ forward passes through the decoder to compute $p_\theta(x|z_k)$ for each $z_k$.
• You need to compute $q_\phi(z_k|x)$ for each $z_k$.

This means the computational cost scales linearly with $K$. If $K=50$, a training epoch will take roughly 50 times longer than for a standard VAE, assuming the decoder is the bottleneck.

Choice of $K$: The number of samples $K$ is a hyperparameter.

• Small $K$ (e.g., 5 to 10): Offers a modest improvement in the bound with a manageable increase in computation.
• Large $K$ (e.g., 50, 100, or even 5000 for evaluation): Provides a much tighter bound, approaching the true log-likelihood, but at a significant computational cost. In practice, training is often done with a moderate $K$ (e.g., $K=5$ or $K=50$), and evaluation might use a larger $K$ for a more accurate log-likelihood estimate.

Gradient Variance: While $L_K$ is a tighter bound, the variance of its gradient estimator can sometimes be an issue, particularly for very large $K$ or when importance weights are highly skewed. However, the benefits of a tighter bound often outweigh this concern, and techniques like using the reparameterization trick for $z_k$ are essential for keeping gradient variance manageable.

Training with the IWAE Objective

Training a VAE with the IWAE objective is similar to training a standard VAE, with the main difference being the loss calculation:

1. Forward Pass:
   • For an input $x$, obtain the parameters of $q_\phi(z|x)$ (e.g., $\mu(x), \sigma(x)$ if $q$ is Gaussian).
   • Draw $K$ samples $z_k \sim q_\phi(z|x)$ using the reparameterization trick. For each $z_k$:
     • Compute $\log q_\phi(z_k|x)$.
     • Pass $z_k$ through the decoder to get $p_\theta(x|z_k)$.
     • Compute $\log p(z_k)$ (from the prior).
     • Calculate the unnormalized log-weight: $\log \tilde{w}_k = \log p_\theta(x|z_k) + \log p(z_k) - \log q_\phi(z_k|x)$.
2. Objective Calculation:
   • Compute the IWAE objective $L_K$. To avoid numerical issues with very small weights, it's common to use the log-sum-exp trick:
   $$L_K(x) = \text{LogSumExp}(\log \tilde{w}_1, \dots, \log \tilde{w}_K) - \log K$$ where $\text{LogSumExp}(a_1, \dots, a_K) = \log \sum_{j=1}^K \exp(a_j)$.
3. Backward Pass: Compute gradients of $-L_K(x)$ with respect to $\theta$ and $\phi$, and update the parameters.

The rest of the VAE machinery, including network architectures for the encoder and decoder, remains largely the same.

Summary and When to Use IWAEs

Importance Weighted Autoencoders provide a principled way to obtain tighter lower bounds on the log-likelihood for VAEs by using multiple importance samples. This often translates into improved model performance, such as better quality generated samples and more meaningful latent representations, without explicitly increasing the complexity of the inference network $q_\phi(z|x)$.

Consider using IWAEs when:

• You need more accurate estimates of the data log-likelihood.
• You suspect that the looseness of the standard ELBO is a limiting factor for your model's performance.
• You have the computational budget to accommodate the $K$-fold increase in computation per example.
• You want to improve VAE performance without resorting to more complex inference network architectures immediately.

IWAEs represent a significant step up from the basic VAE objective. They demonstrate how rethinking the estimation of the objective itself, rather than just the model components, can lead to substantial gains in generative modeling. Next, we will explore techniques that directly enhance the expressiveness of the approximate posterior $q_\phi(z|x)$.