Multi-Rate VAE: Train Once, Get the Full Rate-Distortion Curve

• Published in: arXiv.org 2023-08
• Main Authors: Bae, Juhan, Zhang, Michael R, Ruan, Michael, Wang, Eric, Hasegawa, So, Ba, Jimmy, Grosse, Roger
• Format: Article
• Language: English
• Subjects: Distortion; Divergence; Learning; Parameters; Response functions; Training
• ISSN: 2331-8422

Variational autoencoders (VAEs) are powerful tools for learning latent representations of data used in a wide range of applications. In practice, VAEs usually require multiple training rounds to choose the amount of information the latent variable should retain. This trade-off between the reconstruction error (distortion) and the KL divergence (rate) is typically parameterized by a hyperparameter \(\beta\). In this paper, we introduce Multi-Rate VAE (MR-VAE), a computationally efficient framework for learning optimal parameters corresponding to various \(\beta\) in a single training run. The key idea is to explicitly formulate a response function that maps \(\beta\) to the optimal parameters using hypernetworks. MR-VAEs construct a compact response hypernetwork where the pre-activations are conditionally gated based on \(\beta\). We justify the proposed architecture by analyzing linear VAEs and showing that it can represent response functions exactly for linear VAEs. With the learned hypernetwork, MR-VAEs can construct the rate-distortion curve without additional training and can be deployed with significantly less hyperparameter tuning. Empirically, our approach is competitive and often exceeds the performance of multiple \(\beta\)-VAEs training with minimal computation and memory overheads.