Denoising Multi-Beta VAE: Representation Learning for Disentanglement and Generation
This work addresses a core challenge in generative modeling for researchers and practitioners by enabling better balance between interpretability and quality, though it is incremental as it builds on existing β-VAE and diffusion methods.
The paper tackles the trade-off between disentangled latent representations and generation quality in VAEs by proposing a framework that learns multiple latent representations using a range of β values and a diffusion model to denoise towards more informative representations, achieving improved disentanglement and sharp reconstructions with concrete gains in metrics like reconstruction error and disentanglement scores.
Disentangled and interpretable latent representations in generative models typically come at the cost of generation quality. The $β$-VAE framework introduces a hyperparameter $β$ to balance disentanglement and reconstruction quality, where setting $β> 1$ introduces an information bottleneck that favors disentanglement over sharp, accurate reconstructions. To address this trade-off, we propose a novel generative modeling framework that leverages a range of $β$ values to learn multiple corresponding latent representations. First, we obtain a slew of representations by training a single variational autoencoder (VAE), with a new loss function that controls the information retained in each latent representation such that the higher $β$ value prioritize disentanglement over reconstruction fidelity. We then, introduce a non-linear diffusion model that smoothly transitions latent representations corresponding to different $β$ values. This model denoises towards less disentangled and more informative representations, ultimately leading to (almost) lossless representations, enabling sharp reconstructions. Furthermore, our model supports sample generation without input images, functioning as a standalone generative model. We evaluate our framework in terms of both disentanglement and generation quality. Additionally, we observe smooth transitions in the latent spaces with respect to changes in $β$, facilitating consistent manipulation of generated outputs.