Model Selection for Bayesian Autoencoders

TL;DR

This paper introduces a new model selection method for Bayesian autoencoders using prior hyper-parameter optimization based on distributional sliced-Wasserstein distance, improving performance especially in small-data regimes.

Contribution

It proposes a novel approach to model selection for BAEs by optimizing DSWD, enabling effective high-dimensional data handling and providing a competitive alternative to variational autoencoders.

Findings

Achieves state-of-the-art results in small-data regimes.

Outperforms multiple baselines in unsupervised learning tasks.

Provides a scalable method for model selection in high-dimensional settings.

Abstract

We develop a novel method for carrying out model selection for Bayesian autoencoders (BAEs) by means of prior hyper-parameter optimization. Inspired by the common practice of type-II maximum likelihood optimization and its equivalence to Kullback-Leibler divergence minimization, we propose to optimize the distributional sliced-Wasserstein distance (DSWD) between the output of the autoencoder and the empirical data distribution. The advantages of this formulation are that we can estimate the DSWD based on samples and handle high-dimensional problems. We carry out posterior estimation of the BAE parameters via stochastic gradient Hamiltonian Monte Carlo and turn our BAE into a generative model by fitting a flexible Dirichlet mixture model in the latent space. Consequently, we obtain a powerful alternative to variational autoencoders, which are the preferred choice in modern applications of autoencoders for representation learning with uncertainty. We evaluate our approach qualitatively and quantitatively using a vast experimental campaign on a number of unsupervised learning tasks and show that, in small-data regimes where priors matter, our approach provides state-of-the-art results, outperforming multiple competitive baselines.