[Paper Review] Lagging Inference Networks and Posterior Collapse in Variational Autoencoders

Junxian He, Daniel Spokoyny+2 more|arXiv (Cornell University)|Jan 16, 2019

Generative Adversarial Networks and Image SynthesisComputer Science30 references127 citations

TL;DR

The paper analyzes posterior collapse in VAEs as a training-dynamics issue caused by lagging inference networks and introduces a simple, aggressive inference-update strategy that mitigates collapse without changing the model, improving held-out likelihood and training speed.

ABSTRACT

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.

Motivation & Objective

• Motivate and understand why posterior collapse occurs in VAEs with powerful decoders.
• Analyze training dynamics of the true posterior vs. approximate posterior during early training.
• Propose a simple, aggressive inference-network update procedure to reduce inference lag.
• Show that the method preserves ELBO optimization and does not require new model components.
• Demonstrate competitive density estimation performance with faster training across text and image benchmarks.

Proposed method

• Reformulate ELBO to emphasize the agreement term D_KL(q_phi(z|x)||p_theta(z|x)) and the model posterior p_theta(z|x).
• Identify posterior lag as a driver of collapse using posterior-mean trajectory visualizations in a low-dimensional space.
• Introduce aggressive inference-network updates as an inner loop that updates phi multiple times before updating theta, without modifying the objective.
• Use mutual information I_q between z and x under q to decide when to stop aggressive updates, keeping amortization.
• Maintain standard ELBO optimization while substantially increasing the frequency of inference-network updates.
• Compare against autoregressive baselines and prior collapse-avoidance methods, highlighting training-time efficiency.

Experimental results

Research questions

• RQ1 Why does basic VAE training with strong decoders often lead to posterior collapse?
• RQ2 Can a simple alteration to the training dynamics, without changing the model, prevent latent-variable collapse?
• RQ3 Does aggressively updating the inference network early in training help the model utilize latent codes while preserving ELBO optimization?
• RQ4 How does the proposed method perform in density estimation and training efficiency relative to existing collapse-avoidance techniques?

Key findings

• Posterior collapse is linked to lagging inference networks during early training, which drives the model toward ignoring latent variables.
• Aggressively updating the inference network before each generator update reduces lag and moves the trajectory toward utilizing latent codes.
• The method maintains ELBO optimization and yields competitive or superior held-out likelihoods on text and image benchmarks.
• The approach requires modest extra training time (2–3x vs. standard VAE) and is significantly faster than SA-VAE while avoiding collapse.
• Mutual-information-based stopping criteria effectively prevent unnecessary inner-loop updates once the lag is mitigated.
• Experiments on Yahoo, Yelp, and OMNIGLOT show comparable performance to state-of-the-art methods with substantially reduced complexity.