I don't think the title (and the emphasis on HMC in the abstract and introduction) is entirely adequate.

Thanks for this comment. Please see the general response. Briefly, we propose to change the title and revise the introduction to clarify that the method is generally applicable, but mainly focus on HMC for the scope of the paper for the reasons described in the general response..

... Better replace "HMC" by "no marginalization". A number of symbols are not or not clearly defined ... "recovery" and "ancestral sampling" ... should be more clearly explained ... the distinction between "classifications" and "groups" could be made more clear ... Check the bibliography

We appreciate your meticulous efforts in reviewing our submission. We will make a complete revision of this manuscript to incorporate your suggestions.

What does "marginalization in a vectorized way" mean?

We refer to our method as vectorized marginalization, in contrast to the automatic marginalization in [1], which works only with scalar random variables. Compared to that baseline, we achieve the same goal of marginalization, but with complete vectorization in the computation. As far as we know, no existing works in automatic marginalization solve LMMs with vectorization.

Can this be extended to multivariate observations?

For multivariate observations, sparsity still exists in the model structure and our method should work with no issue.

... "no marginalization" ... slower than ... marginalization. ... Do the authors have an explanation for this?

There are two benefits from marginalization that may improve the sampling speed. The first is the reduced state dimension $H$, which leads to less discretization error for the numerical integrators and thus larger step size. The second is better geometry. There are certain structures like funnels that make adaptive HMC converge to a small step size. Such structures can be removed by marginalization, which makes HMC more efficient.

[1] Lai, J., Burroni, J., Guan, H., & Sheldon, D. (2023, July). Automatically marginalized MCMC in probabilistic programming. In International Conference on Machine Learning (pp. 18301-18318). PMLR.

Thank you for the questions about the experiments. We respond to your questions below, where we will reference additional results uploaded in the PDF response. Broadly, these additional results show additional metrics and experiments consistent with our original findings that marginalization improves inference.

Can you provide more detailed convergence diagnostics for your experiments, such as trace plots, for example?

In the additional PDF, we have included trace plots for the ETH instructor evaluation model as Figure 1. We will include more trace plots to analyze the convergence in the next version of the paper. We also compute $\hat{R}$ for shorter chains on the model and present the results in Table 1 of the additional PDF. The results show that proper marginalization is effective in improving the efficiency of HMC for the model with a denser trace plot and better R-hat values.

Are divergences grouped in specific regions of the posterior distribution?

We notice the divergences of R1,R2 are related to specific combinations of values for $\sigma_2$ and some of the parameters in $**u**_2$. Figure 3 in the additional PDF shows the locations of divergences. Also, the divergence numbers vary with different random seeds. While M2 and M2,R1 never report divergence, the numbers of divergences for R1,R2 are 12, 12, 528, 1373, 233 out of 10,000 samples for the first five random seeds. See Table 2 in the additional PDF for more details. We will include those results of divergence in Section 6.2.

Does the approach rely solely on NumPyro's adaptation routine? … What is the impact of varying HMC parameters (e.g., step size, or adaptation routine parameters) on the number of divergences and the effective sample size?

We mainly use the default adaptive routine of NUTS in NumPyro, which has an adaptive step size with dual averaging, adaptive and diagonal mass matrix, target acceptance probability of 0.8 and maximum tree depth of 10. For the ETH instructor evaluation model, hardness is observed without marginalization, so we set the maximum tree depth to 12. It is possible to reduce the step size by increasing the target acceptance probability. We have tested with different target acceptance probabilities. On the ETH instructor evaluation model, we find benefits of marginalization are consistent with a different choice of the target acceptance probability (Figure 2 in the additional PDF). Notibally, marginalizing $**u**_1$ may not be beneficial in ESS in the original setting, but there are improvements when the target acceptance probability is 0.9. On the grouse ticks model, divergences still exist after increasing the target acceptance probabilities to 0.95 (Table 2 in the additional PDF).

Could you elaborate on the choice of priors in your experiments? Why were these particular priors selected?

For datasets with existing models such as those in [1], we follow their choices of the priors. For other models, we use generic and weakly informative priors [2]. Our conclusions are orthogonal to the choice of priors.

The paper mentions Rao-Blackwellization, but doesn’t go into details.

Thanks for this question. We will expand on the relationship to Rao-Blackwellization in the paper. In short, Rao-Blackwellization reduces the variance of an estimator for some quantity by taking the conditional expectation over some variables and sampling the remaining variables. In this paper, on the other hand, we improve the speed of obtaining samples from the remaining variables (by improving mixing times and/or reducing the cost per iteration). So while Rao-Blackwellization is spiritually similar, the Rao-Blackwell theorem does not apply to our situation. The point of our method is to obtain samples from the non-marginalized variables faster, not to integrate over the marginalized variables in some expectation. If one were interested in an expectation (of all the variables) one could apply Rao-Blackwellization on top of the samples returned by our method.

More formally, if one is interested in some expectation $\mathbb{E}[f(X,Y)]$ one can view Rao-Blackwellization as transforming the estimator as

$$\frac{1}{N} \sum_{i=1}^N f(X_i, Y_i) \to \frac{1}{N} \sum_{i=1}^N \mathbb{E}[ f(X,Y) | X=X_i],$$

where $N$ is the number of samples (if the samples are coming from MCMC, essentially the same equation holds, with $N$ being the effective sample size).

In our paper, by marginalizing the $Y$ variables we aim to have a more effective MCMC method for $p(X,Y)$. Then, for each sample $X_i$, we recover $Y_i$ via exact sampling from $p(Y_i | X_i)$. So, essentially, we would improve the expectation of some function through the transformation

$$\frac{1}{N} \sum_{i=1}^N f(X_i, Y_i) \to \frac{1}{M} \sum_{i=1}^M f(X_i, Y_i),$$

where $M > N$ is a larger number of samples. That is, the improvement simply comes from more (effective) samples, due to faster mixing and/or more iterations per second.

Finally since the distribution of $Y$ given $X$ is tractable, one could imagine applying both our technique and Rao-Blackwellization, with the transformation

$$\frac{1}{N} \sum_{i=1}^N f(X_i, Y_i) \to \frac{1}{M} \sum_{i=1}^M \mathbb{E}[f(X,Y) | X=X_i],$$

where again $M>N$. (In general it might be necessary to use ancestral sampling from $p(Y|X)$ to estimate the inner expectation, but in many cases this should be quite cheap.) This would be interesting to explore more fully in future work.

[1] Nicenboim, B., Schad, D., & Vasishth, S. (2021). An introduction to Bayesian data analysis for cognitive science. Under contract with Chapman and Hall/CRC statistics in the social and behavioral sciences series.

[2] Andrew Gelman et al. (2024) Prior Choice Recommendations.

Dear authors, thank you for your response. I appreciate the extensive rebuttal and the additional results you presented in the attached pdf. Thank you, in particular, for the plots visualising where divergences happen - I think these really help highlight that there is indeed a problem with the geometry that your method seems to address.

This has addressed all major concerns I had and I will increase my score accordingly. Looking forward to seeing the updated version of the paper when available.

Overall, reviewer UJHx suggests that similar techniques are known or standard and faults the paper for not comparing to them, but provides no details or citations. It is impossible to accept these claims without justification. We respond to specific comments below.

Such techniques are widely used in Gaussian processes community.

Although the first step is applying the matrix lemmas, we use additional structures specific to LMMs afterward to obtain a fast running time. This goes beyond what is commonly done in GPs.

Only models with Gaussian or log-Gaussian likelihood is studied.

We have discussed the scope in our general responses. One point of the paper is that normal marginalization covers many cases with continuous observations, given proper transformation functions such as the log function.

In the experiments there is only naive baselines. Other alternative automatic marginalization methods are discussed in the related work section but not compared against.

Automatic marginalization methods exist in the literature, but all except [1] have very different scopes and don’t apply to our setting, e.g., they marginalize discrete variables only or are tailored to sequential Monte Carlo or algebraic exact inference algorithms. The comparison with [1] is in Section 6.3. We are not aware of other baselines that are applicable to the models we consider. Reviewer UJHx provides no details or citations to claim that baselines are missing.

Only 1 dataset (section 6.1) and 1 toy target (section 6.2) are studied.

Reviewer UJHx counts 2 datasets but in our experiments, there are in total 13 datasets from Section 6.1 to Section 6.4.

Is the method specific to HMC?

Our method works for any MCMC algorithm that accepts a black-box model. We highlight using HMC from practical perspectives. See the general responses.

[1] Lai, J., Burroni, J., Guan, H., & Sheldon, D. (2023, July). Automatically marginalized MCMC in probabilistic programming. In International Conference on Machine Learning (pp. 18301-18318). PMLR.

What sort of prior structures allow for this to be applied?

In the context of LMMs, the distributions of the effects are assumed to be normal. We do not restrict the priors of other variables.

Are there methods that can be developed that perform approximate marginalization in some way?

In the general response, we briefly discussed one direction of generalization to models with discrete likelihoods via data augmentation. We think approximation, such as Laplace approximation, is another way to generalize exact marginalization. A pioneering work in this direction is [1]. However, the approximation leads to biased MCMC algorithms. Pseudo-marginalization could be used to correct the bias, but in our preliminary work, the computational advantages and tradeoffs are quite unclear; for example, embedding Newton-based optimization for Laplace’s method inside pseudo-marginal HMC required the computation of third order derivatives.

How hard is it to automate marginalization in practice?

There are two levels of automatic marginalization. The first is to marginalize given high level abstractions of the model, such as R formulas. It is possible to compile the abstractions into concrete models (e.g., probabilistic programs) while marginalizing one or more of the effects using our methods. The second is to marginalize given user-written probabilistic programs of LMMs. In such a case, some compilation or tracing technique is needed to convert the user’s program to a model representation suitable for manipulation. For example, the authors of [2] used program tracing to construct a graphical model representation that could be programmatically analyzed and transformed. To apply this methodology to LMMs, a special parser would also be needed to match the models to LMMs. We think these are interesting extensions of our work in the aspect of probabilistic programming languages.

[1] Margossian, C., Vehtari, A., Simpson, D., & Agrawal, R. (2020). Hamiltonian Monte Carlo using an adjoint-differentiated Laplace approximation: Bayesian inference for latent Gaussian models and beyond. Advances in Neural Information Processing Systems, 33, 9086-9097.

[2] Lai, J., Burroni, J., Guan, H., & Sheldon, D. (2023, July). Automatically marginalized MCMC in probabilistic programming. In International Conference on Machine Learning (pp. 18301-18318). PMLR.