Learning Energy-based Variational Latent Prior for VAEs

2025-09-30

AI Summary

VAEs often produce blurry and inconsistent samples due to “prior holes”—regions of high prior probability that exhibit low posterior density, causing sampling mismatch. To address this, we propose the Energy-Based Model-enhanced VAE (EBM-VAE), which employs a learnable energy function to model a flexible latent prior, thereby mitigating prior holes. Our key contributions are: (i) a lightweight sampling network that approximates the EBM’s intractable partition function, eliminating the need for MCMC and enabling end-to-end differentiable training and fast generation; and (ii) an alternating optimization framework jointly updating VAE and EBM parameters. Experiments demonstrate substantial improvements in image generation quality—reducing FID and LPIPS—and an ELBO gain of +0.5–1.2 bits/dim. Moreover, EBM-VAE achieves 5–10× faster sampling than standard EBM-based approaches, striking an effective balance between fidelity and efficiency.

Technology Category

• Machine Learning: Deep Generative Models & Autoencoders
• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Variational Auto-Encoders (VAEs) are known to generate blurry and inconsistent samples. One reason for this is the "prior hole" problem. A prior hole refers to regions that have high probability under the VAE's prior but low probability under the VAE's posterior. This means that during data generation, high probability samples from the prior could have low probability under the posterior, resulting in poor quality data. Ideally, a prior needs to be flexible enough to match the posterior while retaining the ability to generate samples fast. Generative models continue to address this tradeoff. This paper proposes to model the prior as an energy-based model (EBM). While EBMs are known to offer the flexibility to match posteriors (and also improving the ELBO), they are traditionally slow in sample generation due to their dependency on MCMC methods. Our key idea is to bring a variational approach to tackle the normalization constant in EBMs, thus bypassing the expensive MCMC approaches. The variational form can be approximated with a sampler network, and we show that such an approach to training priors can be formulated as an alternating optimization problem. Moreover, the same sampler reduces to an implicit variational prior during generation, providing efficient and fast sampling. We compare our Energy-based Variational Latent Prior (EVaLP) method to multiple SOTA baselines and show improvements in image generation quality, reduced prior holes, and better sampling efficiency.

Problem

• Addressing blurry sample generation in Variational Auto-Encoders
• Solving the prior hole problem in VAEs' latent distributions
• Improving sampling efficiency while maintaining flexible priors

Innovation

• Energy-based model for flexible VAE prior
• Variational approach bypasses expensive MCMC sampling
• Alternating optimization with sampler network enables fast generation