# Training β-VAE by Aggregating a Learned Gaussian Posterior with a Decoupled Decoder

arXiv: https://arxiv.org/abs/2209.14783 · HTML: https://mlsanity.com/paper/2209.14783/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 21 of its 36 references that are in the map.
- First posted on arXiv 29 Sep 2022.

## Builds on (most cited first, 21 of 21)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Generative Adversarial Networks](https://mlsanity.com/paper/2203.00667/) (2022)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [Auto-Encoding Variational Bayes](https://mlsanity.com/paper/1312.6114/) (2013)
- [Generating Diverse High-Fidelity Images with VQ-VAE-2](https://mlsanity.com/paper/1906.00446/) (2019)
- [Deep Variational Information Bottleneck](https://mlsanity.com/paper/1612.00410/) (2016)
- [InfoGAN: Interpretable Representation Learning by Information Maximizing Generative Adversarial Nets](https://mlsanity.com/paper/1606.03657/) (2016)
- [Disentangling by Factorising](https://mlsanity.com/paper/1802.05983/) (2018)
- [Generating Sentences from a Continuous Space](https://mlsanity.com/paper/1511.06349/) (2015)
- [Adversarial Autoencoders](https://mlsanity.com/paper/1511.05644/) (2015)
- [Understanding disentangling in β-VAE](https://mlsanity.com/paper/1804.03599/) (2018)
- [Deep Unsupervised Clustering with Gaussian Mixture Variational Autoencoders](https://mlsanity.com/paper/1611.02648/) (2016)
- [Diagnosing and Enhancing VAE Models](https://mlsanity.com/paper/1903.05789/) (2019)
- [Avoiding Latent Variable Collapse With Generative Skip Models](https://mlsanity.com/paper/1807.04863/) (2018)
- [Denoising Criterion for Variational Auto-Encoding Framework](https://mlsanity.com/paper/1511.06406/) (2015)
- [Balancing reconstruction error and Kullback-Leibler divergence in Variational Autoencoders](https://mlsanity.com/paper/2002.07514/) (2020)
- [Learning Representations by Maximizing Mutual Information in Variational Autoencoders](https://mlsanity.com/paper/1912.13361/) (2019)
- [A Baseline Approach for AutoImplant: the MICCAI 2020 Cranial Implant Design Challenge](https://mlsanity.com/paper/2006.12449/) (2020)
- [Review of Disentanglement Approaches for Medical Applications -- Towards Solving the Gordian Knot of Generative Models in Healthcare](https://mlsanity.com/paper/2203.11132/) (2022)
- [VMI-VAE: Variational Mutual Information Maximization Framework for VAE With Discrete and Continuous Priors](https://mlsanity.com/paper/2005.13953/) (2020)
- [Structured Uncertainty in the Observation Space of Variational Autoencoders](https://mlsanity.com/paper/2205.12533/) (2022)

## Built on by (most cited first, 2 of 2)

- [MedShapeNet -- A Large-Scale Dataset of 3D Medical Shapes for Computer Vision](https://mlsanity.com/paper/2308.16139/) (2023)
- [Open-Source Skull Reconstruction with MONAI](https://mlsanity.com/paper/2211.14051/) (2022)

## Abstract (arXiv)

The reconstruction loss and the Kullback-Leibler divergence (KLD) loss in a variational autoencoder (VAE) often play antagonistic roles, and tuning the weight of the KLD loss in β-VAE to achieve a balance between the two losses is a tricky and dataset-specific task. As a result, current practices in VAE training often result in a trade-off between the reconstruction fidelity and the continuity/disentanglement of the latent space, if the weight β is not carefully tuned. In this paper, we present intuitions and a careful analysis of the antagonistic mechanism of the two losses, and propose, based on the insights, a simple yet effective two-stage method for training a VAE. Specifically, the method aggregates a learned Gaussian posterior z ~ qθ (z|x) with a decoder decoupled from the KLD loss, which is trained to learn a new conditional distribution pφ (x|z) of the input data x. Experimentally, we show that the aggregated VAE maximally satisfies the Gaussian assumption about the latent space, while still achieves a reconstruction error comparable to when the latent space is only loosely regularized by N(0,I). The proposed approach does not require hyperparameter (i.e., the KLD weight β) tuning given a specific dataset as required in common VAE training practices. We evaluate the method using a medical dataset intended for 3D skull reconstruction and shape completion, and the results indicate promising generative capabilities of the VAE trained using the proposed method. Besides, through guided manipulation of the latent variables, we establish a connection between existing autoencoder (AE)-based approaches and generative approaches, such as VAE, for the shape completion problem. Codes and pre-trained weights are available at https://github.com/Jianningli/skullVAE