TY - GEN
T1 - Training β -VAE by Aggregating a Learned Gaussian Posterior with a Decoupled Decoder
AU - Li, Jianning
AU - Fragemann, Jana
AU - Ahmadi, Seyed Ahmad
AU - Kleesiek, Jens
AU - Egger, Jan
N1 - Publisher Copyright:
© 2023, The Author(s), under exclusive license to Springer Nature Switzerland AG.
PY - 2023
Y1 - 2023
N2 - 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.
AB - 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.
KW - Disentanglement
KW - Latent representation
KW - Shape completion
KW - Skull reconstruction
KW - VAE
UR - http://www.scopus.com/inward/record.url?scp=85151059789&partnerID=8YFLogxK
U2 - 10.1007/978-3-031-25046-0_7
DO - 10.1007/978-3-031-25046-0_7
M3 - Conference paper
AN - SCOPUS:85151059789
SN - 9783031250453
T3 - Lecture Notes in Computer Science (including subseries Lecture Notes in Artificial Intelligence and Lecture Notes in Bioinformatics)
SP - 70
EP - 92
BT - Medical Applications with Disentanglements - First MICCAI Workshop, MAD 2022, Held in Conjunction with MICCAI 2022, Proceedings
A2 - Fragemann, Jana
A2 - Li, Jianning
A2 - Egger, Jan
A2 - Liu, Xiao
A2 - Tsaftaris, Sotirios A.
A2 - Kleesiek, Jens
PB - Springer Science and Business Media Deutschland GmbH
T2 - 1st MICCAI Workshop on Medical Applications with Disentanglements, held in conjunction with MICCAI 2022
Y2 - 22 September 2022 through 22 September 2022
ER -