The Power Spherical distribution

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