PaperlayerView exact revision on arXiv

Source-linked AI summary

Semi-Amortized Variational Autoencoders

Yoon Kim, Sam Wiseman, Andrew C. Miller, David Sontag, Alexander M. Rush

arXiv:1802.02550v7stat.MLcs.CLcs.LG

TL;DR

Amortized inference can yield suboptimal variational parameters, motivating a hybrid approach. SA-VAE initializes them with an inference network, refines them through differentiable SVI, and trains end-to-end; experiments report stronger text and image performance and non-collapsed latent representations.

  • Problem

    Amortized inference is fast but can impose a restrictive parametric mapping that produces suboptimal variational parameters and an amortization gap.

  • Method

    SA-VAE initializes variational parameters with an inference network, refines them using SVI, and differentiates through those updates to train end-to-end.

  • Results

    SA-VAE outperforms autoregressive, VAE, SVI, and non-end-to-end VAE-plus-SVI baselines on text and image models while retaining non-trivial latent representations.

  • Takeaways & Limitations

    The approach enables expressive autoregressive generative models for text without the latent-code neglect associated with posterior collapse.

  • Takeaways & Limitations

    Each training step requires multiple backpropagations through the generative model, and the approach mainly applies to variational families supporting differentiable optimization.

Abstract

from arXiv · show

Amortized variational inference (AVI) replaces instance-specific local inference with a global inference network. While AVI has enabled efficient training of deep generative models such as variational autoencoders (VAE), recent empirical work suggests that inference networks can produce suboptimal variational parameters. We propose a hybrid approach, to use AVI to initialize the variational parameters and run stochastic variational inference (SVI) to refine them. Crucially, the local SVI procedure is itself differentiable, so the inference network and generative model can be trained end-to-end with gradient-based optimization. This semi-amortized approach enables the use of rich generative models without experiencing the posterior-collapse phenomenon common in training VAEs for problems like text generation. Experiments show this approach outperforms strong autoregressive and variational baselines on standard text and image datasets.

1. Introduction

Modern variational inference scales to large datasets and rich, non-conjugate models, but amortized inference can produce suboptimal instance-specific parameters. SA-VAE addresses this gap by differentiating through iterative SVI refinement during end-to-end training.

  • 1. Introduction: Stochastic variational inference provides good instance-specific distributions but requires optimization for each data point, whereas amortized inference offers fast inference through a shared network.AVI’s parametric restriction can create an amortization gap, especially on complex datasets.
  • 1. Introduction: The amortization gap can cause generative-model parameters to be updated using suboptimal variational parameters.
  • 1. Introduction: Prior hybrid methods initialized local parameters with an encoder and refined them iteratively, but trained using separate objectives.
  • 1. Introduction: SA-VAE differentiates through SVI while training the inference network and generative model end-to-end.The method improves variational-parameter estimation and produces better generative models.
  • 1. Introduction: On text and image models, SA-VAE outperforms autoregressive, VAE, SVI, and non-end-to-end VAE-plus-SVI baselines.It also uses an LSTM generative model with non-trivial latent representations without posterior collapse.

2. Background

Variational inference optimizes tractable distributions to approximate intractable posteriors. SVI optimizes local parameters iteratively, while AVI predicts them with a shared encoder, trading speed for possible amortization error.

  • 2. Background: The generative model defines p(z) and p(x | z; θ), while VI uses q(z; λ) and maximizes the ELBO because log p(x; θ) is usually intractable.
  • 2. Background: Posterior collapse occurs when KL[q(z; λ) ∥ p(z)] is approximately zero.
  • 2. Background: SVI initializes each data point’s variational parameters and uses gradient ascent on data subsets, enabling scalable inference.Its iterative updates can require many steps and separate local-parameter optimization from generative-parameter optimization.
  • 2. Background: AVI predicts local variational parameters with a shared parametric inference network, making inference faster than iterative ELBO optimization.The input-dependent parametric restriction can create an amortization gap and hinder generative-model learning when parameters are suboptimal.

3. Semi-Amortized Variational Autoencoders

SA-VAE combines encoder-based initialization with local SVI refinement, then differentiates through the refinement procedure to train the encoder and generative model jointly. Its implementation uses momentum updates, Hessian-vector products, finite differences, and gradient clipping.

  • 3. Semi-Amortized Variational Autoencoders: SA-VAE uses an inference network to initialize variational parameters and subsequently runs SVI to refine them.
  • 3. Semi-Amortized Variational Autoencoders: Training requires total derivatives through the SVI updates so encoder and generative-model parameters receive gradients from the final ELBO.This entails backpropagating through gradient ascent.
  • 3. Semi-Amortized Variational Autoencoders: Algorithm 1 performs momentum-based variational updates, evaluates the final loss, and backpropagates through the iterative refinement steps.
  • 3. Semi-Amortized Variational Autoencoders: Finite-difference Hessian-vector products improve memory efficiency for scaling to rich inference networks and generative models.The implementation uses ϵ = 10^-5 and clips gradients to mitigate exploding gradients.
  • 3. Semi-Amortized Variational Autoencoders: Without gradient clipping, the model empirically relied too much on iterative inference, leaving the inference network’s initial parameters poor.
  • 3. Semi-Amortized Variational Autoencoders: In an ELBO landscape for one test point, SA-VAE’s variational parameters are closest to the grid-search optimum after 20 refinement iterations.The figure compares VAE, SVI, and SA-VAE initial, intermediate, and final parameters.

4. Experiments

Experiments on synthetic data, text, and images show that semi-amortized inference improves variational estimation and generative modeling, while revealing both interpretability benefits and computational and optimization constraints.

  • Experimental setup: SA-VAE uses a diagonal-Gaussian posterior and trains by stochastic gradient descent with momentum on the negative ELBO.The prior is N(0, I), and variational parameters are the mean vector and diagonal log variance vector.
  • Synthetic Data: On synthetic data, SA-VAE produces variational parameters closest to the optimum and outperforms VAE and SVI under both oracle and learned generators.The comparison uses 20 SVI steps for SVI and SA-VAE, with the optimum estimated by grid search.
  • Text: On Yahoo text modeling, SA-VAE uses the latent space and modestly outperforms a language model, whereas VAE/SVI nearly collapse to language-model behavior.VAE + SVI and VAE + SVI + KL do not outperform the language model; semi-amortized models have KL terms significantly above zero.
  • Text: Inference-network initialization gives better test-time SVI starting points than random initialization, while additional refinement beyond training-time steps provides no observed improvement.SA-VAE without test-time refinement still has KL = 6.65 and PPL = 62.0.
  • Images: On OMNIGLOT, semi-amortized approaches outperform VAE/SVI baselines and use latent representations more, but do not outperform models with hierarchical latents or more sophisticated priors.Even 80 SVI steps do not match SA-VAE trained with 10 refinement steps.
  • Images: Saliency analyses associate latent usage with question type, content words, sequence length, rare tokens, and earlier positions, suggesting potentially interpretable language representations.Output saliency has a negative correlation with log-likelihood of -0.51, although the relationship does not always hold.
  • Limitations: The approach is costly because training backpropagates through the generative model multiple times and mainly applies to variational families supporting differentiable ELBO optimization.The authors suggest meta-learning and more efficient gradient use as possible mitigations.

6. Related Work

The work builds on iterative refinement of amortized variational parameters and differentiable optimization, while targeting powerful generative models for images and text.

  • Differentiable optimization: Differentiating through inference or optimization has been explored as a differentiable layer of deep models and for hyperparameter optimization.
  • Image modeling: Powerful autoregressive image models such as PixelCNN substantially improved log-likelihood and sample quality over earlier conditionally independent pixel models.
  • Text modeling: Text VAEs remain difficult because LSTM generative models can cause posterior collapse, leading the model to ignore the latent code.
  • Scope: The proposed approach addresses the amortization gap but not the approximation gap, so it can be combined with richer posterior or prior distributions.

7. Conclusion

Semi-amortized variational autoencoders combine amortized inference with local iterative refinement for deep generative models of text and images. The approach supports expressive text generators that use the latent code, while the model remains non-identifiable in principle.

  • Conclusion: Semi-amortized variational autoencoders combine amortized inference with local iterative refinement to train deep generative models of text and images.
  • Conclusion: The approach trains deep latent-variable text models with an expressive autoregressive generator that does not ignore the latent code.
  • Conclusion: An autoregressive model can factorize p(x) while remaining non-identifiable because it does not have to use the latent variable.
  • Future work: In finite-data regimes, using the latent variable might improve generalization through a potentially better inductive bias, but learning good models and representations remains future work.

A. Training Semi-Amortized Variational Autoencoders with Gradient Clipping

Stable training modifies the semi-amortized procedure to clip gradients at multiple stages. The clipping function uses parameter η, set to 5 in all experiments.

  • Gradient clipping: Gradient clipping at various stages was crucial for stable training in the modified algorithm.
  • Gradient clipping: The clipping function is controlled by a clipping parameter η in Algorithm 2.
  • Gradient clipping: η = 5 was used in all experiments, and finite-difference estimation also applied gradient clipping.

B. Experimental Details

The variational models use a spherical Gaussian prior and a diagonal Gaussian variational family. SVI models initialize their variational parameters randomly from a Gaussian with standard deviation 0.1.

  • Variational setup: All variational models use a spherical Gaussian prior and a diagonal Gaussian variational family parameterized by means and log variances.
  • Variational setup: SVI models randomly initialize variational parameters from a Gaussian with standard deviation 0.1.

B.1. Synthetic Data

The experiments use recurrent and convolutional inference/generative architectures across synthetic, text, and image settings. Training includes stochastic optimization and, for iterative-inference models, SVI refinement with gradient clipping.

  • Synthetic setup: Synthetic examples use a one-layer LSTM with 100 hidden units and five generated time steps.The vocabulary contains 1000 tokens, and embeddings are 100-dimensional.
  • Semi-amortized training: SA-VAE models combine an encoder initialization with iterative SVI updates, using momentum, learning rate 1.0, and gradient clipping.The iterative procedure performs forward and reverse updates through the variational objective.
  • Synthetic setup: The synthetic-data inference and generative models are trained with stochastic gradient descent using batch size 50 and learning rate 1.0.Learning-rate decay begins after validation performance stops improving, with a five-epoch delay before decay can trigger.
  • Image setup: The image model uses a three-block residual inference network and a 12-layer Gated PixelCNN generative model.The 32-dimensional variational sample is projected to image resolution, concatenated with the input image, and passed to PixelCNN.

C. Data Size/Model Capacity

The study varies inference-network capacity, training-set size, and generative-model capacity on OMNIGLOT. Iterative inference helps especially with less powerful models, while SA-VAE maintains a similar advantage across training-set sizes and uses latent representations more strongly.

  • Model capacity: Iterative inference helps more when the inference network and generative model are less powerful.This pattern is reported in the Table 4 capacity experiments.
  • Latent usage: Across all scenarios, SA-VAE has a much higher KL portion of the loss than the compared models.The authors interpret this as indicating greater use of latent representations by the learned generative models.

D. Input Saliency Analysis

The input-saliency analysis examines how token contribution varies by part of speech, position, and frequency. Part-of-speech tags are obtained with NLTK.

  • Saliency dimensions: Figure 5 analyzes input saliency by part-of-speech tag, token position, and frequency.The figure presents these analyses in left, center, and right panels, respectively.
  • Saliency definition: Input saliency for token x_t is defined using the encoder word embedding w_t.The supplied passage introduces the definition but does not include the full displayed expression.
  • Analysis procedure: Part-of-speech tagging in the input-saliency analysis is performed with NLTK.
Loading 1802.02550v7…