A better unit vector

Suppose we want to sample x \in \mathbb{S}^n \subset \mathbb{R}^{n+1} with the spherical embedding approach, i.e. sampling y \in \mathbb{R}^{n+1}, where x = \frac{y}{\lVert y \rVert}. As @betanalpha explained, the log-density “correction” of -\lVert y \rVert^2/2 that Stan uses puts an implicit density on r=\lVert y \rVert.

Specifically, we can think of the transformation from x to y as two steps:

  1. augment the density \pi_x(x) with a prior \pi_r(r) for r>0 to get a new joint density \pi_x(x)\pi_r( r).
  2. bijectively map (x, r) \in \mathbb{S}^n \times R_{>0} to \mathbb{R}^{n+1} \backslash \{0\} with y = r x, and use the Jacobian correction -n \log r to get the log-density \log\pi_y(y) = \log\pi_x(y/\lVert y \rVert) + \log\pi_r(\lVert y \rVert) - n\log \lVert y \rVert.

We are free to choose any continuous proper density for \pi_r(r). Stan implicitly chooses the Chi distribution with n+1 degrees of freedom, which has the log-density \log\pi_r(r) = n \log(r) -\frac{1}{2} r^2. The n \log r term, which attracts draws to 0 in one expression and repulses in the other, perfectly cancel, so that we are left with \log \pi_y(y) = -\frac{1}{2} \lVert y\rVert^2.

We still have a singularity at y=0 that we need to avoid, and as noted above, when \pi_x(x) is concentrated, then \pi_y(y) has a wedge geometry that is challenging to sample. Due to concentration of measure, for large n, draws near y=0 should be rare, so this shouldn’t be as much of a problem. But for low n, this will manifest with divergences.

An alternative solution to those mentioned so far is to use \log\pi_r(r) = \frac{1}{2} r^2 + (n + a) \log r for a \ge 0, i.e. a Chi distribution with n+a+1 degrees of freedom. This repels y from y=0, and the degree of repulsion can be tuned by the user by increasing a.

Here we see the log-density for a von Mises distribution with concentration of 100 transformed to the latent space for various values of a:

Yet another alternative is to use \log \pi_r(r) = -\frac{1}{2} (r-m)^2 + n \log r for m \ge 0:

I’m not sure sure how this could be supported in Stan.