Relaxed-Responsibility Hierarchical Discrete VAEs

TL;DR

This paper introduces a novel stable training method for deep hierarchical discrete VAEs using Relaxed-Responsibility Vector-Quantisation, enabling end-to-end training of models with up to 32 layers of discrete latent variables and achieving state-of-the-art results.

Contribution

It proposes a new parameterization technique for discrete latent variables that improves stability and performance in hierarchical VAEs, allowing end-to-end training with many layers.

Findings

Achieved state-of-the-art bits-per-dim on standard datasets.

Enabled end-to-end training of models with up to 32 layers of discrete latents.

Produced samples via ancestral sampling without additional autoregressive models.

Abstract

Successfully training Variational Autoencoders (VAEs) with a hierarchy of discrete latent variables remains an area of active research. Vector-Quantised VAEs are a powerful approach to discrete VAEs, but naive hierarchical extensions can be unstable when training. Leveraging insights from classical methods of inference we introduce \textit{Relaxed-Responsibility Vector-Quantisation}, a novel way to parameterise discrete latent variables, a refinement of relaxed Vector-Quantisation that gives better performance and more stable training. This enables a novel approach to hierarchical discrete variational autoencoders with numerous layers of latent variables (here up to 32) that we train end-to-end. Within hierarchical probabilistic deep generative models with discrete latent variables trained end-to-end, we achieve state-of-the-art bits-per-dim results for various standard datasets. % Unlike discrete VAEs with a single layer of latent variables, we can produce samples by ancestral sampling: it is not essential to train a second autoregressive generative model over the learnt latent representations to then sample from and then decode. % Moreover, that latter approach in these deep hierarchical models would require thousands of forward passes to generate a single sample. Further, we observe different layers of our model become associated with different aspects of the data.