This post is from Bob.
Chenyang Zhong, a stats professor at Columbia, presented the following paper at our Bayesian computation reading group on Friday.
- Chenyang Zhong, Shouxuan Ji, Tian Zheng. 2024. Graph-Enabled Fast MCMC Sampling with an Unknown High-Dimensional Prior Distribution. arXiv 2408.02122.
The goal is to be able to chain Bayesian inferences on a data stream in situations where there’s no analytic form of the posterior. The problem is that the textbook solution of using analytic posteriors (e.g., chaining binomial likelihoods with beta priors), only works for simple conjugate models.
The model
The unknown prior in this case is constructed from draws from the posterior of a previous model. To ground us with some notation, suppose our joint model is the product of a likelihood and prior,
p(y, theta | x) = p(y | theta, x) * p(theta).
Sequential data
For example, consider the case where we receive a sequence of data sets
(x1, y1), (x2, y2), …, (xn, yn), ….
In Chenyang’s case, the data is private, so you can view the problem as a kind of federated learning or as a kind of meta-analysis. We know we can fit a model p(theta, y | x) and get posterior draws
theta_post(1), …, theta_post(M) ~ p(theta | x1, y1).
Chenyang assumes the posterior draws may be shared. What we’d like to do is use the posterior p(theta | x1, y1) as the prior for theta when analyzing data x2, y2, i.e.,
p(theta | x1, x2, y1, y2) propto p(y2 | theta, x2) * p(theta | y1, x1).
But we don’t have a closed form expression for p(theta | y1, x1), so what do we do?
Kernel density estimate as a posterior approximation
Chenyang’s idea is to use a kernel density estimate with a normal basis as a proxy for p(theta | y1, x1). Specifically, what he’s going to do is write an empirical prior that penalizes squared distance from the posterior draws.
p(theta | y1, x1) approx 1/M SUM_m normal(theta | theta_post(m), h * I),
where I is the identity matrix and h > 0 is a variance parameter. Most of Chenyang’s paper is about how to compute this efficiently. By “high dimensions” he’s talking about a modest 6 to 20 dimensions. He takes M to be about 10,000. He then goes about constructing a really neat way to Metropolis sample exactly using only nearest neighbors of theta.
It’s easy to establish that the maximum likelihood estimate and posterior mean of the prior will be the sample mean of the posterior draws theta(1), …, theta(M). How strongly it concentrates around that mean will depend on how spread out the posterior draws are. What’s interesting is that how hard it pulls does not matter how large M is—the more the merrier in terms of accuracy. But it will depend on the variance term h.
What if we just use Stan?
Stan’s pretty darn fast at normal approximations, so what if we just coded the approximate posterior directly rather than trying to use a graph of nearest neighbors and adjust? Turns out it works very cleanly. Given the data set in his paper, which is a logistic regression with N = 1500, it takes Stan 2s to fit the posterior for (x1, y1), and 35s to fit the posterior for (x2, y2) using 10,000 posterior draws from the approximate posterior (this is my 2017 iMac Pro, which is very slow compared to current ARM-based Macs).
I coded this all up as a Stan case study that you can find here, along with the results:
- Stan case study for Zhong’s kernel posterior approximation.
Should I add this as a new technique to the third section of the Stan User’s Guide? This problem comes up all the time on our forums.
What’s left to do?
I have two four questions left after fitting the model.
- How to set the variance term h? It doesn’t affect the mean, but it does affect how strongly the prior concentrates. Chenyang mentioned something about potentially wanting to discount the past, which you can do. Is there a way to set that by tweaking h? Alternatively, is there a way to set it optimally so that the posterior for the final fit is closest to p(theta | x1, y1, x2, y2) in cases where we want to equally weight the past?
- How many posterior draws do we need? In simple cases like this, probably not 10,000!
- Is this just an easier way to do the computation rather than estimating a covariance matrix and using a multivariate normal? The reason I ask is that I know you can generate from the empirical covariance by differencing in this way—it’s equation (11) in Goodman and Weare’s affine invariant sampling paper.
- How can this approach handle constrained parameters? We can just keep exactly the same code and keep the constraints on the parameters and everything should work, but it seems more natural to lay down a multivariate normal approximation on the unconstrained scale (e.g., after log transforming positive constrained parameters).
[edit: added third and fourth question]
> Is this just an easier way to do the computation rather than estimating a covariance matrix and using a multivariate normal?
I don’t think so. I mean, I can use a KDE with a normal kernel on some data that’s very non normal and will get an estimate of a density that’s very non normal.
I think that the same is true here. Basically what we’re doing is a radial basis function approximation to a density. And the locations of the centers are the locations of the parameter samples.
>Should I add this as a new technique to the third section of the Stan User’s Guide? This problem comes up all the time on our forums.
Yes, please. I was just struggling with this issue this morning. Really interesting.
I vote yes on adding to the user guide… I have this problem every month now it seems.
The kernel‐density proxy is just one way to “lift” a discrete posterior‐sample cloud into a smooth prior. An alternative is to lean on optimal‐transport (Earth Mover’s) ideas. We could use Wasserstein barycenters as priors.
See also https://jmlr.csail.mit.edu/papers/volume19/17-084/17-084.pdf
See also https://arxiv.org/pdf/1310.4375
Would potentially be a good use case for copulas, if I understand correctly.
> By “high dimensions” he’s talking about a modest 6 to 20 dimensions.
For kernel density estimates 20 dimensions is high. How to choose h for each dimension gets difficult. Based on my own experiments and papers I’ve seen, I expect that the best methods might get something useful up to 50 dimensions for the usual sample sizes.
1. There are tons of papers for selecting h for KDE and related methods, but they are not considering the prior case where you would prefer having too wide approximation than too narrow. Inflating h and using Student-t kernel make the prior weaker.
2. The number of posterior draws needed increases with the number of dimensions, unless you are ok letting the approximation get worse, which might be fine. You can reduce the computational cost of making a normal mixture model with fewer components, but using all the draws to fit that mixture model.
Thanks to everyone for:
Daniel Lakeland: I was forgetting that it’s additive, not multiplicative. If we multiplied all the probabilities, you’d get something where the mean of the posterior draws would be the mode of the posterior.
Gautam: I’ll check out the citations.
Aki: Thanks for the tips on dimensionality. I think this paper’s just using a single fixed h. The problem with multivariate normal mixtures in high dimensions is that performance is cubic in dimension times the number of components. The nice part about Zhong’s method here is that it’s linear in dimensionality per point selected times B points. But if B has to grow so fast that we get B >> N^3, then that’s not so helpful.
Elliot Morris and Jeff: I think I’m going to need to understand the method better before recommending it. Evaluating this method and coming up with recommendations would make a nice MS or undergraduate thesis.
Interesting to see all those approaches to make Bayes sequential. In biostatistics the use of meta-analytic-predictive priors are used quite a bit with the aim that an informative prior reduces the sample size in a clinical trial. The idea is to use a meta-analysis to predict parameters in a future study. In order to pre-specify now the design it is common to approximate these MCMC posteriors using mixtures of normals/betas/gammas (depending on the endpoint).
These distributions are usually uni-modal and relatively nice. Fitting the mixtures works well with EM, see https://opensource.nibr.com/RBesT/articles/introduction.html for a simple example.
By now this technique has been expanded to multi-variate normal mixtures… and the best thing is that these multivariate normal mixtures can be plugged as priors into brms model without any sweat, see an example here:
https://opensource.nibr.com/bamdd/src/02ac_meta_analysis_strata.html
The function enabling this is the mixstanvar function from RBesT https://opensource.nibr.com/RBesT/reference/mixstanvar.html … this stuff is super useful for biostats and I would love to see non-biostats applications!
Thanks for reading our paper and the post. I want to raise a few points.
Regarding time complexity:
The complexity for implementing one step of the brute force algorithm is order B where B is the number of prior samples. On the other hand, the complexity for one step of our algorithm is constant order. In practical applications where B is typically large, our new algorithm will run much faster than the brute force algorithm.
Regarding comparison of running time:
I do not think the comparison presented in the blog is a fair comparison. Our paper is focused on introducing a novel methodology and framework. We aimed to tackle the question of how to perform efficient posterior sampling when the priors are estimated, even when the size of prior samples is large.
The paper is not focused on how to optimize the code implementation. In our paper, both the brute force method and the new method are implemented using R, thus allowing for fair comparison. In the coding example presented in the blog, the algorithm is implemented using stan, which is based on C++ and is already heavily optimized. We are comparing apples to oranges if we are just comparing heavily numerically optimized C code to plain R code.
See posts here: https://stackoverflow.com/questions/13390152/how-much-faster-is-c-than-r-in-practice and https://darrenjw.wordpress.com/2011/07/16/gibbs-sampler-in-various-languages-revisited/. In the second link, the author benchmarked the same algorithmic formulation of Gibbs sampler in various languages, and C-based code showed about ~55 times speedup than R, despite both implementations using the same methodological pipeline.
In principle, if we optimize the code and use C++, we expect our algorithm to be much faster than based on R; the point is that optimizing code is not the focus of our paper. For example, besides the computation of the Metropolis-Hasting ratio (a similar computation of order B is also involved in the stan example), the brute force Metropolis-Hastings algorithm only involves sampling one Gaussian variable, which involves less computation than HMC and should be very fast to implement once the code is optimized.