# Reduce\_sum in a Bayesian PCA

**URL:** <https://discourse.mc-stan.org/t/reduce-sum-in-a-bayesian-pca/33369>\
**Category:** Algorithms\
**Created:** [November 17, 2023, 11:04pm UTC](https://discourse.mc-stan.org/t/reduce-sum-in-a-bayesian-pca/33369 "2023-11-17T23:04:48Z")\
**Posts on this page:** 8\
**Page:** 1

<div class="post-metadata">

**Author:** ![kokorap27](https://yyz2.discourse-cdn.com/flex030/user_avatar/discourse.mc-stan.org/kokorap27/32/11389_2.png) [@kokorap27](https://discourse.mc-stan.org/u/kokorap27)\
**Post date:** [November 17, 2023, 11:04pm UTC](https://discourse.mc-stan.org/t/reduce-sum-in-a-bayesian-pca/33369/1 "2023-11-17T23:04:48Z")

</div>

Hello Stan community!

This week I’ve been supporting a researcher with his implementation in Stan of a Bayesian model in which a PCA is included.

Regarding the Bayesian PCA, we are assuming that

\mathbf{y}\_{n}=\mathbf{W}\mathbf{z}\_{n}+\mathbf{\epsilon}\_{n},

where \mathbf{W}\in\mathbb{R}^{P \times Q} with Q\<P, \mathbf{z}\_{n} \sim \mathcal{N}\_{Q}(\mathbf{0},\mathbf{I}), and \mathbf{\epsilon}\_{n} \sim \mathcal{N}\_{P}(\mathbf{0},\sigma^{2}\mathbf{I}) for n=1,2,\ldots,N.

Furthermore, proposing some priors for the latent vectors and hyperparameters, the Bayesian hierarchical model is as follows:

\mathbf{y}\_{n} \sim \mathcal{N}\_{P}(\mathbf{W}\mathbf{z}\_{n},\sigma^{2}\mathbf{I})\\ \mathbf{w}\_{p}\mid\tau,\mathbf{\Omega} \sim \mathcal{N}\_{P}(\mathbf{0},\tau\mathbf{\Omega})\quad p=1,2,\ldots,P\\ \mathbf{z}\_{n} \sim \mathcal{N}\_{Q}(\mathbf{0},\mathbf{I})\\ \sigma \sim \mathcal{N}(0,100^{2})\\ \tau \sim \mathrm{Cauchy}(0, 2.5)\\ \mathbf{\Omega} \sim \mathrm{LKJ}(2)

The ultimate goal is to implement within-chain parallelization aiming at speeding up the Stan code via the function _reduce\_sum_, however, it turned out that computing the log-likelihood of \lbrace\mathbf{y}\_{n}\rbrace\_{n=1}^{N} with _reduce\_sum_ is slower than without it.

```no-highlight
functions {
  real partial_sum_lpdf(array[,] real y_slice,
                      int start, int end,
                      array[,] real mu,
                      real sigma) {
                     return normal_lupdf(to_array_1d(y_slice) | to_array_1d(mu[start:end,:]),sigma);
                   }
}

data {
  int<lower=0> N;         
  int<lower=0> P;         
  int<lower=0> Q;   
  matrix[N,P] Y;
}

parameters {
  matrix[P, Q] W;
  real<lower=1e-6> sigma;     
  matrix[N, Q] Z;     
  cholesky_factor_corr[Q] L_Omega;
  vector<lower=1e-6>[Q] tau;
}

transformed parameters { 
  matrix[N, P] mu = (Z * W');
}

model {
  // Log-likelihood alternatives
  for (n in 1:N) {
    Y[n] ~ normal(mu[n], sigma);
  }
  to_vector(Y) ~ normal(to_vector(mu),sigma);
  int grainsize = 1;
  target += reduce_sum(partial_sum_lpdf,to_array_2d(Y),grainsize,to_array_2d(mu),sigma);
}

```

Above are the relevant pieces of the Stan code. As you can see, I’m trying three alternatives for computing the log-likelihood: the first uses a _for_ loop, the second a vectorized version of the fist one, and the third deploys _reduce\_sum_. This last one is really slow compare with the other ones. Could someone please help me out by detecting issues or suggesting ways of improving the implementation of _reduce\_sum_?

Regards, Román.

---

<div class="post-metadata">

**Author:** ![avehtari](https://yyz2.discourse-cdn.com/flex030/user_avatar/discourse.mc-stan.org/avehtari/32/5935_2.png) [@avehtari](https://discourse.mc-stan.org/u/avehtari)\
**Post date:** [November 18, 2023, 2:42pm UTC](https://discourse.mc-stan.org/t/reduce-sum-in-a-bayesian-pca/33369/2 "2023-11-18T14:42:59Z")

</div>

> [@kokorap27](#):
>
> `grainsize = 1`

With `grainsize = 1` you probably have quite high overhead, which would be reduce with bigger grainsizes.

> [@kokorap27](#):
>
> ` matrix[N, P] mu = (Z * W');`

Why don’t you define Wt as `matrix[Q, P]` to avoid repeated transpose?

> [@kokorap27](#):
>
> ```no-highlight
> for (n in 1:N) {
> Y[n] ~ normal(mu[n], sigma);
> }
> to_vector(Y) ~ normal(to_vector(mu),sigma);
> 
> ```

You seem to have the normal model twice, and you should drop the for loop lines.

If N and P are big, see also the blog post for speed optimizations

> **[Options for improving Stan sampling speed](https://blog.mc-stan.org/2022/08/03/options-for-improving-stan-sampling-speed/)**
>
> In this blog post, I (Aki) discuss different optimizations that can speed up Stan sampling. I approach this from the user viewpoint, and you can thank numerous (10+) Stan developers for many of the…

---

<div class="post-metadata">

**Author:** ![edm](https://avatars.discourse-cdn.com/v4/letter/e/f08c70/32.png) [@edm](https://discourse.mc-stan.org/u/edm)\
**Post date:** [November 19, 2023, 4:24pm UTC](https://discourse.mc-stan.org/t/reduce-sum-in-a-bayesian-pca/33369/3 "2023-11-19T16:24:56Z")

</div>

For this particular model, I think you will get the most speedup by marginalizing out `z` and writing the resulting multivariate normal as a function of the sample mean and covariance matrix.

If you marginalize over `z` you get

\mathbf{y}\_n \sim \mathcal{N}\_{P}(\mathbf{0}, \mathbf{WW}' + \sigma^{2}\mathbf{I})

Then the \mathbf{y}\_n are iid and you can evaluate the likelihood using sufficient statistics (sample mean vector and covariance matrix). This means you don’t need to loop over `n`, you can compute the sample mean vector and covariance matrix once in the generated data block (or they can be sent in as data). The Stan code to evaluate the multivariate normal via sample mean and covariance matrix is [at this post.](https://discourse.mc-stan.org/t/cfa-model-cannot-recover-simulated-values/27857/4)

If you want the \mathbf{z}\_n, you can obtain them in generated quantities by using the multivariate normal rng. You have that (\mathbf{y}\_n'\ \mathbf{z}\_n')' is multivariate normal and, from there, can obtain the distribution of \mathbf{z}\_n \mid\ \mathbf{y}\_n as another multivariate normal. I am working from memory, but I believe the result is

\mathbf{z}\_n \mid\ \mathbf{y}\_n \sim \mathcal{N}\_{Q}(\mathbf{W}'(\mathbf{WW}' + \sigma^{2}\mathbf{I})^{-1}\mathbf{y}\_n,\ \mathbf{I} - \mathbf{W}'(\mathbf{WW}' + \sigma^{2}\mathbf{I})^{-1}\mathbf{W})

For what it’s worth, I would call this a factor analysis instead of a PCA because it involves a multivariate normal model, whereas PCA is typically a matrix algorithm without a likelihood.

---

<div class="post-metadata">

**Author:** ![kokorap27](https://yyz2.discourse-cdn.com/flex030/user_avatar/discourse.mc-stan.org/kokorap27/32/11389_2.png) [@kokorap27](https://discourse.mc-stan.org/u/kokorap27)\
**Post date:** [December 1, 2023, 3:28pm UTC](https://discourse.mc-stan.org/t/reduce-sum-in-a-bayesian-pca/33369/4 "2023-12-01T15:28:17Z")

</div>

> [@avehtari](#):
>
> With `grainsize = 1` you probably have quite high overhead, which would be reduce with bigger grainsizes.

You were right… It’s better to test different values of this parameter

> [@avehtari](#):
>
> Why don’t you define Wt as `matrix[Q, P]` to avoid repeated transpose?

I did it and it also helps

> [@avehtari](#):
>
> You seem to have the normal model twice, and you should drop the for loop lines.

I should mention that I was commenting out two of the three alternatives when compiling the Stan code

> [@avehtari](#):
>
> If N and P are big, see also the blog post for speed optimizations

Thanks a lot, I definitely will read the post thoroughly

---

<div class="post-metadata">

**Author:** ![kokorap27](https://yyz2.discourse-cdn.com/flex030/user_avatar/discourse.mc-stan.org/kokorap27/32/11389_2.png) [@kokorap27](https://discourse.mc-stan.org/u/kokorap27)\
**Post date:** [December 1, 2023, 3:42pm UTC](https://discourse.mc-stan.org/t/reduce-sum-in-a-bayesian-pca/33369/5 "2023-12-01T15:42:28Z")

</div>

You are right, if I marginalize out \mathbf{z}\_{n} I get the marginal distribution you wrote. Nevertheless, if I’m not wrong, Stan needs the conditional distribution of \mathbf{y}\_{n} not the marginal distribution of \mathbf{y}\_{n}.

The conditional distribution \mathbf{z}\_{n} given \mathbf{y}\_{n} assumes that \mathbf{W} is given. Within the Bayesian approach, \mathbf{W} is also a matrix of random variables.

---

<div class="post-metadata">

**Author:** ![kokorap27](https://yyz2.discourse-cdn.com/flex030/user_avatar/discourse.mc-stan.org/kokorap27/32/11389_2.png) [@kokorap27](https://discourse.mc-stan.org/u/kokorap27)\
**Post date:** [December 1, 2023, 3:48pm UTC](https://discourse.mc-stan.org/t/reduce-sum-in-a-bayesian-pca/33369/6 "2023-12-01T15:48:20Z")

</div>

Thanks a lot @avehtari and @edm for your comments and suggestions - they were really useful for better understanding my model.

I’ve have implemented a Bayesian PCA from scratch to test `reduce_sum`, which allowed me to progress on finding a more efficient way to implement within-chain parallelization in the model. I’ll post it today.

---

<div class="post-metadata">

**Author:** ![edm](https://avatars.discourse-cdn.com/v4/letter/e/f08c70/32.png) [@edm](https://discourse.mc-stan.org/u/edm)\
**Post date:** [December 1, 2023, 3:57pm UTC](https://discourse.mc-stan.org/t/reduce-sum-in-a-bayesian-pca/33369/7 "2023-12-01T15:57:46Z")

</div>

I would be wiling to bet a beer that you will obtain the same posterior distribution regardless of whether you use marginal or conditional likelihood.

About \mathbf{W} being given: in generated quantities, \mathbf{W} has already been sampled so is basically given (the “given” just changes at each iteration). Sampling \mathbf{z} in generated quantities is like a Gibbs sampling step.

---

<div class="post-metadata">

**Author:** ![kokorap27](https://yyz2.discourse-cdn.com/flex030/user_avatar/discourse.mc-stan.org/kokorap27/32/11389_2.png) [@kokorap27](https://discourse.mc-stan.org/u/kokorap27)\
**Post date:** [December 1, 2023, 4:23pm UTC](https://discourse.mc-stan.org/t/reduce-sum-in-a-bayesian-pca/33369/8 "2023-12-01T16:23:30Z")

</div>

> [@edm](#):
>
> I would be wiling to bet a beer that you will obtain the same posterior distribution regardless of whether you use marginal or conditional likelihood.

This a good exercise… Even better if there is a beer in play :)  
I’ll try this out once I finish the implementation of the within-chain parallelization via `reduce_sum`

> [@edm](#):
>
> About \mathbf{W} being given: in generated quantities, \mathbf{W} has already been sampled so is basically given (the “given” just changes at each iteration). Sampling \mathbf{z} in generated quantities is like a Gibbs sampling step.

I see your point… I consider it’s worth trying this approach. As I mentioned above, I’ll try this out.