Non-centered Parameterization for MCMC

Adapted from an appendix of my MS thesis.

Non-centered Parameterization

For relatively simple models, samplers mostly just work, providing us with an accurate estimate of the posterior. However, certain posterior geometries are challenging for samplers. A common example is Neal’s funnel. The shape at the top of the funnel is quite wide before narrowing into a small neck. Samplers function by taking steps from one set of parameter values to another, and a key setting is how large of a step to take when exploring the posterior surface. In complex geometries, such as with Neal’s funnel, a step size that works well in one area fails in another [1].

In hierarchical models the posterior geometry is largely defined by the correlation of hyperpriors to other parameters, which can result in funnel geometry that are difficult to sample. These inconvenient posteriors can create inefficient MCMC caused by steep curvature in the negative log-likelihood and lead to divergent transitions. Transforming the prior to be smoother helps. There is a relatively simple tweak to models, referred to as a non-centered parameterization, that helps alleviate the issue [1].

Example: Devil’s Funnel prior

$$\begin{aligned} v &\sim \mathrm{Normal}(0,\sigma_ v) \\ x &\sim \mathrm{Normal}(0, \exp(v)).\end{aligned}$$

As $\sigma_ v$ increases, a trough forms in the prior probability surface known as centered parameterization that is difficult for Hamiltonian dynamics to sample [2, 3].

To handle this, we could configure a smaller steps size to manage the steepness better, but then sampling would take longer to explore the posterior. Alternatively, we can re-parameterize (transform) to make the surface smoother known non-centered parameterization. For this purpose, we add an auxiliary variable $z$ that has a smooth probability surface. We then sample that auxiliary variable, and transform it to obtain the target variable distribution. For the case of the devil’s funnel prior we do the following [2, 3].

$$\begin{aligned} v &\sim \mathrm{Normal}(0, \sigma_ v) \\ z &\sim \mathrm{Normal}(0, 1) \\ x &= z\exp(v).\end{aligned}$$

We can also transform sampling using Cholesky factors. The Cholesky factor $\boldsymbol{L}$ provides an efficient means for encoding a correlation matrix $\boldsymbol{\Omega}$ since it requires fewer floating point numbers than the full correlation matrix. It is defined by $\boldsymbol{\Omega}=\boldsymbol{L}\boldsymbol{L}^ \top$ where $\boldsymbol{L}$ is a lower-triangular matrix. We can sample data with correlation $\boldsymbol{\Omega}$ , and standard deviation $\sigma_ i$ , using the Cholesky factorization $\boldsymbol{L}$ of $\boldsymbol{\Omega}$ as follows where $\boldsymbol{Z}$ is a matrix of $z$ -scores sampled from a standard normal distribution [2, 3].

$$\begin{aligned} \boldsymbol{X_ \Omega} \sim \operatorname{diag}(\sigma_ i)\boldsymbol{L}\boldsymbol{Z}.\end{aligned}$$

References

1. Martin, Osvaldo A., Kumar, Ravin, Lao, Junpeng (2021) Bayesian Modeling and Computation in Python. Chapman and Hall/CRC.
2. McElreath, Richard (2020) Statistical Rethinking: A Bayesian Course with Examples in R and STAN. Chapman and Hall/CRC.
3. Dustin Stansbury, Richard McElreath (2024) Correlated Features. Zenodo.