Hamiltonian Monte Carlo Inference of Marginalized Linear Mixed-Effects Models

Summary

This work considers Bayesian inference via MCMC methods (specifically: Hamiltonian Monte Carlo) in linear mixed-effects models (LMMs). Due to linearity and Gaussianity, the random effects can be analytically integrated out so that the MCMC algorithm needs to only target a distribution on a smaller space which does not include the random effects. This typically leads to more efficient MCMC algorithms. However, as the authors show, a naive marginalisation is costly, incurring a complexity that is cubed in the number of observations and the size of the vector of random effects. The main contribution is to show that, by exploiting sparsity, the random effects can be integrated out in a way that incurs (typically) the same computational complexity as evaluating the likelihood without any marginalization. The authors demonstrate that marginalization in this way improves the performance of the HMC algorithm compared with "no marginalization" both in terms of the effective sample size and effective sample size per second.

Strengths

1. I think this work makes a good contribution to the field: The proposed implementation of the marginalization (especially when applied to all random effects) does seem to greatly enhance the efficiency of the HMC algorithm in some applications and, as the authors prove rigorously, this does not come at the cost of an increased computational complexity. 2. The presentation is good and the writing is clear. 3. There is sufficient empirical demonstration, in my view.

Weaknesses

**Title** I don't think the title (and the emphasis on HMC in the abstract and introduction) is entirely adequate. The wording "marginalized Hamiltonian Monte Carlo" implies to me that some component specific to HMC (e.g. the momentum variable) is analytically integrated out. However, the marginalization here is completely orthogonal to HMC. Indeed, it could be exploited within any other MCMC algorithm or even within other Monte Carlo schemes like sequential Monte Carlo samplers). **Presentation** Overall, I think the presentation is good. Just a few minor comments: 1. In Table 1 and Figures 2, 3 and 4, referring to the "no marginalization" strategy as "HMC" is confusing since all methods use HMC. This is also inconsistent with, e.g., Table 2. Better replace "HMC" by "no marginalization". 2. A number of symbols are not or not clearly defined, e.g., on Page 2, define $\mathbf{u}_j$, $\mathbf{u}$ (is only defined much later), and $k$ (should $k$ be $M$?). More generally, the fact that symbols without the subscript $i$ denote the collection over all $i$ should be mentioned (unless I missed it). 3. From the context, it is clear what the authors mean by "recovery" and "ancestral sampling" but unless these terms are somehow standard in the literature, I think they should be more clearly explained to avoid confusion. 4. Perhaps the distinction between "classifications" and "groups" (around Lines 88 to 92) could be made more clear, especially because the examples in parentheses say "subject", "gender", and "age" in both cases which is a bit confusing. 5. I couldn't quite understand Section 6.3. What does "marginalization in a vectorized way" mean? 6. Check the bibliography; it has many mistakes, e.g., missing capital letters in proper nouns and journal names. **Other** * L241: document -> documentation

Questions

1. Can this be extended to multivariate observations? 2. As the authors mention, the "no marginalization" implementation seems to be slower than the proposed implementation that performs (some) marginalization. This seems surprising. Do the authors have an explanation for this?

Rating

6

Confidence

3

Soundness

4

Presentation

3

Contribution

3

Limitations

I don't foresee any negative societal impact.

Summary

The paper addresses the challenge of marginalizing latent variables in linear mixed-effects models (LMMs) to improve Hamiltonian Monte Carlo (HMC) sampling. The authors propose an algorithm for automatic marginalization of random effects in LMMs, leveraging fast linear algebra techniques. The experiments evaluate the performance of the proposed marginalisation on LMMs from several disciplines by using the No-U-Turn sampler in NumPyro and assessing performance via effective sample size (ESS), number of divergent transitions, and running time.

Strengths

This work focuses on an important problem in Bayesian modelling and inference: specifically the automatic marginalization of continuous latent variables in probabilistic models. While prior work exists that explores such approaches for discrete random variables, this is less so the case for continuous variables, so I believe the contribution is valuable to the community. To the best of my knowledge this contribution is original, as I am not familiar with other work that explores such continuous variables automatic marginalization for probabilistic programming. I also thought the idea that marginalization, similarly to reparameterization, can alleviate difficult characteristics such as funnels in HMC is well articulated and personally it gave me a new perspective.

Weaknesses

My two main concerns about the paper are regarding the experiments and about the significance of the results. ### Experimental Details: Overall, I find that the experiments lack sufficient details and metrics: * There is not enough information about hyperparameters used for each method and model, e.g. the step size for NUTS. Does the approach rely solely on NumPyro's adaptation routine? If yes, what parameters were used for that? * The evaluation can use additional diagnostics and metrics to provide more evidence that marginalization improves inference: for example trace plots, and visualisation of divergences can be very useful. Effective sample size alone is not enough for properly assessing the quality of inference. * I don’t think the number of divergences is a reliable metric, without additional context, such as their location in the posterior distribution or the step size used. In the case with section 6.2, the divergences are not neccesarily indicative of any problems. In my opinion, the 12 divergences observed for R1, R2 are just 0.012% of all drawn samples, which is negligible. Where are those divergences located in the posterior plot? Are they grouped in one place? * The choice of priors for various models seems to me quite broad and lacks justification. ### Significance of Results The claim that marginalization is always beneficial is not novel, as it follows from the Rao-Blackwell theorem. The paper mentions Rao-Blackwellization, but doesn’t go into details. Edit: the rebuttal clarified on all of these points.

Questions

1. Can you provide more detailed convergence diagnostics for your experiments, such as trace plots, for example? 2. Are divergences grouped in specific regions of the posterior distribution? 3. 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? 4. Could you elaborate on the choice of priors in your experiments? Why were these particular priors selected? Edit: the rebuttal has answered all my questions.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

4

Limitations

The last section of the paper addresses a few limitations of the approach, including addressing distributions beyon normal. However, the paper can benefit of further discussion on one, in my opinion, fundament limitation of this approch, which is that it only considers cases where analytical marginalization is possible. How many practical cases are covered by this constraint?

Summary

The paper studies marginalization in linear mixed-effects models with a Gaussian likelihood. Marginalizing out variables is a common trick to improve the convergence of MCMC. As pointed out, for the linear mixed-effects models with a Gaussian likelihood, naively doing so will actually increase the computation as there are a few linear algebra operations that is costly. The paper introduces two tricks using the matrix inversion lemma to speed them up. Some experiments are performed to demonstrate the effectiveness and to compare against the reparameterization alternatives.

Strengths

## originality The paper presents a simple improvement to a known problem. ## quality The technique presented in the paper is relatively simple and straightforward. Such techniques are widely used in Gaussian processes community. ## clarity The paper content is easy to follow however the title is a bit misleading as the technique is not specific to HMC but generally applicable to all MCMC algorithms. ## significance The model family studied in this paper is pretty narrowed and even in this scope it's unclear how the presented method works compared to other automatic marginalized methods.

Weaknesses

## Models studied are narrowed Only models with Gaussian or log-Gaussian likelihood is studied. ## Experiments are not enough 1. In the experiments there is only naive baselines. Other alternative automatic marginalization methods are discussed in the related work section but not compared against. 2. Only 1 dataset (section 6.1) and 1 toy target (section 6.2) are studied.

Questions

Is the method specific to HMC?

Rating

5

Confidence

4

Soundness

3

Presentation

2

Contribution

1

Limitations

Limitations are discussed in section 7 however I do think those limitations make the paper a bit weak (not enough contributions) for acceptance.

Summary

This paper considers the problem of Bayesian inference in linear mixed-effects models. This is a popular class of models used throughout many areas of application. Inference in such models depends on being able to determine values of latent variables (fixed and observed effects) given the observed data. This is typically done by writing a model of interest in a probabilistic programming language and then running an automated inference algorithm on it. For example, the Stan package runs a variant of Hamiltonian Monte Carlo on the model. In general, this can be a complex problem from a sampling point of view. The main contribution of this paper is to speed up vectorized random effects marginalization, which simplifies the sampling problem and leads to faster sampling. The authors note that marginalization tends to always be beneficial, and illustrate their proposed methodology on a well-chosen set of real-life examples.

Strengths

I like the clear motivation and formulation of the problem, including the challenges and limitations of the proposed methodology. The problem is of obvious interest and practical relevance and would be of interest to a wide audience. The presented solution is clean and straightforward in its use of determinant computation and matrix inversion lemma. The experimental results are convincing and of interest.

Weaknesses

The proposed methods are limited to Gaussian models, which the authors note. I would have appreciated some more commentary in this direction. The marginalization approach seems sufficiently interesting to be used widely, whenever possible.

Questions

What sort of prior structures allow for this to be applied? Are there methods that can be developed that perform approximate marginalization in some way? How hard is it to automate marginalization in practice?

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes.