[Paper Review] The Power Spherical distribution

Nicola De Cao, Wilker Aziz|arXiv (Cornell University)|Jun 8, 2020

Gaussian Processes and Bayesian Inference29 references17 citations

TL;DR

The paper proposes the Power Spherical distribution, a numerically stable, reparameterizable alternative to the von Mises-Fisher (vMF) distribution for hyperspherical data. By designing a univariate marginal with an analytical inverse CDF, it enables fast, rejection-free sampling and exact gradient computation, outperforming vMF in sampling speed and stability while maintaining comparable performance in variational autoencoders.

ABSTRACT

There is a growing interest in probabilistic models defined in hyper-spherical spaces, be it to accommodate observed data or latent structure. The von Mises-Fisher (vMF) distribution, often regarded as the Normal distribution on the hyper-sphere, is a standard modeling choice: it is an exponential family and thus enjoys important statistical results, for example, known Kullback-Leibler (KL) divergence from other vMF distributions. Sampling from a vMF distribution, however, requires a rejection sampling procedure which besides being slow poses difficulties in the context of stochastic backpropagation via the reparameterization trick. Moreover, this procedure is numerically unstable for certain vMFs, e.g., those with high concentration and/or in high dimensions. We propose a novel distribution, the Power Spherical distribution, which retains some of the important aspects of the vMF (e.g., support on the hyper-sphere, symmetry about its mean direction parameter, known KL from other vMF distributions) while addressing its main drawbacks (i.e., scalability and numerical stability). We demonstrate the stability of Power Spherical distributions with a numerical experiment and further apply it to a variational auto-encoder trained on MNIST. Code at: https://github.com/nicola-decao/power_spherical

Motivation & Objective

• Address the numerical instability and inefficiency of rejection sampling in the von Mises-Fisher (vMF) distribution for high-dimensional or high-concentration settings.
• Enable exact reparameterization for stochastic backpropagation by eliminating the need for high-variance correction terms.
• Develop a distribution on the hypersphere that supports efficient, stable sampling without relying on rejection sampling.
• Ensure compatibility with deep learning frameworks by enabling differentiable sampling and gradient computation.
• Demonstrate practical utility by replacing vMF in a variational autoencoder on MNIST with improved training speed and equivalent performance.

Proposed method

• Leverage the tangent-normal decomposition of hyperspherical vectors to separate the distribution into a radial component (t) and a tangent component (v), reducing the problem to modeling the marginal distribution of t.
• Define the Power Spherical distribution via a power-law form for the marginal density of t, ensuring an analytical inverse cumulative distribution function (CDF).
• Construct the full hyperspherical distribution by combining the analytically invertible marginal with a uniform distribution over the tangent space.
• Enable reparameterization by sampling t directly from the inverse CDF, allowing backpropagation through the sampling process without correction terms.
• Derive closed-form expressions for key statistical properties, including mean, mode, variance, entropy, and KL divergence from vMF and uniform distributions.
• Validate the method through numerical experiments on sampling stability and performance benchmarks in a variational autoencoder setting.

Experimental results

Research questions

• RQ1 Can a hyperspherical distribution be designed to enable fast, stable, and reparameterizable sampling without rejection sampling?
• RQ2 Does a distribution with an analytically invertible CDF on the hypersphere maintain sufficient expressivity for deep generative modeling?
• RQ3 How does the proposed distribution compare to the vMF in terms of numerical stability under high concentration and high dimensionality?
• RQ4 Can the Power Spherical distribution be used as a direct replacement for vMF in variational autoencoders without performance degradation?
• RQ5 What is the impact of the new sampling method on training speed and gradient quality in deep learning models?

Key findings

• The Power Spherical distribution enables sampling that is at least 6× faster than the vMF and up to 20× faster for high-concentration cases, with constant sampling time regardless of concentration.
• The Power Spherical distribution is numerically stable across high-dimensional spaces (e.g., d=64) and high concentration parameters, unlike the vMF which suffers from numerical instability.
• Sampling from the Power Spherical distribution requires no rejection sampling, enabling exact reparameterization without high-variance correction terms.
• In a variational autoencoder on MNIST, the Power Spherical achieved comparable log-likelihood and ELBO scores to the vMF across all latent dimensions (d=5 to d=40).
• Despite equivalent performance, training with the Power Spherical was more than 2× faster than with the vMF, demonstrating improved computational efficiency.
• The results confirm that the Power Spherical is a viable, scalable, and stable alternative to the vMF for deep generative models on hyperspherical latent spaces.