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
