Sampling Constrained Probability Distributions Using Spherical Augmentation

Abstract

Statistical models with constrained probability distributions are abundant in machine learning. Some examples include regression models with norm constraints (e.g., Lasso), probit, many copula models, and latent Dirichlet allocation (LDA). Bayesian inference involving probability distributions confined to constrained domains could be quite challenging for commonly used sampling algorithms. In this work, we propose a novel augmentation technique that handles a wide range of constraints by mapping the constrained domain to a sphere in the augmented space. By moving freely on the surface of this sphere, sampling algorithms handle constraints implicitly and generate proposals that remain within boundaries when mapped back to the original space. Our proposed method, called Spherical Augmentation, provides a mathematically natural and computationally efficient framework for sampling from constrained probability distributions . We show the advantages of our method over state-of-the-art sampling algorithms, such as exact Hamiltonian Monte Carlo, using several examples including truncated Gaussian distributions, Bayesian Lasso, Bayesian bridge regression, reconstruction of quantized stationary Gaussian process, and LDA for topic modeling.

Appendix

Spherical Geometry

We first discuss the geometry of the D-dimensional sphere $\mathscr {S}^D:=\{\tilde{\varvec{\theta }} \in \mathbb R^{D+1}: \Vert \tilde{\varvec{\theta }}\Vert _2= 1\}\hookrightarrow \mathbb R^{D+1}$ under different coordinate systems, namely, Cartesian coordinates and the spherical coordinates. Since $\mathscr {S}^D$ can be embedded (injectively and differentiably mapped to) in $\mathbb R^{D+1}$, we first introduce the concept of ‘induced metric’.

Definition 2.1

(induced metric) If $\mathscr {D}^{d}$ can be embedded to $\mathscr {M}^{m}$ $(m>d)$ by $f: U\subset \mathscr {D}\hookrightarrow \mathscr {M}$, then one can define the induced metric, $g_{\mathscr {D}}$, on $T\mathscr {D}$ through the metric $g_{\mathscr {M}}$ defined on $T\mathscr {M}$:

$$\begin{aligned} g_{\mathscr {D}}(\varvec{\theta })(\mathbf{u},\mathbf{{v}})=g_{\mathscr {M}}(f(\varvec{\theta }))(df_{\varvec{\theta }}(\mathbf{u}),df_{\varvec{\theta }}(\mathbf{{v}})),\quad \mathbf{u},\mathbf{{v}}\in T_{\varvec{\theta }}\mathscr {D} \end{aligned}$$

(2.58)

Remark 2.2

For any $f: U\subset \mathscr {S}^D\hookrightarrow \mathbb R^{D+1}$, we can define the induced metric through dot product on $\mathbb R^{D+1}$. More specifically,

$$\begin{aligned} g_{\mathscr {S}}(\mathbf{u},\mathbf{{v}}) = {[(Df)\mathbf{u}]}^{\mathsf {T}} (Df)\mathbf{{v}} = \mathbf{u}^{\mathsf {T}} [{(Df)}^{\mathsf {T}} (Df)] \mathbf{{v}} \end{aligned}$$

(2.59)

where $(Df)_{(D+1)\times D}$ is the Jacobian matrix of the mapping f. A Metric induced from dot product on Euclidean space is called a “canonical metric ”. This observation leads to the following simple fact that lays down the foundation of Spherical HMC.

Proposition 2.7.1

(Energy invariance ) Kinetic energy $\frac{1}{2}\langle \mathbf{{v}},\mathbf{{v}}\rangle _{\mathbf{G}(\varvec{\theta })}$ is invariant to the choice of coordinate systems.

Proof

For any $\mathbf{{v}}\in T_{\varvec{\theta }}\mathscr {D}$, suppose $\varvec{\theta }(t)$ such that $\dot{\varvec{\theta }}(0)=\mathbf{{v}}$. Denote the pushforward of $\mathbf{{v}}$ by embedding map $f:\mathscr {D}\rightarrow \mathscr {M}$ as $\tilde{\mathbf{{v}}}:=f_*(\mathbf{v})=\frac{d}{dt}(f\circ \varvec{\theta })(0)$. Then we have

$$\begin{aligned} \frac{1}{2}\langle \mathbf{{v}},\mathbf{{v}}\rangle _{\mathbf{G}(\varvec{\theta })} = \frac{1}{2} g_{\mathscr {M}}(f(\varvec{\theta }))(\tilde{\mathbf{{v}}},\tilde{\mathbf{{v}}}) \end{aligned}$$

(2.60)

That is, regardless of the form of the energy under a coordinate system, its value is the same as the one in the embedded manifold. In particular, when $\mathscr {M}=\mathbb R^{D+1}$, the right hand side simplifies to $\frac{1}{2} \Vert \tilde{\mathbf{{v}}}\Vert _2^2$. $\square$

Canonical Metric in Cartesian Coordinates

Now consider the D-dimensional ball $\mathscr {B}_\mathbf{0}^D(1):=\{\varvec{\theta }\in \mathbb R^D: \Vert \varvec{\theta }\Vert _2\le 1\}$. Here, $\{\varvec{\theta }, \mathscr {B}_\mathbf{0}^D(1)\}$ can be viewed as the Cartesian coordinate system for ${\mathscr {S}}^D$. The coordinate mapping $T_{\mathscr {B}\rightarrow \mathscr {S}_+}:\varvec{\theta }\mapsto \tilde{\varvec{\theta }}=(\varvec{\theta },\theta _{D+1})$ in (2.5) can be viewed as the embedding map into $\mathbb R^{D+1}$, and the Jacobian matrix of $T_{\mathscr {B}\rightarrow \mathscr {S}_+}$ is $dT_{\mathscr {B}\rightarrow \mathscr {S}_+}=\frac{d\tilde{\varvec{\theta }}}{d{\varvec{\theta }}^{\mathsf {T}}}=\begin{bmatrix}{} \mathbf{I}_{D}\\tp{\varvec{\theta }}/\theta _{D+1}\end{bmatrix}$. Therefore the canonical metric of ${\mathscr {S}}^D$ in Cartesian coordinates, $\mathbf{G}_{\mathscr {S}_c}(\varvec{\theta })$, is

$$\begin{aligned} \mathbf{G}_{\mathscr {S}_c}(\varvec{\theta }) = {dT}^{\mathsf {T}}_{\mathscr {B}\rightarrow \mathscr {S}_+} dT_{\mathscr {B}\rightarrow \mathscr {S}_+} = \mathbf{I}_D + \frac{\varvec{\theta } {\varvec{\theta }}^{\mathsf {T}}}{\theta _{D+1}^2} = \mathbf{I}_D + \frac{\varvec{\theta } {\varvec{\theta }}^{\mathsf {T}}}{1-\Vert \varvec{\theta }\Vert _2^2} \end{aligned}$$

(2.61)

Another way to obtain the metric is through the first fundamental form $ds^{2}$ (i.e., squared infinitesimal length of a curve) for $\mathscr {S}^D$, which can be expressed in terms of the differential form $d\varvec{\theta }$ and the canonical metric $\mathbf{G}_{\mathscr {S}_c}(\varvec{\theta })$,

$$\begin{aligned} ds^{2} = \langle d\varvec{\theta }, d\varvec{\theta } \rangle _{\mathbf{G}_{\mathscr {S}_c}} = d{\varvec{\theta }}^{\mathsf {T}}{} \mathbf{G}_{\mathscr {S}_c}(\varvec{\theta })d\varvec{\theta } \end{aligned}$$

On the other hand, $ds^{2}$ can also be obtained as follows [59]:

$$\begin{aligned} ds^2 = \sum _{i=1}^D d\theta _i^2 + (d(\theta _{D+1}(\varvec{\theta })))^2 = d{\varvec{\theta }}^{\mathsf {T}} d\varvec{\theta } + \frac{({\varvec{\theta }}^{\mathsf {T}} d\varvec{\theta })^2}{1-\Vert \varvec{\theta }\Vert _2^2} = d{\varvec{\theta }}^{\mathsf {T}} [\mathbf{I} + \varvec{\theta } {\varvec{\theta }}^{\mathsf {T}}/\theta _{D+1}^2] d\varvec{\theta } \end{aligned}$$

Equating the above two quantities yields the form of the canonical metric $\mathbf{G}_{\mathscr {S}_c}(\varvec{\theta })$ as in Eq. (2.61). This viewpoint provides a natural way to explain the length of tangent vector. For any vector $\tilde{\mathbf{{v}}}=(\mathbf{{v}},v_{D+1})\in T_{\tilde{\varvec{\theta }}}{\mathscr {S}}^D=\{\tilde{\mathbf{{v}}}\in \mathbb R^{D+1}: {\tilde{\varvec{\theta }}}^{\mathsf {T}}\tilde{\mathbf{{v}}}=0\}$, one could think of $\mathbf{G}_{\mathscr {S}_c}(\varvec{\theta })$ as a mean to express the length of $\tilde{\mathbf{v}}$ in terms of $\mathbf{{v}}$,

$$\begin{aligned} \mathbf{{v}}^{\mathsf {T}}{} \mathbf{G}_{\mathscr {S}_c}(\varvec{\theta })\mathbf{{v}} = \Vert \mathbf{{v}}\Vert _2^2 + \frac{\mathbf{{v}}^{\mathsf {T}}\varvec{\theta } {\varvec{\theta }}^{\mathsf {T}} \mathbf{{v}}}{\theta _{D+1}^2} = \Vert \mathbf{{v}}\Vert _2^2 + \frac{ (-\theta _{D+1} v_{D+1})^2}{\theta _{D+1}^2} = \Vert \mathbf{{v}}\Vert _2^2 + v_{D+1}^2 = \Vert \tilde{\mathbf{{v}}}\Vert _2^2 \end{aligned}$$

(2.62)

This indeed verifies the energy invariance Proposition 2.7.1.

The following proposition provides the analytic forms of the determinant and the inverse of $\mathbf{G}_{\mathscr {S}_c}(\varvec{\theta })$.

Proposition 2.7.2

The determinant and the inverse of the canonical metric are as follows:

$$\begin{aligned} |\mathbf{G}_{\mathscr {S}_c}(\varvec{\theta })| = \theta _{D+1}^{-2}, \qquad \mathbf{G}_{\mathscr {S}_c}(\varvec{\theta })^{-1} = \mathbf{I}_D - \varvec{\theta } {\varvec{\theta }}^{\mathsf {T}} \end{aligned}$$

(2.63)

Proof

The determinant of the canonical metric $\mathbf{G}_{\mathscr {S}_c}(\varvec{\theta })$ is given by the matrix determinant lemma

$$\begin{aligned} |\mathbf{G}_{\mathscr {S}_c}(\varvec{\theta })| = \det \left[ \mathbf{I}_D + \frac{\varvec{\theta } {\varvec{\theta }}^{\mathsf {T}}}{\theta _{D+1}^2}\right] = 1+ \frac{{\varvec{\theta }}^{\mathsf {T}} \varvec{\theta }}{\theta _{D+1}^2} = \frac{1}{\theta _{D+1}^2} \end{aligned}$$

The inverse of $\mathbf{G}_{\mathscr {S}_c}(\varvec{\theta })$ is obtained by the Sherman-Morrison-Woodbury formula [29]

$$\begin{aligned} \mathbf{G}_{\mathscr {S}_c}(\varvec{\theta })^{-1} = \left[ \mathbf{I}_D + \frac{\varvec{\theta } {\varvec{\theta }}^{\mathsf {T}}}{\theta _{D+1}^2} \right] ^{-1} = \mathbf{I}_D - \frac{\varvec{\theta } {\varvec{\theta }}^{\mathsf {T}}/\theta _{D+1}^2}{1+{\varvec{\theta }}^{\mathsf {T}}\varvec{\theta }/\theta _{D+1}^2} = \mathbf{I}_D - \varvec{\theta } {\varvec{\theta }}^{\mathsf {T}} \end{aligned}$$

$\square$

Corollary 2.1

The volume adjustment of changing measure in (2.6) is

$$\begin{aligned} \left| \frac{d\varvec{\theta }_{\mathscr {B}}}{d\varvec{\theta }_{\mathscr {S}_c}}\right| =|\mathbf{G}_{\mathscr {S}_c}(\varvec{\theta })|^{-\frac{1}{2}}=|\theta _{D+1}| \end{aligned}$$

(2.64)

Proof

Canonical measure can be defined through the Riesz representation theorem by using a positive linear functional on the space $C_0(\mathscr {S}^D)$ of compactly supported continuous functions on $\mathscr {S}^D$ [21, 59]. More precisely, there is a unique positive Borel measure $\mu _c$ such that for (any) coordinate chart $(\mathscr {B}_\mathbf{0}^D(1),T_{\mathscr {B}\rightarrow \mathscr {S}_+})$,

$$\begin{aligned} \int _{\mathscr {S}_+^D} f(\tilde{\varvec{\theta }}) d\varvec{\theta }_{\mathscr {S}_c} = \int _{\mathscr {B}_\mathbf{0}^D(1)} f(\varvec{\theta }) \sqrt{|\mathbf{G}_{\mathscr {S}_c}(\varvec{\theta })|} d\varvec{\theta }_{\mathscr {B}} \end{aligned}$$

where $\mu _c=d\varvec{\theta }_{\mathscr {S}_c}$, and $d\varvec{\theta }_{\mathscr {B}}$ is the Euclidean measure. Therefore we have

$$\begin{aligned} \left| \frac{d\varvec{\theta }_{\mathscr {S}_c}}{d\varvec{\theta }_{\mathscr {B}}}\right| = |\mathbf{G}_{\mathscr {S}_c}(\varvec{\theta })|^{\frac{1}{2}} = |\theta _{D+1}|^{-1} \end{aligned}$$

Alternatively, $\left| \frac{d\varvec{\theta }_{\mathscr {B}}}{d\varvec{\theta }_{\mathscr {S}_c}}\right| =|\theta _{D+1}|$. $\square$

Geodesic on a Sphere in Cartesian Coordinates

To find the geodesic on a sphere , we need to solve the following equations:

$$\begin{aligned} \dot{\varvec{\theta }}&= \mathbf{{v}} \end{aligned}$$

(2.65)

$$\begin{aligned} \dot{\mathbf{{v}}}&= -\mathbf{{v}}^{\mathsf {T}}\varvec{\varGamma }_{\mathscr {S}_c}(\varvec{\theta })\mathbf{{v}} \end{aligned}$$

(2.66)

for which we need to calculate the Christoffel symbols, $\varvec{\varGamma }_{\mathscr {S}_c}(\varvec{\theta })$, first. Note that the (i, j)-th element of $\mathbf{G}_{\mathscr {S}_c}$ is $g_{ij} = \delta _{ij} + \theta _i\theta _j/\theta _{D+1}^2$, and the (i, j, k)-th element of $d\mathbf{G}_{\mathscr {S}_c}$ is $g_{ij,k} = (\delta _{ik}\theta _j + \theta _i\delta _{jk})/\theta _{D+1}^2 + 2\theta _i\theta _j\theta _k/\theta _{D+1}^4$. Therefore

$$\begin{aligned} \begin{aligned} \varGamma _{ij}^k&= \frac{1}{2}g^{kl}[g_{lj,i}+g_{il,j}-g_{ij,l}]\\&= (\delta ^{kl}-\theta ^k\theta ^l)\theta _l/\theta _{D+1}^2[\delta _{ij}+\theta _i\theta _j/\theta _{D+1}^2]\\&= \theta _k[\delta _{ij}+\theta _i\theta _j/\theta _{D+1}^2] = [\mathbf{G}_{\mathscr {S}_c}(\varvec{\theta })\otimes \varvec{\theta }]_{ijk} \end{aligned} \end{aligned}$$

Using these results, we can write Eq. (2.66) as $\dot{\mathbf{{v}}} = -\mathbf{{v}}^{\mathsf {T}} \mathbf{G}_{\mathscr {S}_c}(\varvec{\theta })\mathbf{{v}}\varvec{\theta } = -\Vert \tilde{\mathbf{v}}\Vert _2^2\varvec{\theta }$. Further, we have

Therefore, we can rewrite the geodesic equations (2.65), (2.66) with augmented components as

$$\begin{aligned} \dot{\tilde{\varvec{\theta }}}&= \tilde{\mathbf{{v}}} \end{aligned}$$

(2.67)

$$\begin{aligned} \dot{\tilde{\mathbf{{v}}}}&= -\Vert \tilde{\mathbf{{v}}}\Vert _2^2 \tilde{\varvec{\theta }} \end{aligned}$$

(2.68)

Multiplying both sides of Eq. (2.68) by ${\tilde{\mathbf{{v}}}}^{\mathsf {T}}$ to obtain $\frac{d}{dt}\Vert \tilde{\mathbf{{v}}}\Vert _2^2=0$, we can solve the above system of differential equations as follows:

$$\begin{aligned} \tilde{\varvec{\theta }}(t)&= \tilde{\varvec{\theta }}(0) \cos (\Vert \tilde{\mathbf{{v}}}(0)\Vert _{2} t) + \frac{\tilde{\mathbf{{v}}}(0)}{\Vert \tilde{\mathbf{{v}}}(0)\Vert _{2}} \sin (\Vert \tilde{\mathbf{{v}}}(0)\Vert _{2} t)\\ \tilde{\mathbf{{v}}}(t)&= -\tilde{\varvec{\theta }}(0) \Vert \tilde{\mathbf{{v}}}(0)\Vert _{2} \sin (\Vert \tilde{\mathbf{{v}}}(0)\Vert _{2} t) + \tilde{\mathbf{{v}}}(0) \cos (\Vert \tilde{\mathbf{{v}}}(0)\Vert _{2} t) \end{aligned}$$

Round Metric in the Spherical Coordinates

Consider the D-dimensional hyperrectangle $\mathscr {R}_\mathbf{0}^D:=[0,\pi ]^{D-1}\times [0,2\pi )$ and the corresponding spherical coordinate system, $\{\varvec{\theta }, \mathscr {R}_\mathbf{0}^D\}$, for $\mathscr {S}^D$. The coordinate mapping $T_{\mathscr {R}_\mathbf{0}\rightarrow \mathscr {S}}: \varvec{\theta }\mapsto \mathbf{x},\; x_d = \cos (\theta _d)\prod _{i=1}^{d-1}\sin (\theta _i),\, d=1,\ldots , D+1$, ($\theta _{D+1}=0$) can be viewed as the embedding map into $\mathbb R^{D+1}$, and the Jacobian matrix of $T_{\mathscr {R}_\mathbf{0}\rightarrow \mathscr {S}}$ is $\frac{d\mathbf{x}}{d{\varvec{\theta }}^{\mathsf {T}}}$ with the (d, j)-th element $[-\tan (\theta _d)\delta _{dj}+\cot (\theta _j)I(j<d)]x_d$. The induced metric of ${\mathscr {S}}^D$ in the spherical coordinates is called round metric , denoted as $\mathbf{G}_{\mathscr {S}_r}(\varvec{\theta })$, whose (i, j)-th element is as follows

$$\begin{aligned} \begin{aligned}&\mathbf{G}_{\mathscr {S}_r}(\varvec{\theta })_{ij}\\&= \sum _{d=1}^{D+1} [-\tan (\theta _d)\delta _{di}+\cot (\theta _j)I(i<d)][-\tan (\theta _d)\delta _{dj}+\cot (\theta _j)I(j<d)]x_d^2\\&={\left\{ \begin{array}{ll} -\tan (\theta _j)\cot (\theta _i)x_j^2 + \cot (\theta _i)\cot (\theta _j)\sum _{d>j} x_d^2 =0,&{} i<j\\ \tan ^2(\theta _i)x_i^2 + \cot ^2(\theta _i) \sum _{d>i} x_d^2 = (\tan ^2(\theta _i)+1)x_i^2 = \prod _{i=1}^{d-1}\sin ^2(\theta _i),&{} i=j \end{array}\right. }\\&= \prod _{i=1}^{d-1}\sin ^2(\theta _i) \delta _{ij} \end{aligned} \end{aligned}$$

(2.69)

Therefore, $\mathbf{G}_{\mathscr {S}_r}(\varvec{\theta })={{\mathrm{diag}}}[1,\sin ^2(\theta _1),\ldots ,\prod _{d=1}^{D-1}\sin ^2(\theta _d)]$. Another way to obtain $\mathbf{G}_{\mathscr {S}_r}(\varvec{\theta })$ is through the coordinate change

$$\begin{aligned} \mathbf{G}_{\mathscr {S}_r}(\varvec{\theta }) = \frac{d{\varvec{\theta }}^{\mathsf {T}}_{\mathscr {S}_c}}{d\varvec{\theta }_{\mathscr {S}_r}}{} \mathbf{G}_{\mathscr {S}_c}(\varvec{\theta }) \frac{d\varvec{\theta }_{\mathscr {S}_c}}{d{\varvec{\theta }}^{\mathsf {T}}_{\mathscr {S}_r}} \end{aligned}$$

(2.70)

Similar to Corollary (2.1), we have

Proposition 2.7.3

The volume adjustment of changing measure in (2.9) is

$$\begin{aligned} \left| \frac{d\varvec{\theta }_{\mathscr {R}_\mathbf{0}}}{d\varvec{\theta }_{\mathscr {S}_r}}\right| =|\mathbf{G}_{\mathscr {S}_r}(\varvec{\theta })|^{-\frac{1}{2}}=\prod _{d=1}^{D-1}\sin ^{-(D-d)}(\theta _{d}) \end{aligned}$$

(2.71)

Jacobian of the Transformation Between q -Norm Domains

The following proposition gives the weights needed for the transformation from $\mathscr {Q}^D$ to $\mathscr {B}_\mathbf{0}^D(1)$.

Proposition 2.7.4

The Jacobian determinant (weight) of $T_{\mathscr {B}\rightarrow \mathscr {Q}}$ is as follows:

$$\begin{aligned} |dT_{\mathscr {S}\rightarrow \mathscr {Q}}| = \left( \frac{2}{q}\right) ^D \left( \prod _{i=1}^{D}|\theta _i|\right) ^{2/q-1} \end{aligned}$$

(2.72)

Proof

Note

$$\begin{aligned} T_{\mathscr {B}\rightarrow \mathscr {Q}}:\, \varvec{\theta }\mapsto \varvec{\beta }=\mathrm {sgn}(\varvec{\theta })|\varvec{\theta }|^{2/q} \end{aligned}$$

The Jacobian matrix for $T_{\mathscr {B}\rightarrow \mathscr {Q}}$ is

$$\begin{aligned} \frac{d\varvec{\beta }}{d{\varvec{\theta }}^{\mathsf {T}}} = \frac{2}{q}\mathrm {diag}(|\varvec{\theta }|^{2/q-1}) \end{aligned}$$

Therefore the Jacobian determinant of $T_{\mathscr {B}\rightarrow \mathscr {Q}}$ is

$$\begin{aligned} |dT_{\mathscr {B}\rightarrow \mathscr {Q}}| = \left| \frac{d\varvec{\beta }}{d{\varvec{\theta }}^{\mathsf {T}}}\right| = \left( \frac{2}{q}\right) ^D \left( \prod _{i=1}^{D}|\theta _i|\right) ^{2/q-1} \end{aligned}$$

$\square$

The following proposition gives the weights needed for the change of domains from $\mathscr {R}^D$ to $\mathscr {B}_\mathbf{0}^D(1)$.

Proposition 2.7.5

The Jacobian determinant (weight) of $T_{\mathscr {B}\rightarrow \mathscr {R}}$ is as follows:

$$\begin{aligned} |dT_{\mathscr {B}\rightarrow \mathscr {R}}| = \frac{\Vert \varvec{\theta }\Vert _2^D}{\Vert \varvec{\theta }\Vert _{\infty }^D} \prod _{i=1}^D \frac{u_i-l_i}{2} \end{aligned}$$

(2.73)

Proof

First, we note

$$\begin{aligned} T_{\mathscr {B}\rightarrow \mathscr {R}} = T_{\mathscr {C}\rightarrow \mathscr {R}}\circ T_{\mathscr {B}\rightarrow \mathscr {C}}:\, \varvec{\theta }\mapsto \varvec{\beta }'=\varvec{\theta } \frac{\Vert \varvec{\theta }\Vert _2}{\Vert \varvec{\theta }\Vert _{\infty }}\mapsto \varvec{\beta }=\frac{\mathbf{u}-\mathbf{l}}{2}\varvec{\beta }'+\frac{\mathbf{u}+\mathbf{l}}{2} \end{aligned}$$

The corresponding Jacobian matrices are

where $\mathbf{e}_{\arg \max |\varvec{\theta }|}$ is a vector with $(\arg \max |\varvec{\theta }|)$-th element 1 and all others 0. Therefore,

$$\begin{aligned} |dT_{\mathscr {B}\rightarrow \mathscr {R}}| = |dT_{\mathscr {C}\rightarrow \mathscr {R}}|\, |dT_{\mathscr {B}\rightarrow \mathscr {C}}| = \left| \frac{d\varvec{\beta }}{d{(\varvec{\beta }')}^{\mathsf {T}}}\right| \left| \frac{d\varvec{\beta }'}{d{\varvec{\theta }}^{\mathsf {T}}}\right| = \frac{\Vert \varvec{\theta }\Vert _2^D}{\Vert \varvec{\theta }\Vert _{\infty }^D} \prod _{i=1}^D \frac{u_i-l_i}{2} \end{aligned}$$

$\square$

Splitting Hamiltonian (Lagrangian) Dynamics on ${\varvec{\mathscr {S}}}^D$

Splitting the Hamiltonian dynamics and its usefulness in improving HMC is a well-studied topic of research [13, 36, 56]. Splitting the Lagrangian dynamics (used in our approach), on the other hand, has not been discussed in the literature, to the best of our knowledge. Therefore, we prove the validity of our splitting method by starting with the well-understood method of splitting Hamiltonian [13],

$$\begin{aligned} H^*(\varvec{\theta },\mathbf{p}) = \frac{1}{2} U(\varvec{\theta }) + \frac{1}{2} \mathbf{p}^{\mathsf {T}} \mathbf{G}_{\mathscr {S}_c}(\varvec{\theta })^{-1}{} \mathbf{p} + \frac{1}{2}U(\varvec{\theta }) \end{aligned}$$

The corresponding systems of differential equations,

can be written in terms of Lagrangian dynamics in $(\varvec{\theta }, \mathbf{{v}})$ as follows:

We have solved the second dynamics (on the right) in Appendix “Geodesic on a Sphere in Cartesian Coordinates”. To solve the first dynamics, we note that

Therefore, we have

$$\begin{aligned} \tilde{\varvec{\theta }}(t)&= \tilde{\varvec{\theta }}(0)\\ \tilde{\mathbf{{v}}}(t)&= \tilde{\mathbf{{v}}}(0) -\frac{t}{2} \begin{bmatrix} \mathbf{I}\\ -\frac{{\varvec{\theta }(0)}^{\mathsf {T}}}{\theta _{D+1}(0)}\end{bmatrix} [\mathbf{I}-\varvec{\theta }(0){\varvec{\theta }(0)}^{\mathsf {T}}] \nabla _{\varvec{\theta }} U(\varvec{\theta }) \end{aligned}$$

where $\begin{bmatrix} \mathbf{I}\\ -\frac{{\varvec{\theta }(0)}^{\mathsf {T}}}{\theta _{D+1}(0)}\end{bmatrix} [\mathbf{I}-\varvec{\theta }(0){\varvec{\theta }(0)}^{\mathsf {T}}]=\begin{bmatrix} \mathbf{I}-\varvec{\theta }(0){\varvec{\theta }(0)}^{\mathsf {T}}\\ -\theta _{D+1}(0) {\varvec{\theta }(0)}^{\mathsf {T}}\end{bmatrix} = \begin{bmatrix} \mathbf{I}\\ \mathbf{0}^{\mathsf {T}}\end{bmatrix} - \tilde{\varvec{\theta }}(0) {\varvec{\theta }(0)}^{\mathsf {T}}$.

Finally, we note that $\Vert \tilde{\varvec{\theta }}(t)\Vert _2=1$ if $\Vert \tilde{\varvec{\theta }}(0)\Vert _2=1$ and $\tilde{\mathbf{{v}}}(t)\in T_{\tilde{\varvec{\theta }}(t)} \mathscr {S}_c^D$ if $\tilde{\mathbf{{v}}}(0)\in T_{\tilde{\varvec{\theta }}(0)} \mathscr {S}_c^D$.

Error Analysis of Spherical HMC

Following [36], we now show that the discretization error $e_n=\Vert \mathbf{{z}}(t_n) - \mathbf{{z}}^{(n)}\Vert =\Vert ({\varvec{\theta }}(t_n),\mathbf{{v}}(t_n)) - ({\varvec{\theta }}^{(n)},\mathbf{{v}}^{(n)})\Vert$ (i.e., the difference between the true solution and the numerical solution) is ${\mathscr {O}}(\varepsilon ^3)$ locally and $\mathscr {O}(\varepsilon ^2)$ globally, where $\varepsilon$ is the discretization step size. Here, we assume that $\mathbf{{f}}({\varvec{\theta }},\mathbf{{v}}):= \mathbf{{v}}^{\textsf {T}}\varvec{\varGamma }({\varvec{\theta }}) \mathbf{{v}} + \mathbf{{G}}({\varvec{\theta }})^{-1} \nabla _{\varvec{\theta }}U({\varvec{\theta }})$ is smooth; hence, $\mathbf{{f}}$ and its derivatives are uniformly bounded as $\mathbf{{z}}=({\varvec{\theta }},\mathbf{{v}})$ evolves within finite time duration T. We expand the true solution $\mathbf{{z}}(t_{n+1})$ at $t_n$:

$$\begin{aligned} \mathbf{{z}}(t_{n+1})&= \mathbf{{z}}(t_n)+\dot{\mathbf{{z}}}(t_n)\varepsilon +\frac{1}{2}\ddot{\mathbf{{z}}}(t_n)\varepsilon ^2 +\mathscr {O}(\varepsilon ^3)\nonumber \\&= \left[ \begin{array}{ll}{\varvec{\theta }}(t_n)\\ \mathbf{{v}}(t_n)\end{array}\right] + \left[ \begin{array}{cl}{} \mathbf{{v}}(t_n)\\ -\mathbf{{f}}({\varvec{\theta }}(t_n),\mathbf{{v}}(t_n))\end{array}\right] \varepsilon + \frac{1}{2} \left[ \begin{array}{ll}-\mathbf{{f}}({\varvec{\theta }}(t_n),\mathbf{{v}}(t_n)) \\ -\dot{\mathbf{{f}}}({\varvec{\theta }}(t_n),\mathbf{{v}}(t_n))\end{array}\right] \varepsilon ^2 +\mathscr {O}(\varepsilon ^3) \end{aligned}$$

(2.76)

We first consider Spherical HMC in Cartesian coordinates, where $\mathbf{{f}}({\varvec{\theta }},\mathbf{{v}})=\Vert \tilde{\mathbf{{v}}}\Vert ^2\varvec{\theta } + [\mathbf{I}-\varvec{\theta }{\varvec{\theta }}^{\mathsf {T}}] \nabla _{\varvec{\theta }} U(\varvec{\theta })$. From Eq. (2.34) we have

$$\begin{aligned} \mathbf{{v}}^{(n+1/2)}&= \mathbf{{v}}^{(n)} -\frac{\varepsilon }{2} (\mathbf{I}-\varvec{\theta }^{(n)}{(\varvec{\theta }^{(n)})}^{\mathsf {T}}) \nabla _{\varvec{\theta }} U(\varvec{\theta }^{(n)})\nonumber \\ \Vert \tilde{\mathbf{{v}}}^{(n+1/2)}\Vert ^2&= \Vert \tilde{\mathbf{{v}}}^{(n)}\Vert ^2 -\varepsilon {(\mathbf{{v}}^{(n)})}^{\mathsf {T}} \nabla _{\varvec{\theta }} U(\varvec{\theta }^{(n)}) + \mathscr {O}(\varepsilon ^2) \end{aligned}$$

(2.77)

Now we expand Eq. (2.35) using Taylor series as follows:

$$\begin{aligned} \varvec{\theta }^{(n+1)} =&\, \varvec{\theta }^{(n)}[1-\Vert \tilde{\mathbf{{v}}}^{(n+1/2)}\Vert ^2\varepsilon ^2/2+\mathscr {O}(\varepsilon ^4)]\\&+ \mathbf{{v}}^{(n+1/2)}\varepsilon [1 - \Vert \tilde{\mathbf{{v}}}^{(n+1/2)}\Vert ^2\varepsilon ^2/3! + \mathscr {O}(\varepsilon ^4)]\\ \mathbf{{v}}^{(n+3/4)} =&-\varvec{\theta }^{(n)} \Vert \tilde{\mathbf{{v}}}^{(n+1/2)}\Vert ^2 \varepsilon [1 - \Vert \tilde{\mathbf{{v}}}^{(n+1/2)}\Vert ^2\varepsilon ^2/3! + \mathscr {O}(\varepsilon ^4)]\\&+ \mathbf{{v}}^{(n+1/2)} [1-\Vert \tilde{\mathbf{{v}}}^{(n+1/2)}\Vert ^2\varepsilon ^2/2+\mathscr {O}(\varepsilon ^4)] \end{aligned}$$

Substituting (2.77) in the above equations yields

$$\begin{aligned} \begin{aligned} \varvec{\theta }^{(n+1)}&=\varvec{\theta }^{(n)} + \mathbf{{v}}^{(n+1/2)}\varepsilon - \varvec{\theta }^{(n)}\Vert \tilde{\mathbf{{v}}}^{(n+1/2)}\Vert ^2\varepsilon ^2/2 +\mathscr {O}(\varepsilon ^3)\\&= \varvec{\theta }^{(n)} + \mathbf{{v}}^{(n)}\varepsilon - \frac{1}{2} \mathbf{f}({\varvec{\theta }}^{(n)},\mathbf{{v}}^{(n)}) \varepsilon ^2 + \mathscr {O}(\varepsilon ^3)\\ \mathbf{{v}}^{(n+3/4)}&= \mathbf{{v}}^{(n+1/2)} -\varvec{\theta }^{(n)} \Vert \tilde{\mathbf{{v}}}^{(n+1/2)}\Vert ^2 \varepsilon - \mathbf{{v}}^{(n+1/2)} \Vert \tilde{\mathbf{{v}}}^{(n+1/2)}\Vert ^2\varepsilon ^2/2 +\mathscr {O}(\varepsilon ^3)\\&= \mathbf{{v}}^{(n)} - [(\mathbf{I}-\varvec{\theta }^{(n)}{(\varvec{\theta }^{(n)})}^{\mathsf {T}}) \nabla _{\varvec{\theta }} U(\varvec{\theta }^{(n)})/2 + \varvec{\theta }^{(n)} \Vert \tilde{\mathbf{{v}}}^{(n)}\Vert ^2]\varepsilon \\&{= \mathbf{{v}}^{(n)}\;} + [\varvec{\theta }^{(n)} {(\mathbf{{v}}^{(n)})}^{\mathsf {T}} \nabla _{\varvec{\theta }} U(\varvec{\theta }^{(n)}) - \mathbf{{v}}^{(n)} \Vert \tilde{\mathbf{{v}}}^{(n)}\Vert ^2/2]\varepsilon ^2 + \mathscr {O}(\varepsilon ^3) \end{aligned} \end{aligned}$$

With the above results, we have

$$\begin{aligned} \mathbf{{v}}^{(n+1)}&=\, \mathbf{{v}}^{(n+3/4)} -\frac{\varepsilon }{2} (\mathbf{I}-\varvec{\theta }^{(n+1)}{(\varvec{\theta }^{(n+1)})}^{\mathsf {T}}) \nabla _{\varvec{\theta }} U(\varvec{\theta }^{(n+1)})\\&= \mathbf{{v}}^{(n)} - \mathbf{f}({\varvec{\theta }}^{(n)},\mathbf{{v}}^{(n)})\varepsilon + [\varvec{\theta }^{(n)} {(\mathbf{{v}}^{(n)})}^{\mathsf {T}} \nabla _{\varvec{\theta }} U(\varvec{\theta }^{(n)}) - \mathbf{{v}}^{(n)} \Vert \tilde{\mathbf{{v}}}^{(n)}\Vert ^2/2]\varepsilon ^2 + \mathscr {O}(\varepsilon ^3)\\&- \frac{1}{2}[ (\mathbf{I}-\varvec{\theta }^{(n)}{(\varvec{\theta }^{(n)})}^{\mathsf {T}}) \nabla _{\varvec{\theta }}^2 U(\varvec{\theta }^{(n)}) \mathbf{{v}}^{(n)} - (\varvec{\theta }^{(n)}{(\mathbf{{v}}^{(n)})}^{\mathsf {T}}+\mathbf{{v}}^{(n)}{(\varvec{\theta }^{(n)})}^{\mathsf {T}}) \nabla _{\varvec{\theta }} U(\varvec{\theta }^{(n)}) ]\varepsilon ^2 \\&= \mathbf{{v}}^{(n)} - \mathbf{f}({\varvec{\theta }}^{(n)},\mathbf{{v}}^{(n)})\varepsilon - \frac{1}{2} \dot{\mathbf{f}}({\varvec{\theta }}^{(n)},\mathbf{{v}}^{(n)})\varepsilon ^2 + \mathscr {O}(\varepsilon ^3) \end{aligned}$$

where for the last equality we need to show ${(\mathbf{{v}}^{(n)})}^{\mathsf {T}} \nabla _{\varvec{\theta }} U(\varvec{\theta }^{(n)})= -2\frac{d}{dt}\Vert \tilde{\mathbf{{v}}}^{(n)}\Vert ^2$. This can be proved as follows:

$$\begin{aligned} \frac{d}{dt}\Vert \tilde{\mathbf{{v}}}\Vert ^2&= \frac{d}{dt} [ \Vert \tilde{\mathbf{{v}}}\Vert ^2 + v_{D+1}^2] = 2[-\mathbf{{v}}^{\mathsf {T}}{} \mathbf{f} + v_{D+1}\dot{v}_{D+1}]\\&= 2\left[ -\mathbf{{v}}^{\mathsf {T}}{} \mathbf{f} +\left( -\frac{{\dot{\varvec{\theta }}}^{\mathsf {T}} \mathbf{{v}}+ {\varvec{\theta }}^{\mathsf {T}}\dot{\mathbf{{v}}}}{\theta _{D+1}} + \frac{{\varvec{\theta }}^{\mathsf {T}}{} \mathbf{{v}}}{\theta _{D+1}^{2}}\dot{\theta }_{D+1}\right) v_{D+1}\right] \\&= -2\left[ {\left( \mathbf{{v}} -\frac{v_{D+1}}{\theta _{D+1}}\varvec{\theta }\right) }^{\mathsf {T}}{} \mathbf{f} + \frac{v_{D+1}}{\theta _{D+1}} \Vert \tilde{\mathbf{{v}}}\Vert ^2 \right] \\&= -2\left[ \left( \mathbf{{v}}^{\mathsf {T}}\varvec{\theta } -\frac{v_{D+1}}{\theta _{D+1}}(\Vert \varvec{\theta }\Vert ^2-1)\right) \Vert \tilde{\mathbf{{v}}}\Vert ^2 + {\left( \mathbf{{v}} -\frac{v_{D+1}}{\theta _{D+1}}\varvec{\theta }\right) }^{\mathsf {T}}[\mathbf{I}-\varvec{\theta }{\varvec{\theta }}^{\mathsf {T}}] \nabla _{\varvec{\theta }} U(\varvec{\theta })] \right] \\&= -2\left[ \mathbf{{v}}^{\mathsf {T}}\nabla _{\varvec{\theta }} U(\varvec{\theta }) + \left( -\mathbf{{v}}^{\mathsf {T}}\varvec{\theta } -\frac{v_{D+1}}{\theta _{D+1}}(1-\Vert \varvec{\theta }\Vert ^2)\right) {\varvec{\theta }}^{\mathsf {T}}\nabla _{\varvec{\theta }} U(\varvec{\theta })\right] \\&=-2 \mathbf{{v}}^{\mathsf {T}}\nabla _{\varvec{\theta }} U(\varvec{\theta }) \end{aligned}$$

Therefore we have

$$\begin{aligned} \mathbf{z}^{(n+1)} := \begin{bmatrix}{\varvec{\theta }}^{(n+1)}\\\mathbf{{v}}^{(n+1)}\end{bmatrix} = \begin{bmatrix}{\varvec{\theta }}^{(n)}\\\mathbf{{v}}^{(n)}\end{bmatrix} +\begin{bmatrix}{} \mathbf{{v}}^{(n)}\\ -\mathbf{f}({\varvec{\theta }}^{(n)},\mathbf{{v}}^{(n)})\end{bmatrix}\varepsilon + \frac{1}{2} \begin{bmatrix}-\mathbf{f}({\varvec{\theta }}^{(n)},\mathbf{{v}}^{(n)})\\ -\dot{\mathbf{f}}({\varvec{\theta }}^{(n)},\mathbf{{v}}^{(n)})\end{bmatrix}\varepsilon ^2 + \mathscr {O}(\varepsilon ^3) \end{aligned}$$

(2.78)

The local error is

$$\begin{aligned} \begin{aligned} e_{n+1}&=\Vert \mathbf{z}(t_{n+1}) - \mathbf{z}^{(n+1)}\Vert \\&= \left\| \begin{bmatrix}{\varvec{\theta }}(t_n)-{\varvec{\theta }}^{(n)}\\ \mathbf{{v}}(t_n)-\mathbf{{v}}^{(n)}\end{bmatrix} + \begin{bmatrix}{} \mathbf{{v}}(t_n)-\mathbf{{v}}^{(n)}\\ -[\mathbf{f}(t_n)- \mathbf{f}^{(n)}]\end{bmatrix}\varepsilon + \frac{1}{2} \begin{bmatrix}-[\mathbf{f}(t_n)- \mathbf{f}^{(n)}]\\ -[\dot{\mathbf{f}}(t_n)- \dot{\mathbf{f}}^{(n)}]\end{bmatrix}\varepsilon ^2 + \mathscr {O}(\varepsilon ^3) \right\| \\&\le (1+M_1\varepsilon +M_2\varepsilon ^2) e_n + \mathscr {O}(\varepsilon ^3) \end{aligned} \end{aligned}$$

(2.79)

where $M_k = c_k\sup _{t\in [0,T]}\Vert \nabla ^k \mathbf{f}({\varvec{\theta }}(t),\mathbf{{v}}(t))\Vert ,\; k=1,2$ for some constants $c_k>0$. Accumulating the local errors by iterating the above inequality for $L=T/\varepsilon$ steps provides the following global error:

$$\begin{aligned} \begin{aligned} e_{L+1}&\le (1+M_1\varepsilon +M_2\varepsilon ^2) e_L + \mathscr {O}(\varepsilon ^3) \le (1+M_1\varepsilon +M_2\varepsilon ^2)^2 e_{L-1} + 3\mathscr {O}(\varepsilon ^3)\le \cdots \\&\le (1+M_1\varepsilon +M_2\varepsilon ^2)^L e_1 + L\mathscr {O}(\varepsilon ^3) \le (e^{M_1T}+T)\varepsilon ^2 \rightarrow 0,\quad as\; \varepsilon \rightarrow 0 \end{aligned} \end{aligned}$$

(2.80)

For Spherical HMC in the spherical coordinates, we conjecture that the integrator of Algorithm 2 still has order 3 local error and order 2 global error. One can follow the same argument as above to verify this.

Bounce in Diamond: Wall HMC for 1-Norm Constraint

Reference [46] discusses the Wall HMC method for $\infty$-norm constraint only. We can however derive a similar approach for 1-norm constraint. As shown in the left panel of Fig. 2.13, given the current state $\varvec{\theta }_0$, HMC makes a proposal $\varvec{\theta }$. It will hit the boundary to move from $\varvec{\theta }_0$ towards $\varvec{\theta }$. To determine the hit point ‘X’, we are required to solve for $t\in (0,1)$ such that

$$\begin{aligned} \Vert \varvec{\theta }_0+(\varvec{\theta }-\varvec{\theta }_0)t\Vert _1=\sum _{d=1}^D|\theta _0^d+(\theta ^d-\theta _0^d)t|=1 \end{aligned}$$

(2.81)

One can find the hitting time using the bisection method. However, a more efficient method is to find the orthant in which the sampler hits the boundary, i.e., find the normal direction $\mathbf{n}$ with elements being $\pm 1$. Then, we can find t,

$$\begin{aligned} \Vert \varvec{\theta }_0+(\varvec{\theta }-\varvec{\theta }_0)t\Vert _1 = \mathbf{n}^{\mathsf {T}}[\varvec{\theta }_0+(\varvec{\theta }-\varvec{\theta }_0)t]=1 \implies t^* = \frac{1-\mathbf{n}^{\mathsf {T}}\varvec{\theta }_0}{\mathbf{n}^{\mathsf {T}}(\varvec{\theta }-\varvec{\theta }_0)} \end{aligned}$$

(2.82)

Fig. 2.13

Wall HMC bounces in the 1-norm constraint domain. Left given the current state $\varvec{\theta }_0$, Wall HMC proposes $\varvec{\theta }$, but bounces of the boundary and reaches $\varvec{\theta }'$ instead. Right determining the hitting time by monitoring the first intersection point with coordinate planes that violates the constraint

Therefore the hit point is $\varvec{\theta }_0'=\varvec{\theta }_0+(\varvec{\theta }-\varvec{\theta }_0)t^*$ and consequently the reflection point is

$$\begin{aligned} \varvec{\theta }'=\varvec{\theta }-2\mathbf{n}^*\langle \mathbf{n}^*,\varvec{\theta }-\varvec{\theta }_0'\rangle = \varvec{\theta }-2\mathbf{n}(\mathbf{n}^{\mathsf {T}}\varvec{\theta }-1)/D \end{aligned}$$

(2.83)

where $\mathbf{n}^*:=\mathbf{n}/\Vert \mathbf{n}\Vert _2$ and $\mathbf{n}^{\mathsf {T}}\varvec{\theta }_0'=1$ because $\varvec{\theta }_0'$ is on the boundary with the normal direction $\mathbf{n}^*$.

It is in general difficult to directly determine the intersection of $\varvec{\theta }-\varvec{\theta }_0$ with boundary. Instead, we can find its intersections with coordinate planes $\{\varvec{\pi }_d\}_{d=1}^D$, where $\varvec{\pi }_d:=\{\varvec{\theta }\in \mathbb R^D|\theta ^d=0\}$. The intersection times are defined as $\mathbf{T}=\{\theta _0^d/(\theta _0^d-\theta ^d)|\theta _0^d\ne \theta ^d\}$. We keep those between 0 and 1 and sort them in ascending order (Fig. 2.13, right panel). Then, we find the intersection points $\{\varvec{\theta }_k:=\varvec{\theta }_0+(\varvec{\theta }-\varvec{\theta }_0)T_k\}$ that violate the constraint $\Vert \varvec{\theta }\Vert \le 1$. Denote the first intersection point outside the constrained domain as $\varvec{\theta }_k$. The signs of $\varvec{\theta }_k$ and $\varvec{\theta }_{k-1}$ determine the orthant of the hitting point $\varvec{\theta }_0'$.

Note, for each $d\in \{1,\ldots D\}$, $({{\mathrm{sign}}}(\varvec{\theta }_k^d),{{\mathrm{sign}}}(\varvec{\theta }_{k-1}^d))$ cannot be $(+,-)$ or $(-,+)$, otherwise there exists an intersection point $\varvec{\theta }^*:=\varvec{\theta }_0+(\varvec{\theta }-\varvec{\theta }_0)T^*$ with some coordinate plane $\varvec{\pi }_{d^*}$ between $\varvec{\theta }_k$ and $\varvec{\theta }_{k-1}$. Then $T_{k-1}<T^*<T_k$ contradicts the order of $\mathbf{T}$.⁴ Therefore any point (including $\varvec{\theta }_0'$) between $\varvec{\theta }_k$ and $\varvec{\theta }_{k-1}$ must have the same sign as ${{\mathrm{sign}}}({{\mathrm{sign}}}(\varvec{\theta }_k)+{{\mathrm{sign}}}(\varvec{\theta }_{k-1}))$; that is

$$\begin{aligned} \mathbf{n} = {{\mathrm{sign}}}({{\mathrm{sign}}}(\varvec{\theta }_k)+{{\mathrm{sign}}}(\varvec{\theta }_{k-1})) \end{aligned}$$

(2.84)

After moving from $\varvec{\theta }$ to $\varvec{\theta }'$, we examine whether $\varvec{\theta }'$ satisfies the constraint. If it does not satisfy the constraint, we repeat above procedure with $\varvec{\theta }_0\leftarrow \varvec{\theta }_0'$ and $\varvec{\theta }\leftarrow \varvec{\theta }'$ until the final state is inside the constrained domain. Then we adjust the velocity direction by

$$\begin{aligned} \mathbf{{v}} \leftarrow (\varvec{\theta }'-\varvec{\theta }_0')\frac{\Vert \mathbf{{v}}\Vert }{\Vert \varvec{\theta }'-\varvec{\theta }_0'\Vert } \end{aligned}$$

(2.85)

Algorithm 4 summarizes the above steps.