[Paper Review] Increasing Expressivity of a Hyperspherical VAE

Tim R. Davidson, Jakub M. Tomczak+1 more|arXiv (Cornell University)|Oct 7, 2019

Gaussian Processes and Bayesian Inference23 references4 citations

TL;DR

This paper proposes a product-space hyperspherical VAE that enhances expressivity by decomposing a single high-dimensional hypersphere into multiple independent sub-hyperspheres, each with its own concentration parameter. By replacing the scalar concentration parameter of the von Mises-Fisher distribution with per-subspace parameters, the model achieves significantly improved likelihoods and disentanglement on image datasets like MNIST, Caltech, and Omniglot, outperforming standard hyperspherical and Gaussian VAEs.

ABSTRACT

Learning suitable latent representations for observed, high-dimensional data is an important research topic underlying many recent advances in machine learning. While traditionally the Gaussian normal distribution has been the go-to latent parameterization, recently a variety of works have successfully proposed the use of manifold-valued latents. In one such work (Davidson et al., 2018), the authors empirically show the potential benefits of using a hyperspherical von Mises-Fisher (vMF) distribution in low dimensionality. However, due to the unique distributional form of the vMF, expressivity in higher dimensional space is limited as a result of its scalar concentration parameter leading to a 'hyperspherical bottleneck'. In this work we propose to extend the usability of hyperspherical parameterizations to higher dimensions using a product-space instead, showing improved results on a selection of image datasets.

Motivation & Objective

• To address the limited expressivity of the von Mises-Fisher (vMF) distribution in high-dimensional latent spaces due to its single scalar concentration parameter.
• To improve the representational capacity of hyperspherical VAEs by decomposing the latent space into a Cartesian product of lower-dimensional hyperspheres.
• To enable greater flexibility and disentanglement in generative modeling by introducing independent concentration parameters for each sub-hypersphere.
• To empirically validate the method on benchmark image datasets, demonstrating improved likelihood and reconstruction performance.
• To show that the product-space construction mitigates the 'hyperspherical bottleneck' without sacrificing model simplicity.

Proposed method

• The method decomposes the latent space $\mathcal{S}_M$ into a product of $k+1$ independent sub-hyperspheres $\mathcal{S}_{M_0} \times \cdots \times \mathcal{S}_{M_k}$, where $\sum M_i = M$.
• Each sub-hypersphere is modeled using a von Mises-Fisher distribution with its own concentration parameter $\kappa_i$, increasing per-dimension flexibility.
• The joint prior and posterior are factorized as $p(\mathbf{z}) = \prod_i p(\mathbf{z}_i)$ and $q(\mathbf{z}|\mathbf{x}) = \prod_i q(\mathbf{z}_i|\mathbf{x})$, assuming conditional independence.
• The KL divergence decomposes as $\mathrm{KL}(q(\mathbf{z}|\mathbf{x})||p(\mathbf{z})) = \sum_i \mathrm{KL}(q(\mathbf{z}_i|\mathbf{x})||p(\mathbf{z}_i))$, enabling efficient optimization.
• The model is trained via variational inference with reparameterization gradients, using the vMF's tractable normalization constant and Bessel functions.
• The method is evaluated on MNIST, Caltech-101, and Omniglot using log-likelihood, reconstruction error, and KL divergence as metrics.

Experimental results

Research questions

• RQ1 Can a product-space construction of hyperspherical components improve the expressivity of a VAE beyond a single vMF distribution?
• RQ2 Does introducing per-subspace concentration parameters mitigate the 'hyperspherical bottleneck' in high-dimensional latent spaces?
• RQ3 How does the product-space vMF VAE compare to standard Gaussian and single-hypersphere vMF VAEs in terms of likelihood and disentanglement?
• RQ4 Can sub-hyperspheres learn semantically meaningful, disentangled factors of variation in image data?
• RQ5 What is the optimal trade-off between number of sub-hyperspheres and concentration parameters in terms of model performance?

Key findings

• The product-space vMF VAE achieved a log-likelihood of -92.50 on MNIST, outperforming the best single $\mathcal{S}_m$-VAE with -96.32.
• On Caltech-101, the model reached a log-likelihood of -139.84, significantly better than the single $\mathcal{S}_m$-VAE's -143.49.
• On Omniglot, the product-space model achieved -112.58 log-likelihood, surpassing the single $\mathcal{S}_m$-VAE's -113.83.
• The model demonstrated improved disentanglement, with specific sub-hyperspheres controlling semantic attributes like stroke thickness, as shown in interpolation visualizations.
• The method achieved lower reconstruction error and more stable training compared to standard vMF and Gaussian VAEs, particularly in high-dimensional settings.
• The ablation study confirmed that increasing the number of sub-hyperspheres and concentration parameters improved performance, with $[\mathcal{S}_7]_{\times 5} \times \mathcal{S}_5$ yielding strong results on all datasets.