Parameterizations for Gradient-based Markov Chain Monte Carlo on the Stiefel Manifold: A Comparative Study

Jauch et al. (Jauch et al., 2021) parameterized an orthogonal matrix $\boldsymbol{\Upsilon}$ with a $J\times K$ unconstrained matrix $\tilde{\boldsymbol{\Upsilon}}\in\mathbb{R}^{J\times K}$ using the polar expansion, $\boldsymbol{\Upsilon}=\tilde{\boldsymbol{\Upsilon}}\left(\tilde{\boldsymbol{\Upsilon}}^{\top}\tilde{\boldsymbol{\Upsilon}}\right)^{-\frac{1}{2}}$, where $\left(\tilde{\boldsymbol{\Upsilon}}^{\top}\tilde{\boldsymbol{\Upsilon}}\right)^{-\frac{1}{2}}$ denotes the inverse of the symmetric square root of $\tilde{\boldsymbol{\Upsilon}}^{\top}\tilde{\boldsymbol{\Upsilon}}$. Representing the singular value decomposition of $\tilde{\boldsymbol{\Upsilon}}$ as $\tilde{\boldsymbol{\Upsilon}}=\boldsymbol{U}\boldsymbol{S}\boldsymbol{V}^{\top}$, we get $\left(\tilde{\boldsymbol{\Upsilon}}^{\top}\tilde{\boldsymbol{\Upsilon}}\right)^{-\frac{1}{2}}=\boldsymbol{V}\boldsymbol{S}^{-\frac{1}{2}}\boldsymbol{V}^{\top}$. We sample $\boldsymbol{\varphi}=\textrm{vec}\left(\tilde{\boldsymbol{\Upsilon}}\right)$ , where $\textrm{vec}\left(\cdot\right)$ denotes the column-wise vectorization operator. Following (Jauch et al., 2021), we set the intermediate distribution to the Wishart distribution. Then, the target kernel is specified as

Nirwan and Bertschinger (Nirwan and Bertschinger, 2019) proposed a parameterization based on Householder transformations. Orthogonal matrix $\boldsymbol{\Upsilon}$ is written as an ordered product of Householder reflectors:

where $\boldsymbol{O}_{\left(J-K\right)\times K}$ is a $\left(J-K\right)\times K$ matrix of zeros, $\boldsymbol{I}_{J\times K}$ is a $J\times K$ matrix with the identity matrix as its top block and the remaining entries zero, $\boldsymbol{e}_{n,1}=\left(1,\boldsymbol{0}_{n-1}^{\top}\right)^{\top}$, $\mathrm{sgn}\left(\cdot\right)$ is the sign operator, and $\left\|\cdot\right\|$ is the Euclidean norm. The vector to be sampled is $\boldsymbol{\varphi}=\left(\boldsymbol{v}_{J}^{\top},\boldsymbol{v}_{J-1}^{\top},...,\boldsymbol{v}_{J-K+1}^{\top}\right)^{\top}$. Using this approach, we can avoid computing the Jacobian adjustment term.

Jauch et al. (Jauch et al., 2020) parameterized orthogonal matrices using the modified Cayley transform (Shepard et al., 2015). Given a $K\times K$ skew symmetric matrix

the modified Cayley transform of $\boldsymbol{X}$ can be written as

For brevity, hereinafter we call this type of transformation as Cayley transform. $\boldsymbol{X}$ has the following block structure:

where $\boldsymbol{A}\in\mathbb{R}^{\left(J-K\right)\times K}$ and $\boldsymbol{B}\in\mathrm{Skew}\left(K\right)$. Let $\boldsymbol{b}$ be a $K\left(K-1\right)/2$-dimensional vector of independent elements of $\boldsymbol{B}$ obtained by eliminating the diagonal and supradiagonal elements of $\textrm{vec}\left(\boldsymbol{B}\right)$. We sample $\boldsymbol{\varphi}=\left(\boldsymbol{b}^{\top},\textrm{vec}\left(\boldsymbol{A}^{\top}\right)\right)^{\top}$. The Jacobian adjustment term is defined as follows:

where $\boldsymbol{D}_{K}$ is a $K^{2}\times K\left(K-1\right)/2$ matrix such that $\boldsymbol{D}_{K}\boldsymbol{b}=\mathrm{vec}\left(\boldsymbol{B}\right)$ and $\boldsymbol{K}_{J,J}$ is the commutation matrix that satisfies $\boldsymbol{K}_{J,J}\mathrm{vec}\left(\boldsymbol{A}\right)=\mathrm{vec}\left(\boldsymbol{A}\right)^{\top}$. See Appendix B of (Jauch et al., 2020) (pp. 1581-1582) for further details to construct $\boldsymbol{D}_{K}$ and $\boldsymbol{K}_{J,J}$.

Pourzanjani et al. (Pourzanjani et al., 2021) proposed to use a Givens representation of $\boldsymbol{\Upsilon}$:

where $\theta_{k,j}$ is a coordinate variable, $\boldsymbol{R}_{k,j}\left(\theta_{k,j}\right)$ is a $J\times J$ matrix that takes the form of an identity matrix except for the $\left(j,j\right)$ and $\left(k,k\right)$ positions which are replaced by $\cos\theta_{j,k}$ and the $\left(j,k\right)$ and $\left(k,j\right)$ positions which are replaced by $-\sin\theta_{j,k}$ and $\sin\theta_{k,j}$, respectively. $\theta_{1,2},\theta_{2,3},...,\theta_{K,K+1}\in\left(-\pi,\pi\right)$ and the remaining coordinates are in the range $\left(-2\pi,2\pi\right)$. This parameterization can induce the posterior of $\theta_{j,k}$ to be multimodal (Pourzanjani et al., 2021). To address this, Pourzanjani et al. (Pourzanjani et al., 2021) suggested to further reparameterize $\theta_{k,j}$ with an independent auxiliary parameter $r_{k,j}$,

The marginal distribution of $r_{k,j}$ is given by a normal distribution with a mean of one and a standard deviation of 0.1, $r_{k,j}\sim\mathcal{N}\left(1,0.1^{2}\right)$. See Section 4.2 of (Pourzanjani et al., 2021) (pp. 647-653) for further discussion. The Jacobian adjustment term is:

Table 1 summarizes the number of essential parameters, i.e., the dimension of $\boldsymbol{\varphi}$, for the four alternative parameterizations. The polar expansion has the largest number of essential parameters, whereas the Householder and Cayley transforms the smallest.

We compared the alternative parameterizations for uniform sampling on the Stiefel manifold. The dimension of the target distribution was set to combinations of $K\in\left\{3,10\right\}$ and $J\in\left\{10,100,200\right\}$. We did not examine an approach based on the Cayley transform for these situations with $J=K$ because it requires considerable coding effort.⁴⁴4When $J=K$, $\boldsymbol{A}$ disappears and it holds that $\boldsymbol{X}=\boldsymbol{B}$. Thus, conditional executions must be added to virtually all lines involving $\boldsymbol{X}$. Owing to the limitations of the computational budget, we did not test the Givens representation for the most high-dimensional case, $\left(J,K\right)=\left(200,10\right)$.

Table 2 presents the simulation results. Polar expansion was the best parameterization regardless of the efficiency measure. The win margin was larger for high-dimensional cases. For the minESS/iter metric, the remaining three parameterizations were comparable. For minESS/sec, Householder was the most efficient parameterization after polar expansion, whereas the Givens representation was the worst.

The network eigenmodel was first introduced by Hoff (2009) and has been widely used as a testing bench (Byrne and Girolami, 2013; Jauch et al., 2021; Pourzanjani et al., 2021); however, a thorough comparison of alternative parameterizations is absent. A graph matrix represents how proteins in a protein network interact with each other. Probability of a connection between proteins is modeled using a symmetric matrix of probabilities whose rank is much smaller than the dimension of the observed graph matrix. Given a $J\times J$ symmetric graph matrix $\boldsymbol{Y}=\left(y_{j,j^{\prime}}\right)$ with $y_{j,j^{\prime}}\in\left\{0,1\right\}$, the probability of the connections is specified as follows: For $j,j^{\prime}=1,...,J$,

where $\boldsymbol{\Lambda}=\mathrm{diag}\left(\lambda_{1},\lambda_{2},\lambda_{3}\right)$, $\boldsymbol{\Upsilon}\in\mathcal{V}^{J\times K}$, $\mu$, and $\lambda_{k}$ are unknown parameters, and $\Phi\left(\cdot\right)$ is the probit link function. Following previous studies, we selected $K=3$. We assign normal priors to $\mu$ and $\lambda_{k}$: $\mu\sim\mathcal{N}\left(0,10^{2}\right)$ $\lambda_{k}\sim\mathcal{N}\left(0,J\right)$. The dimension of the dataset was $J=270$.

As shown in Table 3, none of the four parameterizations performed well. Polar expansion was the best, but its minESS/iter and minESS/sec were too small to conduct a reliable posterior analysis within a reasonable time. Previous studies reported minESS/iter for only some of the unknown parameters. Although not reported in Table 3, we confirmed the results of (Jauch et al., 2021) and (Pourzanjani et al., 2021). As reported in these papers, some of the unknown parameters (e.g., $\lambda_{k}$s and $c$) are effectively simulated, but the posterior draws of some entries in $\boldsymbol{\Upsilon}$ are strongly autocorrelated.

In probabilistic principal component analysis (Tipping and Bishop, 2002), a $J$-dimensional demeaned vector $\boldsymbol{y}_{i}\in\mathbb{R}^{J}$ was modeled using a linear function of low-dimensional latent state $\boldsymbol{z}_{i}\in\mathbb{R}^{K}$ with $K<J$. The distribution of $\boldsymbol{y}_{i}$ is specified as follows: For $i=1,...,N$,

where $\boldsymbol{\Lambda}=\mathrm{diag}\left(\lambda_{1},...,\lambda_{K}\right)$, $\boldsymbol{W}\in\mathcal{V}^{J\times K}$, and $\lambda_{k}$ are unknown parameters, and $\sigma^{2}$ is the variance of the idiosyncratic shocks. This model is closely related to static latent factor modeling (see (Lopes, 2014) for an overview). $\boldsymbol{z}_{i}$ can be seen as a vector of latent factors and the quantity $\boldsymbol{W}\boldsymbol{\Lambda}$ as a factor loading matrix.

We applied this model to two specifications of synthetic data. The first specification was adopted from (Nirwan and Bertschinger, 2019): $N=150$, $J=5$, $K=2$, $\left(\lambda_{1},\lambda_{2}\right)=\left(9,1\right)$, and $\sigma^{2}=0.01^{2}$. The second specification was adopted from (Jauch et al., 2020; Pourzanjani et al., 2021): $N=100$, $J=50$, $K=3$, $\left(\lambda_{1},\lambda_{2},\lambda_{3}\right)=\left(5,3,1.5\right)$, and $\sigma^{2}=1$. For both specifications, $\boldsymbol{W}$ was generated uniformly from $\mathcal{V}^{50\times 3}$⁵⁵5First, each entry of $\boldsymbol{A}\in\mathbb{R}^{J\times K}$ is simulated independently from the standard normal distribution and then $\boldsymbol{W}$ is obtained through a polar expansion, $\boldsymbol{W}=\boldsymbol{A}\left(\boldsymbol{A}^{\top}\boldsymbol{A}\right)^{-\frac{1}{2}}$ (see Proposition 7.1 of (Eaton, 1989), pp. 100-101).. Non-informative priors were employed for all the parameters and cases.

The first two rows of panels (a) and (b) in Table 4 summarize the results. In terms of minESS/iter, the performance order of the four parameterizations are not clear. For the first dataset, polar expansion performed the best. For the second dataset, Cayley transform performed the best. When the performance is evaluated based on the minESS/sec metric, for the first dataset, polar expansion and Cayley transformation are largely comparable and superior to the other parameterizations. For the second dataset, polar expansion performed the best. Givens representation did not work well for either dataset. Similar to uniform sampling cases, in terms of minESS/iter, Givens representation was marginally inferior to Householder and Cayley transformations, while Givens representation was slow to draw, leading to a much smaller minESS/sec.

Following (Nirwan and Bertschinger, 2019), we applied the model to the breast cancer Wisconsin dataset retrieved from scikit-learn, a machine learning library for the Python programming language (Pedregosa et al., 2011). The model was extended by incorporating a mean vector $\boldsymbol{\mu}$ as

We assigned a non-informative prior to $\boldsymbol{\mu}$: $p\left(\boldsymbol{\mu}\right)\propto 1$. As shown in Table 4, polar expansion performed the best, irrespective of the performance measure. We do not report on Givens representation because the posterior simulation was intolerably slow.

The final testing bench is matrix completion for causal analysis. Matrix completion has been studied extensively in the machine learning community (e.g., (Salakhutdinov and Mnih, 2007; Ding et al., 2011; Babacan et al., 2012; Mai and Alquier, 2015; Farias et al., 2022; Tanaka, 2022; Yuchi et al., 2023; Zhai and Gutman, 2023)). It has been applied to causal analysis with event study designs using panel data (e.g., (Samartsidis et al., 2020; Nethery et al., 2021; Tanaka, 2021; Pang et al., 2022; Samartsidis et al., 2024)). In this framework, outcomes under treatment are removed from an outcome matrix, and the outcome matrix is “completed”, treating the missing entries as unrealized potential outcomes under treatment, or counterfactual outcomes. By comparing between the estimates of the counterfactual and realized outcomes, we can infer the treatment effects.

Let $\boldsymbol{Y}\in\mathbb{R}^{J\times T}$ denote the matrix of outcomes for $J$ entities and $T$ time periods. The elements in $\boldsymbol{Y}$ corresponding to the treated entities and the time periods are treated as missing data. We can infer unrealized potential outcomes by completing $\boldsymbol{Y}$. Let $\boldsymbol{Y}^{\mathrm{miss}}$ denote the set of missing entries in $\boldsymbol{Y}$ and $\boldsymbol{Y}^{\mathrm{obs}}$ denote the set of observed entries.

Given $\boldsymbol{Y}^{\mathrm{miss}}$, the model of “completed” $\boldsymbol{Y}$ is composed of three matrices, $\boldsymbol{Y}=\boldsymbol{\Xi}+\boldsymbol{\Gamma}+\boldsymbol{U}$. First, $\boldsymbol{\Xi}\in\mathbb{R}^{J\times T}$ is a matrix of covariate effects, $\boldsymbol{x}_{j,t}$ is a vector of covariates, $\boldsymbol{\Xi}=\left(\xi_{j,t}\right)$ with $\xi_{j,t}=\boldsymbol{\beta}^{\top}\boldsymbol{x}_{j,t}$. Second, $\boldsymbol{\Gamma}\in\mathbb{R}^{J\times T}$ is a matrix with rank $K<\min\left(J,T\right)$ that is factorized as in singular value decomposition, $\boldsymbol{\Gamma}=\boldsymbol{\Phi}\boldsymbol{\Lambda}\boldsymbol{\Psi}^{\top}$, where $\boldsymbol{\Lambda}=\mathrm{diag}\left(\lambda_{1},...,\lambda_{K}\right)$ is a diagonal matrix, $\boldsymbol{\Phi}=\left(\phi_{j,k}\right)\in\mathcal{V}^{J\times K}$ and $\boldsymbol{\Psi}=\left(\psi_{t,k}\right)\in\mathcal{V}^{T\times K}$ are orthogonal matrices. Third, $\boldsymbol{U}=\left(u_{j,k}\right)\in\mathbb{R}^{J\times T}$ is a matrix of error terms whose entries are distributed according to an independent normal distribution with variance $\sigma^{2}$, $u_{j,t}\sim\mathcal{N}\left(0,\sigma^{2}\right).$

We conducted a posterior simulation treating $\boldsymbol{Y}^{\mathrm{miss}}$ as unknown parameters. Let $\boldsymbol{X}$ denote the set of covariates. Then, the likelihood of the “completed” $\boldsymbol{Y}$ has the standard form:

The target kernel is represented as:

where $p\left(\boldsymbol{Y}^{\mathrm{miss}},\boldsymbol{\Phi},\boldsymbol{\Lambda},\boldsymbol{\Psi},\boldsymbol{\beta},\sigma^{2}\right)$ is a prior. For $\lambda_{k}$, we employed an exponential prior $\lambda_{k}\sim\mathcal{E}\left(\eta\right)$, where $\eta$ is a prefixed rate parameter. This choice can be seen as a Bayesian counterpart to matrix completion with the nuclear norm penalty (Du et al., 2023). We assigned uniform Haar prior to $\boldsymbol{\Phi}$ and $\boldsymbol{\Psi}$: $p\left(\boldsymbol{\Phi}\right)\propto 1$ and $p\left(\boldsymbol{\Psi}\right)\propto 1$, non-informative prior to $\boldsymbol{Y}^{\mathrm{miss}}$ and $\boldsymbol{\beta}$: $p\left(\boldsymbol{Y}^{\mathrm{miss}}\right)\propto 1$, $p\left(\boldsymbol{\beta}\right)\propto 1$, and Jeffreys prior to $\sigma^{2}$: $p\left(\sigma^{2}\right)\propto\sigma^{-2}$.

Using this model, we re-examined the empirical analysis of carbon taxes on $\mathrm{CO}_{2}$ emissions in (Andersson, 2019) which uses synthetic control method (Abadie and Gardeazabal, 2003; Abadie et al., 2010). We used the annual panel data on per capita $\mathrm{CO}_{2}$ emissions from transport, covering the years 1960-2005 for 14 OECD (Organisation for Economic Co-operation and Development) countries: Australia, Belgium, Canada, Denmark, France, Greece, Iceland, Japan, New Zealand, Poland, Portugal, Spain, Sweden, Switzerland, and the United States. Thus, $J=14$ and $T=46$. Sweden implemented carbon taxes in 1990, whereas the others did not. We exploited this event as a quasi-experiment to infer the effects of the carbon taxes on $\mathrm{CO}_{2}$ emissions in Sweden, treating the other countries as a control group. See (Andersson, 2019) for further details on the dataset and a discussion on the choice of the control group. Three cases with $K\in\left\{3,5,10\right\}$ were considered.

Table 5 presents the results. Using Givens representation, for some runs with $K=10$, we obtained extremely autocorrelated chains and could not effectively compute the ESS due to numerical instabilities. Based on the minESS/iter metric, the performance of the four parameterizations were similarly poor. For the minESS/sec metric, polar expansion ranked the best.