Energy-Calibrated VAE with Test Time Free Lunch

TL;DR

EC-VAE presents a principled way to sharpen VAE-generated samples by introducing a conditional energy-based model that calibrates the generative direction during training. By formulating a constrained primal-dual optimization, the method trains the VAE with generated samples and briefly runs MCMC to obtain calibrated targets, while preserving one-step inference at test time. The calibration concept extends to variational learning and normalizing flows, and the framework enables zero-shot image restoration via neural transport priors and range-null theory. Empirically, EC-VAE achieves competitive or superior single-step generation across multiple datasets, delivers substantial gains in sampling efficiency, and demonstrates strong zero-shot restoration performance, highlighting practical impact for efficient high-quality generative modeling.

Abstract

In this paper, we propose a novel generative model that utilizes a conditional Energy-Based Model (EBM) for enhancing Variational Autoencoder (VAE), termed Energy-Calibrated VAE (EC-VAE). Specifically, VAEs often suffer from blurry generated samples due to the lack of a tailored training on the samples generated in the generative direction. On the other hand, EBMs can generate high-quality samples but require expensive Markov Chain Monte Carlo (MCMC) sampling. To address these issues, we introduce a conditional EBM for calibrating the generative direction of VAE during training, without requiring it for the generation at test time. In particular, we train EC-VAE upon both the input data and the calibrated samples with adaptive weight to enhance efficacy while avoiding MCMC sampling at test time. Furthermore, we extend the calibration idea of EC-VAE to variational learning and normalizing flows, and apply EC-VAE to an additional application of zero-shot image restoration via neural transport prior and range-null theory. We evaluate the proposed method with two applications, including image generation and zero-shot image restoration, and the experimental results show that our method achieves competitive performance over single-step non-adversarial generation. Our code is available at https://github.com/DJ-LYH/EC-VAE.

Figures (13)

• Figure 1: Random generated samples from EC-VAE. For CelebA 64 and CelebA-HQ-256, we pick out samples for diversity.
• Figure 2: Comparison of EC-VAE (Top) and EC-VAE (w/ MCMC) (Bottom) on CIFAR-10. Best viewed when zoomed in.
• Figure 3: Qualitative results of zero-shot image restoration (colorization, inpainting, 4$\times$ super-resolution).
• Figure 4: Qualitative results on the 25-Gaussians dataset.
• Figure 5: Qualitative comparison of Energy-Calibrated VAE (Ours) and other models on CIFAR-10. Samples are un-selected.