Iterative Amortized Hierarchical VAE

TL;DR

The paper introduces IA-HVAE, a hybrid inference model combining amortized and iterative methods, achieving faster inference and improved accuracy in inverse problems like deblurring and denoising.

Contribution

It presents a novel hybrid inference scheme with a linearly separable decoder, significantly speeding up hierarchical VAE inference and enhancing reconstruction quality.

Findings

35x faster inference compared to traditional HVAE

Outperforms fully amortized and iterative models in accuracy and speed

Improves reconstruction quality in inverse problems

Abstract

In this paper we propose the Iterative Amortized Hierarchical Variational Autoencoder (IA-HVAE), which expands on amortized inference with a hybrid scheme containing an initial amortized guess and iterative refinement with decoder gradients. We achieve this by creating a linearly separable decoder in a transform domain (e.g. Fourier space), enabling real-time applications with very high model depths. The architectural change leads to a 35x speed-up for iterative inference with respect to the traditional HVAE. We show that our hybrid approach outperforms fully amortized and fully iterative equivalents in accuracy and speed respectively. Moreover, the IAHVAE shows improved reconstruction quality over a vanilla HVAE in inverse problems such as deblurring and denoising.