Multi-Rate VAE: Train Once, Get the Full Rate-Distortion Curve
This reduces hyperparameter tuning and computational cost for practitioners using VAEs in applications like representation learning.
The paper tackles the problem of needing multiple training rounds to choose the trade-off between reconstruction error and KL divergence in VAEs, introducing Multi-Rate VAE (MR-VAE) to learn optimal parameters for various β in a single training run, achieving competitive or better performance than multiple β-VAEs with minimal overhead.
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 $β$. In this paper, we introduce Multi-Rate VAE (MR-VAE), a computationally efficient framework for learning optimal parameters corresponding to various $β$ in a single training run. The key idea is to explicitly formulate a response function that maps $β$ to the optimal parameters using hypernetworks. MR-VAEs construct a compact response hypernetwork where the pre-activations are conditionally gated based on $β$. 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 $β$-VAEs training with minimal computation and memory overheads.