PaperlayerView exact revision on arXiv

Source-linked AI summary

Lagging Inference Networks and Posterior Collapse in Variational Autoencoders

Junxian He, Daniel Spokoyny, Graham Neubig, Taylor Berg-Kirkpatrick

arXiv:1901.05534v2cs.LGstat.ML

TL;DR

VAEs can collapse by ignoring latent variables when strong autoregressive generators are used, motivating an explanation of this failure through training dynamics. The paper identifies lagging inference networks as an initial cause and adds aggressive inference updates before model updates. The resulting method avoids collapse and matches or exceeds strong baselines across text and image benchmarks while remaining substantially cheaper than more complex approaches.

  • Problem

    VAE training often reaches a posterior-collapse local optimum in which latent variables are ignored and both posteriors match the prior, especially with strong autoregressive generators.

  • Method

    The method aggressively optimizes the amortized inference network with extra updates, using mutual information to determine when to stop before returning to basic VAE training.

  • Results

    The approach avoids posterior collapse and achieves comparable or better performance than previous state-of-the-art systems across Yahoo, Yelp, and OMNIGLOT benchmarks.

  • Takeaways & Limitations

    A simple training-procedure change can address posterior collapse without new modeling components or significant complexity over basic VAEs.

  • Takeaways & Limitations

    The posterior-mean-space analysis notes that all points at the origin need not constitute a local optimum, such as when the model is initialized there.

Abstract

from arXiv · show

The variational autoencoder (VAE) is a popular combination of deep latent variable model and accompanying variational learning technique. By using a neural inference network to approximate the model's posterior on latent variables, VAEs efficiently parameterize a lower bound on marginal data likelihood that can be optimized directly via gradient methods. In practice, however, VAE training often results in a degenerate local optimum known as "posterior collapse" where the model learns to ignore the latent variable and the approximate posterior mimics the prior. In this paper, we investigate posterior collapse from the perspective of training dynamics. We find that during the initial stages of training the inference network fails to approximate the model's true posterior, which is a moving target. As a result, the model is encouraged to ignore the latent encoding and posterior collapse occurs. Based on this observation, we propose an extremely simple modification to VAE training to reduce inference lag: depending on the model's current mutual information between latent variable and observation, we aggressively optimize the inference network before performing each model update. Despite introducing neither new model components nor significant complexity over basic VAE, our approach is able to avoid the problem of collapse that has plagued a large amount of previous work. Empirically, our approach outperforms strong autoregressive baselines on text and image benchmarks in terms of held-out likelihood, and is competitive with more complex techniques for avoiding collapse while being substantially faster.

1 INTRODUCTION

VAEs commonly suffer posterior collapse, especially with strong autoregressive generators, because training can settle on solutions that ignore latent variables. This paper instead analyzes the training dynamics and proposes aggressive inference-network optimization without changing the model or standard ELBO.

  • 1 INTRODUCTION: Posterior collapse occurs when VAEs ignore latent variables, particularly with strong autoregressive generators such as LSTMs and PixelCNNs.The problem is especially evident for discrete data and hinders applications such as natural language processing.
  • 1 INTRODUCTION: Earlier approaches weaken the generator or modify the training objective because the collapsed solution can be a reasonably good ELBO local optimum.The paper frames these approaches as responses to a static optimization perspective.
  • 1 INTRODUCTION: The paper studies posterior collapse through training dynamics rather than only static optimization.It investigates why basic VAE training fails to find better latent-using local optima.
  • 1 INTRODUCTION: The proposed procedure aggressively optimizes the inference network with additional updates to reduce lag before reverting to basic VAE training.It introduces neither new modeling components nor additional complexity over basic VAEs.
  • 1 INTRODUCTION: On Yahoo, Yelp, and OMNIGLOT, the method outperforms neural autoregressive baselines and achieves comparable performance to more complex methods at a fraction of the training cost.The reported evaluations cover text and image density-estimation benchmarks.

2 BACKGROUND

VAEs optimize a tractable ELBO using an amortized inference network that approximates the model posterior. Posterior collapse is the degenerate state where both posteriors equal the prior, and posterior mean space visualizes how inference and model posteriors align during training.

  • 2.1 VARIATIONAL AUTOENCODERS: VAEs use a prior and conditional generator, optimizing a tractable ELBO when marginal data likelihood is intractable.The ELBO contains reconstruction and KL-regularization terms.
  • 2.1 VARIATIONAL AUTOENCODERS: The inference network qφ(z|x) approximates pθ(z|x), while shared parameters amortize inferential cost across data samples.This enables efficient training compared with separate local variational inference for each sample.
  • 2.2 POSTERIOR COLLAPSE: Posterior collapse is the ELBO local optimum qφ(z|x) = pθ(z|x) = p(z), where the model ignores latent variables and loses meaningful latent features.Model collapse and inference collapse describe intermediate states in which the model posterior or variational posterior, respectively, equals the prior.
  • 2.3 VISUALIZATION OF POSTERIOR DISTRIBUTION: Posterior mean space tracks the means of pθ(z|x) and qφ(z|x), with the diagonal indicating agreement between their means.The origin represents the collapsed local optimum, while points along the diagonal indicate a better-trained inference network.

3 A LAGGING INFERENCE NETWORK PREVENTS USING LATENT CODES

Basic VAE training can lose latent information because the approximate posterior lags behind the model posterior during optimization. This lag drives the model toward a collapsed local optimum.

  • 3 A LAGGING INFERENCE NETWORK PREVENTS USING LATENT CODES: Posterior collapse occurs when VAEs ignore latent variables and both posteriors move toward the prior.The collapsed solution is a local optimum of the ELBO objective.
  • 3.1 INTUITIONS FROM ELBO: The ELBO balances marginal data likelihood against posterior agreement, with the KL term driving the model posterior toward the approximate posterior.If the approximate posterior is perfect, the agreement force vanishes.
  • 3.2 OBSERVATIONS ON SYNTHETIC DATA: In a synthetic discrete-sequence experiment, a basic VAE used a scalar latent variable with LSTM encoder and decoder to track posterior means during training.The study visualized 500 validation samples at four training stages.
  • 3.2 OBSERVATIONS ON SYNTHETIC DATA: Initially, both posteriors made latent variables nearly independent of observations, with all projected samples located at the origin.This initialization corresponds to simultaneous model and inference collapse in the posterior-mean space.
  • 3.2 OBSERVATIONS ON SYNTHETIC DATA: During basic training, model posterior means spread while approximate posterior means remained horizontal, showing that inference lag preceded convergence to collapse.The model posterior first moved away from the prior, but the inference network failed to catch up.

4 METHOD

The method separates inference and generator optimization by aggressively updating the inference network before generator updates. It stops these extra updates when mutual information no longer increases, retaining efficient amortized training.

  • 4 METHOD: The proposed procedure uses an inner loop that aggressively optimizes qφ(z|x) before each generator update.This separates inference-network and generator optimization to reduce posterior lag.
  • 4 METHOD: The approach remains fully amortized and avoids additional instance-specific SVI steps while matching the simplicity of basic VAE training.The authors contrast it with methods that use costly local inference updates.
  • 4 METHOD: The paper separately examines whether the method works under inference-network initialization away from the origin.This setting is studied in Appendix G.
  • 4 METHOD: Aggressive updates are stopped when validation mutual information Iq stops climbing, allowing training to revert to the basic VAE procedure.Across the reported experiments, reversion usually occurred within 5 epochs.
  • 4.3 OBSERVATIONS ON SYNTHETIC DATASET: On synthetic data, aggressive inference moved posterior-mean points toward the diagonal µx,θ = µx,φ and helped escape inference collapse.The diagonal indicates agreement between model and approximate posterior means.

5 RELATION TO RELATED WORK

Prior work addresses posterior collapse through objective modifications, weaker decoders, alternative regularizers, and hybrid inference, while this approach targets amortization-related inference limitations with lower training cost.

  • Objective and decoder modifications: KL cost annealing increases the KL weight during a warm-up period but has been unable to prevent collapse on complex text datasets with very large LSTM decoders.The cited work identifies this limitation for annealing-based approaches.
  • Inference-based approaches: Amortization gap denotes the ELBO difference caused by parameter sharing in the inference network.Several methods combine instance-specific and amortized inference to narrow this gap.
  • Objective and decoder modifications: Other approaches modify regularizers, decoder strength, input information, latent connections, or the prior to address posterior collapse.These methods represent distinct intervention strategies rather than a single common modification.

6 EXPERIMENTS

Experiments evaluate collapse prevention, predictive log-likelihood, and training efficiency across Yahoo, Yelp, and OMNIGLOT using standard and competing VAE methods. The proposed method performs comparably or better across datasets, avoids collapse without annealing, and is substantially faster than SA-VAE.

  • 6 EXPERIMENTS: Experiments target posterior-collapse prevention, predictive log-likelihood, and training efficiency.The study compares the proposed method with autoregressive models, basic VAE, KL annealing, β-VAE, and SA-VAE.
  • 6.1 SETUP: NLL is estimated with 500 importance-weighted samples, alongside KL, mutual information, and active-unit metrics.A latent dimension is active when Az > 0.01; the NLL estimate variance is smaller than 10^-3 on all datasets.
  • 6.1 SETUP: The evaluation covers text density estimation on Yahoo and Yelp and image density estimation on OMNIGLOT.Text models use single-layer LSTM encoders and decoders; image models use a ResNet encoder and 13-layer Gated PixelCNN decoder.
  • 6.2 RESULTS: The proposed method achieves comparable or better performance than previous state-of-the-art systems on all three datasets.The paper also reports that it avoids posterior collapse without annealing, whereas SA-VAE collapses on both text datasets without annealing.
  • 6.3 TRAINING TIME: The proposed algorithm is only 2–3 times slower than regular VAE training and 3–7 times faster than SA-VAE.Training times are compared across all three datasets.
  • 6.4 ANALYSIS OF BASELINES: On Yelp, annealing maintains small mutual information despite an increasing KL regularizer and eventually collapses, while β-VAE produces a very large aggregated-posterior divergence.The analysis tracks mutual information, the KL regularizer, and DKL(qφ(z)∥p(z)).
  • 6.5 ANALYSIS OF INNER LOOP UPDATE: Sufficient inner-loop inference updates are necessary to address collapse and achieve good performance, while performance saturates near convergence.The implementation typically uses roughly 30–100 inference-network updates per inner-loop update and stops after the ELBO is unchanged or decreases across 10 iterations.

7 CONCLUSION

The paper identifies lagging inference networks as the source of posterior collapse and proposes additional inference-network updates before returning to basic VAE training.

  • 7 CONCLUSION: Posterior collapse is linked to a lagging inference network during the initial stages of training.The synthetic experiment identifies this training-dynamics problem.
  • 7 CONCLUSION: The proposed remedy aggressively optimizes the inference network with more updates before reverting to basic VAE training.The procedure is described as simple yet effective.
  • 7 CONCLUSION: Experiments on text and image modeling demonstrate the effectiveness of the proposed approach.

B EXPERIMENTAL DETAILS

The experiments use KL-weight annealing, a synthetic Gaussian-mixture setup, recurrent architectures for text, and specified optimization and PixelCNN configurations for images.

  • B EXPERIMENTAL DETAILS: The KL weight increases linearly from 0.1 to 1.0 during the first 10 epochs.This follows the annealing strategy of Kim et al. (2018).
  • B EXPERIMENTAL DETAILS: Synthetic data are generated from a two-dimensional latent variable sampled from a four-component Gaussian mixture.The dimension balances distributional simplicity with enough complexity for a one-dimensional latent-variable model.
  • B EXPERIMENTAL DETAILS: Synthetic experiments use LSTM encoders and decoders with 50 hidden units and 50 latent embeddings.
  • B EXPERIMENTAL DETAILS: Text models use single-layer LSTMs with 1024 hidden units and 512-dimensional word embeddings for both encoder and decoder.The final encoder hidden state predicts the latent variable through a linear transformation.
  • B EXPERIMENTAL DETAILS: Image experiments use dynamically binarized inputs, binary likelihood, a ResNet, and a 13-layer PixelCNN architecture.The PixelCNN uses five 7 x 7, four 5 x 5, and four 3 x 3 layers, each with 64 feature maps.

C ADDITIONAL RESULTS CONTAINING ELBO

The additional ELBO results report mean performance across five random restarts for Yahoo, Yelp, and OMNIGLOT, with standard deviations when available.

  • C ADDITIONAL RESULTS CONTAINING ELBO: Yahoo and Yelp results report mean values across 5 different random restarts.Standard deviations are given in parentheses when available; LSTM-LM∗ uses exact negative log likelihood.
  • C ADDITIONAL RESULTS CONTAINING ELBO: OMNIGLOT results report mean values across 5 different random restarts.Standard deviations are given in parentheses when available; PixelCNN∗ uses exact negative log likelihood.

D UNCERTAINTY OF EVALUATION

Evaluation uncertainty is assessed by repeating evaluation with 10 random seeds for models trained by the proposed approach and by basic VAE training.

  • D UNCERTAINTY OF EVALUATION: Evaluation uses 10 different random seeds to measure uncertainty from random Monte Carlo samples.The reported values include means and variances in Tables 7 and 8.
  • D UNCERTAINTY OF EVALUATION: Table 7 evaluates a trained VAE model produced by the proposed approach, reporting means and variances.IW denotes the approximation to NLL used in Section 6.
  • D UNCERTAINTY OF EVALUATION: Table 8 evaluates a trained VAE model produced by basic VAE training, reporting means and variances.IW denotes the approximation to NLL used in Section 6.

E COMPARISON WITH DIFFERENT KL-ANNEALING SCHEDULES

The paper compares KL-annealing schedules and finds that sufficiently slow annealing can avoid posterior collapse, though it does not outperform the proposed approach and requires dataset- and architecture-specific tuning. Separate encoder learning-rate adjustments also fail to resolve collapse and can cause KL-value explosion when too large.

  • KL-annealing schedules: Sufficiently slow KL annealing avoids posterior collapse, but does not produce superior predictive log likelihood to the proposed approach.The annealing duration is varied across 30k, 50k, 100k, and 120k iterations.
  • KL-annealing schedules: The optimal KL-annealing schedule varies across datasets and model architectures, requiring careful task-specific tuning.
  • Separate learning rates: Decoder gradients are much larger than encoder gradients in the reported experiments, motivating separate learning rates.
  • Separate learning rates: All reported encoder-learning-rate settings suffer posterior collapse, while overly large rates cause KL-value explosion.The text states that settings including 8x, 10x, 30x, and 50x produced KL-value explosion and were excluded from Table 10.

G DISCUSSION ABOUT INITIALIZATION OF INFERENCE NETWORKS

The paper examines whether inference-network initialization affects the observed lagging behavior and compares basic VAE training with aggressive inference-network training during optimization.

  • Initialization analysis: The initialization analysis addresses whether different inference-network initializations also suffer from posterior collapse.The reported experiment adds an offset to the uniform initialization used previously.
  • Figure 6: Figure 6 projects 500 synthetic-data samples onto posterior-mean space across generator updates.The figure uses “iter” to denote the number of generator updates.
  • Figure 6: Basic VAE training leaves the approximate posterior far behind the true model posterior, whereas aggressive inference training moves points onto the diagonal.
Loading 1901.05534v2…