Energy-Calibrated VAE with Test Time Free Lunch

TL;DR

This paper introduces Energy-Calibrated VAE (EC-VAE), a novel generative model that improves sample quality by calibrating the VAE with a conditional EBM during training, without MCMC at test time, and applies it to image generation and restoration.

Contribution

The paper proposes EC-VAE, which enhances VAE training with a conditional EBM calibration, extending to variational learning and normalizing flows, and demonstrates its effectiveness in image tasks.

Findings

Achieves competitive image generation quality.

Improves zero-shot image restoration performance.

Avoids MCMC sampling at test time.

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.