Auxiliary Variables and Semi-Amortized Variational Inference

As we've discussed, the standard amortized inference network in VAEs, while efficient, often employs a simple distribution (like a diagonal Gaussian) for $q_\phi(z|x)$. This simplicity can be a bottleneck, preventing $q_\phi(z|x)$ from accurately approximating a potentially complex true posterior $p_\theta(z|x)$. This section explores two powerful strategies to create more expressive and accurate approximate posteriors: leveraging auxiliary variables and employing semi-amortized inference schemes.

Augmenting Inference with Auxiliary Variables

One way to enhance the flexibility of the approximate posterior $q_\phi(z|x)$ without making its direct functional form overly complex is to introduce auxiliary random variables. These variables are not part of the original generative model $p_\theta(x|z)$ but are used within the inference network to help shape a richer distribution for $z$.

Let's denote these auxiliary variables by $a$. Instead of defining $q_\phi(z|x)$ directly, we define a joint distribution $q_\phi(z, a|x)$ over both the original latent variables $z$ and these new auxiliary variables $a$. A common factorization for this joint distribution is hierarchical:

$$q_\phi(z,a|x) = q_\phi(z|x,a) q_\phi(a|x)$$

Here, $q_\phi(a|x)$ is an inference network that maps an input $x$ to the parameters of a distribution over $a$. Then, $q_\phi(z|x,a)$ is another inference network that maps $x$ and a sample $a \sim q_\phi(a|x)$ to the parameters of a distribution over $z$.

The resulting marginal distribution for $z$, $q_\phi(z|x) = \int q_\phi(z|x,a) q_\phi(a|x) da$, can be significantly more complex and flexible than if we had modeled $q_\phi(z|x)$ directly with a simple family (e.g., a single Gaussian). Think of it as using $a$ to "steer" or "refine" the inference for $z$. For example, $q_\phi(a|x)$ could capture some high-level aspects of the posterior, and $q_\phi(z|x,a)$ could then model finer details conditioned on these aspects.

To incorporate this into the VAE framework, we adjust the Evidence Lower Bound (ELBO). We treat $a$ as additional latent variables and assume a simple prior $p(a)$ for them (e.g., a standard Normal distribution, $p(a) = \mathcal{N}(0,I)$). The ELBO for this augmented system is:

$$\mathcal{L}(x; \theta, \phi) = \mathbb{E}_{q_\phi(z,a|x)} [\log p_\theta(x|z)] - D_{KL}(q_\phi(z,a|x) || p(z)p(a))$$

Note that the generative model $p_\theta(x|z)$ still only depends on $z$. The auxiliary variables $a$ are only "seen" by the inference machinery and the prior $p(a)$. Using our chosen factorization $q_\phi(z,a|x) = q_\phi(z|x,a)q_\phi(a|x)$ and assuming $p(z,a) = p(z)p(a)$ (i.e., $z$ and $a$ are independent in the prior), the KL divergence term can be decomposed:

$$D_{KL}(q_\phi(z,a|x) || p(z)p(a)) = \mathbb{E}_{q_\phi(a|x)}[D_{KL}(q_\phi(z|x,a) || p(z))] + D_{KL}(q_\phi(a|x) || p(a))$$

So, the ELBO becomes:

$$\mathcal{L} = \mathbb{E}_{q_\phi(a|x)} \left[ \mathbb{E}_{q_\phi(z|x,a)}[\log p_\theta(x|z)] - D_{KL}(q_\phi(z|x,a) || p(z)) \right] - D_{KL}(q_\phi(a|x) || p(a))$$

This expression looks like a VAE objective where $q_\phi(a|x)$ acts as an "encoder" for $a$, and then, conditioned on $a$, $q_\phi(z|x,a)$ acts as another "encoder" for $z$. The overall structure allows $q_\phi(z|x)$ to implicitly represent a mixture of simpler distributions, leading to a much richer family for the approximate posterior.

Benefits:

• Increased Expressiveness: The primary benefit is a more flexible $q_\phi(z|x)$, capable of better matching the true posterior.
• Tighter ELBO: A better posterior approximation generally leads to a tighter ELBO, which can improve both reconstruction quality and the learned representations.

Costs:

• Increased Complexity: The inference network becomes more complex, involving more parameters and computational steps.
• Optimization Challenges: Training these more complex inference networks can sometimes be more challenging.

Models like Auxiliary Deep Generative Models (ADGMs) and some variants of hierarchical VAEs (when applied to the inference side) are examples of this approach. This technique is distinct from Normalizing Flows (which transform a simple noise distribution into a complex posterior using invertible functions) but can be complementary.

Semi-Amortized Variational Inference: Refining the Posterior Per Datapoint

Amortized variational inference, where a single neural network $q_\phi(z|x)$ directly outputs the parameters of the approximate posterior for any given $x$, is computationally efficient. However, it makes a strong assumption: that a single set of network parameters $\phi$ can provide optimal (or near-optimal) variational parameters for all datapoints. This can lead to an "amortization gap", the difference in ELBO quality between what a fully amortized $q_\phi(z|x)$ can achieve and what could be achieved if we optimized the variational parameters for each datapoint individually.

Semi-amortized variational inference aims to bridge this gap. The core idea is to use the amortized inference network to provide a good initialization for the variational parameters for a specific datapoint $x_i$. Then, these initial parameters are refined through a few steps of optimization, specifically for that $x_i$, by directly maximizing the ELBO with respect to the variational parameters for that instance.

Let $\lambda$ denote the parameters of the approximate posterior for a single datapoint $x_i$ (e.g., if $q(z|x_i)$ is Gaussian, $\lambda = (\mu_i, \sigma_i)$). The process is as follows:

1. Initialization: Obtain initial variational parameters $\lambda_0$ using the amortized encoder: $$\lambda_0 = \text{Encoder}_\phi(x_i)$$
2. Refinement: Iteratively update $\lambda$ to maximize the instance-specific ELBO, $\mathcal{L}(x_i, \lambda)$: For $t = 0, \dots, T-1$: $$\lambda_{t+1} = \lambda_t + \eta \nabla_{\lambda_t} \mathcal{L}(x_i, \lambda_t)$$ where $\eta$ is a learning rate for this refinement process, and $T$ is the number of refinement steps.
3. Final Posterior: Use the refined parameters $\lambda_T$ to define the approximate posterior $q(z|x_i; \lambda_T)$ for downstream tasks (e.g., calculating the ELBO for training the VAE's generative model $p_\theta(x|z)$).

The following diagram illustrates this refinement process:

The semi-amortized inference process: An amortized network provides an initial estimate of posterior parameters, which are then refined through instance-specific optimization.

Benefits:

• Improved Posterior Fit: By optimizing parameters for each datapoint, semi-amortized VI can achieve a much better fit to the true posterior $p_\theta(z|x_i)$ for that specific instance.
• Tighter ELBO: This improved fit usually translates to a tighter ELBO.
• Flexibility: Allows the model to adapt its inference more precisely to individual data characteristics.

Costs:

• Increased Inference Time: The iterative refinement process makes inference significantly slower than purely amortized approaches, as optimization steps are performed for each datapoint.
• Training Complexity: If the parameters $\phi$ of the amortized encoder are to be trained to produce good initializations $\lambda_0$ that lead to fast and effective refinement, one might need to differentiate through the $T$ steps of optimization. This can be computationally intensive and complex (though approximations or simpler training schemes exist, e.g., training $\phi$ to simply produce good $\lambda_0$ without backpropagating through the refinement).

This approach is particularly useful when the true posterior has high variance across different datapoints, making it difficult for a single amortized network to perform well universally. The number of refinement steps $T$ is a hyperparameter; even a small number of steps (e.g., $T=5$ to $10$) can often yield substantial improvements.

Synergies and Practical Considerations

Auxiliary variables and semi-amortized inference are not mutually exclusive. One could, for instance, define an expressive posterior family using auxiliary variables and then use semi-amortized inference to fine-tune the parameters of this richer posterior for each data point.

When deciding whether to use these advanced inference techniques, consider the following:

• Performance vs. Computation: Both methods generally improve VAE performance (tighter ELBO, potentially better samples or representations) but at the cost of increased computation, especially for semi-amortized inference during prediction/generation.
• Complexity of True Posterior: If you suspect the true posterior $p_\theta(z|x)$ is highly complex or varies significantly across data, these methods are strong candidates.
• Implementation Effort: Implementing semi-amortized inference, especially the training part involving backpropagation through optimization, can be more involved. Auxiliary variable models also add a layer of complexity to the inference network design.

In practice, starting with a well-tuned standard VAE and then exploring these techniques can be a good strategy if further improvements in posterior approximation are needed. The choice depends on the specific application, available computational resources, and the desired trade-off between model performance and inference speed. Both methods offer valuable tools for pushing the boundaries of what VAEs can achieve by enabling more accurate and flexible posterior inference.