Iterative Amortized Hierarchical VAE
This work addresses efficiency and accuracy challenges in inverse problems for machine learning practitioners, representing an incremental improvement over existing methods.
The paper tackles the problem of slow iterative inference in hierarchical variational autoencoders by proposing a hybrid scheme with an initial amortized guess and iterative refinement using decoder gradients, achieving a 35x speed-up and improved reconstruction quality in tasks like deblurring and denoising.
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.