Multivariate Variational Autoencoder

TL;DR

The paper introduces the Multivariate Variational Autoencoder (MVAE), a full-covariance extension of VAEs that models correlated latent factors, improving calibration, disentanglement, and reconstruction quality across multiple datasets.

Contribution

It proposes a tractable multivariate Gaussian posterior for VAEs, enabling correlated latent variables while maintaining simple optimization and evaluation procedures.

Findings

MVAE outperforms diagonal VAEs in calibration and clustering metrics.

Provides smoother, more coherent latent traversals.

Achieves better or comparable reconstruction quality across datasets.

Abstract

Learning latent representations that are simultaneously expressive, geometrically well-structured, and reliably calibrated remains a central challenge for Variational Autoencoders (VAEs). Standard VAEs typically assume a diagonal Gaussian posterior, which simplifies optimization but rules out correlated uncertainty and often yields entangled or redundant latent dimensions. We introduce the Multivariate Variational Autoencoder (MVAE), a tractable full-covariance extension of the VAE that augments the encoder with sample-specific diagonal scales and a global coupling matrix. This induces a multivariate Gaussian posterior of the form $N(\mu_\phi(x), C \operatorname{diag}(\sigma_\phi^2(x)) C^\top)$, enabling correlated latent factors while preserving a closed-form KL divergence and a simple reparameterization path. Beyond likelihood, we propose a multi-criterion evaluation protocol that jointly assesses reconstruction quality (MSE, ELBO), downstream discrimination (linear probes), probabilistic calibration (NLL, Brier, ECE), and unsupervised structure (NMI, ARI). Across Larochelle-style MNIST variants, Fashion-MNIST, and CIFAR-10/100, MVAE consistently matches or outperforms diagonal-covariance VAEs of comparable capacity, with particularly notable gains in calibration and clustering metrics at both low and high latent dimensions. Qualitative analyses further show smoother, more semantically coherent latent traversals and sharper reconstructions. All code, dataset splits, and evaluation utilities are released to facilitate reproducible comparison and future extensions of multivariate posterior models.