Hamiltonian Monte Carlo Inference of Marginalized Linear Mixed-Effects Models
Authors: Jinlin Lai
Justin Domke
Daniel Sheldon
Venue: NeurIPS 2024 poster (2024)
Keywords: Hamiltonian Monte Carlo
Bayesian inference
hierarchical models
Primary area: probabilistic_methods

Abstract: Bayesian reasoning in linear mixed-effects models (LMMs) is challenging and often requires advanced sampling techniques like Markov chain Monte Carlo (MCMC).
A common approach is to write the model in a probabilistic programming language and then sample via Hamiltonian Monte Carlo (HMC).
However, there are many ways a user can transform a model that make inference more or less efficient.
In particular, marginalizing some variables can greatly improve inference but is difficult for users to do manually.
We develop an algorithm to easily marginalize random effects in LMMs.
A naive approach introduces cubic time operations within an inference algorithm like HMC, but we reduce the running time to linear using fast linear algebra techniques.
We show that marginalization is always beneficial when applicable and highlight improvements in various models, especially ones from cognitive sciences.

Hamiltonian Monte Carlo Inference of Marginalized Linear Mixed-Effects Models
Jinlin Lai, Justin Domke, Daniel Sheldon
Manning College of Information and Computer Sciences
University of Massachusetts Amherst [jimiinlai, domke, sheldon) lecs.umass.edu
Abstract
Bayesian reasoning in linear mixed-effects models (LMMs) is challenging and often requires advanced sampling techniques like Markov chain Monte Carlo (MCMC). A common approach is to write the model in probabilistic programming language and then sample via Hamiltonian Monte Carlo (HMC). However, there are many ways user can transform a model that make inference more or less efficient. In particular, marginalizing some variables difficult for users to do manually. We develop an algorithm to easily marginalize random effects LMMs. A naive approach introduces cubic time operations within an inference algorithm like HMC, but we reduce the running time to linear using fast linear algebra techniques. We show that marginalization always beneficial when applicable and highlight improvements in various models, especially ones from cognitive sciences
greatly improve inference but is
Introduction
Bayesian hierarchical models account for complicated relationships in data by introducing hierarchical structures [23]. Among hierarchical models, linear mixed effects models (LMMs) are widely used in various scientific disciplines, including ecology B1, medicine (7]. psychology 41], neuroscience [77] and cognitive science [47). Solving LMMs involves inferring latent variables, such as fixed and random effects, based on the observed data. Fixed effects are shared by all observations, while random effects vary across different groups within the data. LMMs are often implemented using probabilistic programming languages (PPLs), which isolate inference from modeling: users write program representing the model and the PPL automatically executes suitable inference algorithm. Variants of Hamiltonian Monte Carlo (HMC) [5] are dominant in many PPLs today and are widely used for LMMs. For example, BRMS 8 is influential R package that allows users to write regression- -tyll formulas that are automatically translated to Stan programs (9] representing an LMM, and then Stan's 's HMC implementation is called to generate posterior samples.
We develop techniques that allow users to easily transform their models to analytically marginalize random effect variables from LMMs to improve the efficiency of HMC. Marginalization has several benefits. First, there are often pathologies in LMMs that hinder efficient HMC sampling. A notable one is the "funnel" shape created by correlation between variance parameters and parameters for fixed or random effects 45) Marginalization
26 have been shown to be useful in addressing such pathologies. Second, marginalization reduces the number H of
( and other program transformations
latent variables for HMC. The complexity of HMC is about 5(544) II 46). so it is desirable to run HMC on a subset of variables if marginalization can be done efficiently. Our methods enable marginalization of random effects in LMMs with linear Gaussian structure, which includes models with normal and log-normal likelihoods as well as other likelihoods for continuous data based on
'The code is available
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

transforming a normal distribution. Note that our methods are not limited to HMC, and could be applied to many inference algorithms.
There are several challenges to efficient marginalization. The automatic marginalization algorithm of 35 can be applied to L.Mss but is limited to scalar random variables, it requires construct the LMM as graphical model with separate variables for each effect and observation. Another alternative is model the relationships between effects and observations with a design matrix and marginalize effects using properties of multivariate normal distributions. We call this the "vectorized approach" since it can leverage vectorization accelerate computations. Unfortunately, vectorized marginalization leads to dense covariance matrix over the observations and thus cubic time for evaluating the log-densitg within HMC, when the log-density of the original could be evaluated in linear time. Our main technical contribution is to accelerate vectorized marginalization for L.Ms using fast linear algehra: we show that marginalization for single random effect can be achieved with linear time complexity and can significantly accelerate HMC compared to both the original model and non-vectorized marginalization.
We implement vectorized marginalization for LMMs in NumPyro S 54] via simple classes users can use express their models. We evaluate our approach on variety of real LMMs from past scientific investigations, including nine models and datasets from cognitive sciences, and find that marginalization is always beneficial. Our findings suggest that practitioners should marginalize group-level effects whenever applicable in Bayesian hicrarchical inference.
2 Background
To motivate our problem, we present an example model. In 72). a set of experiments were run to examine the relationship between human pupil and attention load. A total of N 2228 measurements of pupil sizes from M = 20 subjects were taken under different attention load levels. Specifically, in the ith measurement, the pupil size Y1 E IR+ of subject 9j {1, ..., k) under attention load o; (0. 1, 2, 3, 4, 5} was recorded. Pupil size can be assumed to have linear relationship y; F fg + 0 with respect the attention load C, where both the slope 0, and intercept Og split into fixed and random effects:
y = ug c3(3 + ug.2) tG, t e~ c~,,00'").
where a, 8 are variables for fixed effects and u., are variables for Bayesian hierarchical modeling assigns priors to each unknown variable:
random effects.
(10000 500²), 8 ~ (N 100),
T ~~ (0, diag (1000 1000²)),
L, LKJCholesky(2,1) 1), [uj,1, u,,2] ~ (o TL.L.T), 1,2, ...
half-normal distribution (N*) and an LKJ distribution (LKJCholesky) (36 are used as prior on the covariance matrix. Inference for the unknown parameters determining the relationship between pupil size and attention load can be performed by writing a probabilistic program and running HMC. For example, in NumPyro, the regression model for all measurements may be implemented as below.
numpyro sample ''y' "oyy''"aaa d.ist.Normaa (alpha-u [g] [: 0] +c* (beta+u[g) [: 1]) si sigma) obs-y)
The code above uses advanced indexing and vectorization techniques in numpy, where u, u,g, c,y all vectors or matrices. We further observe that, conditioned on a, 3, a. T, Lu., the distribution of all uj and all y, form multivariate normal distribution. Theoretically it is possible analytically integrate u out from the model to improve inference efficiency. But is not straightforward for users will see, if done in the most obvious way,
to transform the probabilistic program do so, and, may not make the model more efficient for HMC.
To be more clear about how marginalization can be implemented, rearrange the model into a canonical form that focuses on the random effects. All observations and random effects into the vector u
u1
112
collected into the vector y [V/1,.., .... HN [u1,1. U1 44,1, u4 2)
21
Y22
1/3
Then, write
u ~ N(H, y M(Au b, Ey).
Figure 1: A tree-structured model conditioned on e.

---

where Je, Eu. A, Ey are functions of a. 3, T, Lu. 9r. C.. Note that ys only depends on the entry Ug of u. The corresponding graphical model has tree structure, as demonstrated in Figure [ This tree structure has several benefits: first, matrix multiplications like Au and A y can be done efficiently; second, we will see that it leads to block-diagonal structure that facilitates efficient inversion a key matrix that appears later.
For more general LMMs with more than one class of random effects we generalize the canonical form as
+ ~ p(), n (() i= 1= L
En:
y/O, un U2 up N/
Σ (,)u), + b(O), My E,,))
(I)
where p(O) is the distribution for global variables (including fixed effects), p(u; ( is the distribution for random effects and p(y|@, U1; Ur.) is the distribution for observations. Notationally this generalization further adds an index to each random effect to specify its class. A user might specify the model directly in this canonical form, or in another syntax (c.g.. the formula syntax of BRMS) that is compiled to this form. Each pair (u;, A.) specifies a class of random effects for particular classification of the observations (..g.. by subject, age, gender, etc.). Each classification contains multiple groups and different classifications are distinct from one another. Each observation belongs to one group for each classification. The vector u, [u? 2 u,,,,1 1" contains random effects for the ith classification (c.g., subject, age, or gender). consisting of k, groups (c.g.. one subject, age, or gender), with U containing the random effects (c.g., slope and intercept) for the jth group. We denote the number of observations dim(y) = .N, and the number of random effects per group as dim dim(u,,)) = d. Any sovariates-such as c; in the pupil size example-are considered constants and not represented in the notation. In LMMs, the number is related to the number of covariates and is usually small. The total number of random effects for u, is denoted as dim = M, k,d. The matrix A, therefore has size N x M,, and encodes the group structure for u, by mapping random effects (together with covariates) to observations. Each row of A, encodes the assignment of an observation to one group, so it has at most d nonzero elements. Therefore, the complexity of computing A; is (Nd)) as A has at most Nd nonzero elements. Henceforth, we omit the dependence on + for μa, Eu. A, b, My for simplicity.
Marginalizing u; is possible to analytically marginalize variables in this model: since the mean o y is linear in each u, and all of these variables are normally distributed, the joint distribution of (y, up...., uz) also multivariate normal. We will focus for most of the paper on marginalizing the random effects u, for single i in order to leverage the tree structure mentioned earlier, but return in Section 4 to the idea of marginalizing many effects. Locally, u; and y form the conditional distri- p(u, θ p(y(( u-,, u:). Marginalized MCMC rewrites this conditional distribution as p(u, y O u.) ply O u-:))((,/()) y: u_;), which reverses the dependence between u; and y 35 During sampling,
bution p( U1: ② u-;)
is marginalized from the HMC procedure by using p(y O u as the likelihood function and p(O) u.:) as the distribution of latent variables. After HMC sampling, u; is recovered through ancestral sampling from p(u u, O, y, u-;) given posterior samples of (o, u..) The reversal requires analytical forms of ply O u-i) and p(u, O y, u.... which can be obtained via standard marginalization and conditioning operations on multivariate normal distributions [e.g..[]]
100,ui~M
Σ Aj + A.p. + +b, + Ey jti
n 1/O,y,u ~ M
M++M y - (I- (I-MA,)Eu: 1t1
(2)
where M =
+y)) Marginalization introduces the benefit of sampling in a lower dimensional space, but the cost depends on the complexity of evaluating the log-density functions of these two distributions in order to run HMC.
-1,
3

---

Table 1: Time complexities of different HMC approaches for the submodel involved in marginaliza- tion. Initialization done once before the HMC loop. The log density is computed within each step of the leapfrog integrator. Recovery is performed for each sample from HMC. N is the number of observations, M the dimension for one class of random effects, D is the dimension for all classes of random effects, L is the number of classes, d is the dimension for an effect of a group in a class.
Submoddl
Approach
Intialization
density
Recewery
No marginalization
O(M& NTd)
p(u,: (1,,y)
Naive marginalizaaion
O(M 4 N²)
O Nº) O(Ma? + NLd Nw²)
Marginalize with lemmas
O(Ma + NLd + Nd")
No marginalization
O(Di + NLd)
p(v,y(e)
Naive mangimaliiaaiion
O N
0 No) O(D" + NLd)
Marginalize with assumpiions
O + NN
O(D + ld)
2.1 Challenges of multivariate marginalization
In practice, the original model usually has structure that makes evaluating its density very efficient, which is lost by naive marginalization. For example, the observations in y are usually conditionally independent, making Σy diagonal; also, En. is usually block diagonal with blocks of size d x d. So evaluating the density p( u, y yrO, u = p( u, p))p( O u.:L) requires Ok,dd + Ld) = OM,dd + N Ld) time with the main operations being (1) inverting and computing the determinant of En and Σy: (2) computing the mean parameter of y. When En; is diagonal, the complexity goes down ((,d Ld). However, it more expensive to evaluate the density of the reversed model in Equation 2 Computing p(y O u_i) and p(u, O.yy u-;) requires the inverting and computing Σy. which we denote by E for simplicity. For
the determinant of the N x N matrix A, A
the log likelihood, we need to compute log p(y O u-:) = 1 det (E) 1 E + C, where z = y Σjzz A,u; A. b. E is not diagonal and without using additional structure will trigger O(N3) operations within each step of the leapfrog integrator within HMC. For the recovery distribution p(u, e, E will be inverted when calculating M. Also, a Cholesky decomposition for the covariance (I MA,) should be computed for sampling, which takes O(M)) time. These cubic time operations are prohibitively expensive for large datasets. We summarize the complexities of different approaches in Table [ In Section 3 we discuss how to marginalize one group of random effects with lemmas from linear algebra. In Section 4 we discuss how to marginalize all random effects with additional assumptions.
3 Marginalization with fast linear algebra
We now show how to speed up calculations with the marginalized model using fast linear algebra methods. In particular, we use the matrix inversion lemma and matrix determinant lemma together with special structure in the relevant matrices. In this section, we sometimes omit the subscript such as for A, and Σ,; for simplicity. The steps in log density evaluation and recovery are summarized in Algorithm[] []] and in Algorithm in the appendix, with comments about their implementation and cost. We mainly use sparsity and tree-structure in A to make operations faster. As an overview, computing takes O(N Ld) time for L sparse matrix multiplications of time ((Nd) each. Also, evaluating As and Tt both take ((Nd) for any E M and any t E N... With tree-structure, we will see that 1 A is block-diagonal and can be computed efficiently.
A Ey
3.1
Matrix inversion and determinant lemmas in marginalization
The two main bottlenecks when evaluating log p(y O u-s) are computing det and 12. With the matrix determinant lemma (3). we have that
det (E) = (At(AA,,4 + Σy) = det ( + A' A) det(E Eu) dot(Eyy)
(3)
By the matrix inversion lemma or the Woodbury formula 53 we have that
E 1 (AΣ,A"
-
A A
A) -1A'E
Therefore,
T " E'E''z 15-1=220 2" 2''''' -2--2- ,'''' + + A) A
(4)

---

Algorithm Evaluating p u--)). Each A; is an N x M, sparse matrix with Nd elements and tree structure. Σy is N N diagonal. En M x M block-diagonal with block size d.
I: z = y Σjグs u A:p b
Sparse matrix multiplication in (N time Block diagonal computation in O(M + NN)d) time Sparse matrix multiplication in ((N)) time > Determinants in ((Md)) time
2: F Σ, + A by Ly
3: x A Syyy
4: log dct(F) + +og (Eu) + (oggottyy))
5: b z" y z F F'F-xx
Quadratic form in O Md) time
6: return (a + b) +
By using the facts that Eu is block-diagonal, Σy is diagonal, and A has Nd nonzero elements, the quantities det(Eu), det (Ey), z² Eyy'" z, and A A" " can each be calculated in O + NNd) time. Equations (3) and 44) contain the expressions F-1 or det(F) for the M x M matrix F = En A My A, which both require ((M)) time when done naively. The following theorem shows that these quantities can be computed in O((M + N)d) ) for LMMs.
Theorem 1. If E is diagonal. Eu is block-diagon with blocks of size d x d. then F = E,' A A A is also diagooaa with d blocks and computing A AT A takes ((Nd)).
Proof. The proof uses the tree-structure in A. For details, see Appendix B.
Therefore, it is O(MM + N)d) to compute det(F) and F-1. Combined with other parts in the formulas, the overall complexity is O(M + Nd + Ndf). In LMMs, d ís usually small, so the complexity with marginalization can be viewed as the same as the complexity without marginalization.
3.2 Speeding up the recovery step
Different from evaluating log ply|O O u-i), ancestral sampling from p(u, O, u-:) is only per- formed once for each posterior sample. When sampling from p(u O. u.y, u.)). :), computing M directly is also costly. With the matrix inversion lemma, we have
M = E.A T (AE,A" + Ey)-
= E E. E A 1 (''','A)--ATA
(5)
With this expression, the mean variable fa + Mz, then is evaluated in O(M + ), by computing Eg 2-1 +A" A in the same way Line 2 of Algorithm we have from substituting Equation 5 that
For the covariance variable (I- (I-M)) MA)E..
(I - M.) = En Σ,, A Σy AE, + E, A Ey A + + - A A AE
EE A ",,,, +
+A'
Note that E. ATE,'A and F - Σ AT A adding and multiplying them also block diagonal. For block diagonal matrix with blocks of size d x d, the time complexity for Cholesky decomposition is O O(Md). Combined with the complexity of computing z, the recovery step takes O( + NLd Nd)) time.
'A are all block diagonal, so the result of
Marginalizing multiple effects with additional assumptions
We have shown that it is efficient to marginalize one class of random effects. With additional practical assumptions, is possible to marginalize all classes of random effects for efficient HMC inference. Instead of separating different classes of random effects, LMMs can also be written as Ni,, Ev); y N(vv + b, Ey), where define that D = ΣL M,. The matrix inversion and determinant lemmas can still be applied to marginalize v out, but the combined matrix does not have the special structure of A, we exploited in Section 3 More specifically, the computation of del (F) and the evaluation of F-1 for Σ- + Bd Σ BB both become non-trivial. We introduce additional assumptions to show that they can be solved faster in some special cases. For the general case, the discussion section.
[A,..., A1] and v = [uf..., u.," We
5

---

The assumption we make is that Σy = T.I and Σy TyI. where Tv: Ty are scalars that either belong ④ or are fixed non-random parameters. This means that all effects share the same variance and all observations share the same noise scale. These assumptions are not as restrictive as it may appear. If the underlying distribution is u1 Ni,; II where oj is a fixed parameter, it is possible to reparameterize this distribution as u ~ (0, I). A = o,A,, b' = b + Byg, and use u': A b' in place of u, A,, b. Then Σv becomes a scaled identity matrix. Also, in many models, the noise scale for different observations is the same, making My a scaled identity matrix well.
In practice, if the assumptions are satisfied, marginalization can be done in O(D² + Nd) time with O(D NI2dd) preprocessing. Details are provided in Appendix B.3
5 Related Work
73 a number of other works focus on model transformation. Non-centered parameterization 49) is widely used trick among MCMC users to alleviate slow sampling in difficult posterior distributions. However, there is no general way know whether non-centered parameterization will be beneficial (76). Variationally inferred parameterization (26] proposes to learn a model parameterization from a specified family that will lead to effective sampling. In Parno and Marzouk [52] and Hoffman et al. [33, preconditioners for HMC are learned to transform the model to be approximately isotropic Gaussians. Marginalizatioo differs from reparameterization in that it reduces the problem dimension as well as potentially alleviating difficult characteristics such as funnels, so it has two mechanisms to improve MCMC efficiency. The Laplace approximation (LA) is one way to approximately marginalize variables in
While many works aim to improve HMC directly 771 30 58
MCMC 559 40 655,
but it may be difficult to quantify the error or recover the marginalized variables.
Marginalization, Rao- Rao-Blackwellizzttion has been an important topic in Bayesian inference and probabilistic programming, In Gibbs sampling, marginalization is usually called collapsing (37). Collapsed Gibbs sampling has been developed for latent Dirichlet allocation (56) and LMMs [50].
We explore marginalization in the context of HMC, which induces different considerations. Methods with HMC do not have to make the conditional distributions of the marginalized model tractable. Marginalization is also related to symbolic inference in probabilistic programming. Hakaru 44 and are systems for performing exact Bayesian inference by symbolically marginalizing
PSI [2 22
all latent variables. To marginalize discrete variables, Gorinova et al. [27] propose information
flow type system. Another line of related work is delayed sampling 43 R which automates
marginalization of variables within Rao- Bao-Blackwellized particle filters 42 Lai et al. 35 developed
an automatic system for marginalizing variables in HMC, but is limited to scalar variables so cannot leverage vectorization and forces users to write models with univariate distributions.
Linear algebra tricks have been widely utilized in various machine learning algorithms, such as ridge regression 688 Gaussian processes [6] and Kalman filters [60] Recently, frameworks 62 20 57
have been proposed to ease the implementation of fast linear algebras in machine learning algorithms.
Marginalization in Bayesian models may be an interesting application of those frameworks.
Fast and scalable inference for LMMs has been studied in the context of maximum likelihood estimation 19], variational EM (24), Gibbs sampling [51] and numerical integration (28). We are the first to consider speeding up the inference of LMMs with HMC. There is also a recent trend in integrating random effects into deep neural networks for correlated data [67] or personalization
(66 64 74
with parameters estimated by maximum likelihood.
6 Experiments
We conduct experiments on LMMs from various disciplines using the default no-U-turn sampler (NUTS) [34 from NumPyro 55 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, we set the maximum tree depth to 12 to overcome difficulties performing inference without marginalization in preliminary experiments. For all models, we use weakly informative priors unless specified. In general, our conclusion is insensitive to the choice of hyperparameters and priors. For all experiments, we collect 10,000 warm up samples for tuning, and 1 100,000 samples for evaluation, and evaluate performance via effective sample size (ESS) and running time.

---

Table 2: Running time in seconds for HMC, with or without marginalization. Mean and standard deviation over independent runs are reported. Experiments are run on NVIDIA A40.
Method
No marginalization
Marginalize un
Marginalize U2
Marginalize us
Marginalize u
Time (8)
13417 (98)
5004 (1468)
2607 (3)
3071 (4)
631 (12)
200000
Method merginalization Marginalize un
150000
2
100000
Marginalize wq
50000
Marginalize ua
Morginalize
different sample strategics.
Figure 2: Average ESS for each variable on the instruction evaluation model with different HMC strategies. Numbers above the sample size 100,000 indicate effective sampling.
6.1 Marginalization in cross-effects models
Cross-effects models are a type of LMM that have more than one class of random effects (i.e. T, > 1). Usually each observation belongs to one subject group (e.g. individuals, animals) and one item group (e.g. questions, objects). The correlation among latent variables can create severely challenging geometry that slows down the sampling of HMC. With our idea, it is possible marginalize one or more group of effects from the model, reducing the dimension of latent space for faster sampling and better geometry.
ETH instructor evaluations An example cross-effects model describes university lecture evaluations by students at ETH [4]. The dataset records N 73421 ratings, where each rating y/, comes from student 8,, for professor Pa teaching a course from department d,, with in indicating whether the professor is teaching outside their own department. There are total of M1 2972 students, M2 = 1128 professors and M3 = 14 departments. We use a version of the model from the document of Tensorflow probability [13]. The model is
Likelihood V/. ~ N(1,,, + U2pp + U3,d. +a 0²),
Prior: U1,1 ~,,, 1), 22 (O, 1), 43,k 1), a N(0,5)) B 1),, o (O 1), where 1 i M1. 1 S < M2 and 1 < k < M3. Given the dataset, we wish to learn about the latent variables u1, U2, U3. a, 3 and a. HMC the most direct way to sample those variables, but the dimension and complicated relations make it inefficient. Marginalization can be applied to one of the effects, U1, U2 or ug. We report the running time of sampling from the model with and without marginalization in Table We found that marginalizing any group of random effects improves the sampling speed of HMC. However, the improvements are not necessarily predicted by the dimension of marginalized variable: HMC is faster when marginalizing u3 than when marginalizing though u, has 200-times higher dimension than ug. In Figure 2 the ESS for each variable is reported. Without marginalization, sampling U2 and ug are both difficult compared to sampling u1, and HMC becomes more efficient when marginalizing either of these variables, so we conjecture that u2 and u3 are responsible for the difficulty for sampling in the original model. In this model, all random effects are independent and have the same variance, so En is scaled identity matrix and we can marginalize all random effects efficiently. This approach is observed be the most efficient in experiments, despite having quadratic complexity in D. Overall, marginalization never hurts ESS, and runs faster. We expect that any marginalization strategy works better than HMC in the original model, a finding which will be consistent across experiments. Additional results of this experiment, including trace plots and diagnosis, are included in Figure and Table in the Appendix.
even
6.2 Marginalization vs reparameterization
To tackle bad geometry in statistical models, another model transformation is non-centered parame- terization, or reparameterization (49. Reparameterization converts the distribution = ~ N(,, o²) into (0, 1) and hierarchical models. We note that when applicable, marginalization able to solve a broader class of problems. We compare marginalization and reparameteriation the grouse ticks model.
+ J.. Reparameterization is especially useful for funnel shapes in
7

---

23
23
23
25
23
25
2.0
2.
2.8
20
20
20
M divergence 54
21.1
9,1 div divergence 27
R2. civercence
divergence
divergence
divergence
0
20
20
40
29
Figure 3: Distribution of 10,000 samples for variable pains (01, #1.1) and (02, U2.61) on the grouset- ieks model with different methods. We use M1 to represent marginalizing u1, M2 to represent marginalizing ug. RI represent reparameterizing U1 R2 to represent reparameterizing u2. The number of divergences for each case are reported, with locations shown as red dots. We choose U2 to demonstrate the distribution of divergences when reparameterizing U2.
Table 3: Compilation time Te and running time T, in seconds for marginalized MCMC 35 with or without vectorization. Mean and std across 5 independ runs reported.
Model
Tc of 35
T, of 35
To of ours
T, of ours
Electric company
552 (4)
1249 (95)
7 (0)
252 (23)
Pulmonary fibrosis
727 (11)
2208 (80)
10 (1)
178 (3)
Grouse ticks The dataset 4 contains observations y of the the number ticks on the heads of red grouse chicks in the field. Each observation Y/ comes from brood bş in location 4 during year eg at altitude a₁. where year and altitude give fixed effects, and there are random effects u, and u2 corresponding brood and location. There are N 403 observations, M1 118 broods and M2 = 63 locations. We define the hierarchical model as follows:
Likelihood ~ N(u1bb + U2,1 + Beex + Boar, of) Prior: μ사 -O 1), σ1 HalfCauchy(5)) P2 ~ NO0 1), 02 HalfCauchy(5), 8. NO 1), 8n ~ (0, 1), 41,i ~ NNu,, ot), U2 ~ N(u2, o HalfCauchy(5),
where I,.., Mr.j j I,... M2, k = 1,.., . and each y is observed. The correlation between and creates the funnel shape that makes vanilla HMC inefficient. Nevertheless, it is possible to apply either marginalization or reparameterizatioo to each random effect. In Figure 3 we plot the distributions of samples for variable pairs (a (1,1) and (a2,, (2211) with different combinations of marginalization and reparameterization. There is difficult correlation between d2 and u2. After U2, HMC manages explore the funnel region (at
applying marginalization reparameterization low values of 01). However, we find that only samplers that marginalize U2 report zero divergent transitions after warm-up. Such behavior is consistent with different random seeds. See Table in the Appendix. Also, the distribution of divergent samples is related to specific parameters when reparameterizing u2, implying that reparameterization introduces pathologies that create challenges for HMC inference. In addition, we find that reparameterization does not improve the running time of HMC, while marginalizing U2 speeds up sampling by about 20%.
6.3 Benefits from vectorization
In theory, marginalization with LMMs can be done by constructing a graphical model for scalar random variables and performing automatic marginalization as in [35]. But it is more efficient to marginalize in a vectorized way. We demonstrate the benefits from vectorization in Table Both marginalization strategies are performed on two hierarchical linear regression models, the electric
company model 2 and the pulmonary fibrosis model [63] We find that vectorized marginalization much more efficient for sampling from the two models.

---

Table 4: Specifications of the datasets from cognitive sciences. Details of each model are provided in Appendix DJ GPU models run on an NVIDIA RTX 2080ti GPU. CPU models run on one Intel Xeon Gold 6148 processor.
dillonFi
english
ppds
mandarin
mandarinn
pupil
stroop
2855
26176
768
595
141
24
334
48
40
lfgg
48
16
30
16
16
15
15
Likelitood
LogNormal
Nonmal
Normal
Normal
LogNormal
LogNomaal
LogNormal GPI
Nomaa
LogNommal
GPU
100
Itens
No Marginalizatin
Marginalize u.
Marginalize
dillonE1
dutch
eeg
english
2g05
mandarin
mandarin2
pupil
stroop
Average itens sampling
1.00
Method Marginallzation Marginalize Manginalize Ma
0.50
0.00
dillnEE1
dutch
mandarin Average ESS/iter for latent variables
eng lish
9905
mandarin2
Figure 4: Experimental results for the 9 cognitive science datasets with and without marginalization. Each experiment is performed 5 times with different random seeds. Marginalization usually improves sampling speed measured by iterations per second (iter/s) and sample efficiency measured by ESS per iteration (ESS/iter).
6.4 Applications in cognitive sciences
Hierarchical Bayesian inference with LMMs has wide applications in cognitive science
47 We
highlight the effectiveness of marginalization with datasets from cognitive science (Table cover various settings, with one or two random effects, normal log-normal likelihoods, or GPU. Experiments that are slow on CPU are performed on GPU. Each dataset corresponds to an LMM where both the intercept and the coefficient include random effects. Details of all the models can be found in Appendix D Results are summarized in Figure 4 Marginalization usually improves the sampling speed of HMC and consistently improves efficiency measured by ESS per iteration.
4h. They
CPU
7 Discussion
There are several promising directions for future work.
Marginalization vs Blekkkellllizoionnn
Marginalization is related to Rao-Blackwellizattee This paper focuses marginalization. which improves the speed obtaining samples from the remaining variables by improving mixing times, reducing the cost per iteration, an interesting avenue for future work. More formally, if one is interested in some expectation (( u)] in an LMM, there is Monte Carlo estimator
both. Combining marginalization with Rao-Blackwellization
Σ f(O'
where (Q', u') ~ p(O, uly) and N is the sample size. Marginalization trick to improve the efficiency of the posterior sampling, so that we can achieve the same estimation variance with smaller N less runtime. the same time, we also have access to conditional distribution that useful for Rao- soo-Blaakwlllizatoon If the effects variable can be marginalized we have both approximate

---

posterior for p(Oly) and an analytical conditional distribution p( With AhooBlaakwellization we have that pie..)-p((e.lllllf( u)] = Be u)]]. In such case, another Monte Carlo estimator can be constructed:
N
F2
= N Σ
Barpulyy [S(0'
where @' ('>p(Oly) p(yl)). For some functions, such those that are polynomial in the inner expectation can be computed exactly using properties of Gaussians. In other cases, the inner expectation can be estimated cheaply via Monte Carlo using exact samples from p(u(O;,))
7.2 Marginalizing multiple effects in general models
In Section we proposed to marginalize multiple classes of random effects by assuming scaled identity covariance matrix. To marginalize multiple effects in general models, a possibility is to compute 2 and estimate (E) and the corresponding gradients with conjugate gradient (CG) solvers [1 20). However, this approach uses stochastic estimators for the determinant and gradients, which introduce bias into the HMC dynamics. These biases can be corrected through pseudo-maaginnaiizatoo [2], but it is unclear how significantly the extra stochasticity will affect the sampling. Another possible way marginalize multiple effects for LMMs is to introduced the balanced levels assumption [50]. We leave these ideas for future exploration.
7.3 Beyond normal likelihoods
In this work, we only consider normal or log-normal likelihoods, but our method can be easily gener- alized to other deterministic transformation of normal likelihood. This implies that marginalization benefit regression with most continuous predictors given proper link functions. Another potential future direction is to marginalize classification models with probit regressions []] Marginalization will turn probit models into multivariate probit models as A matrix, which may require a simulation-basee interesting to see how ideas from multivariate probit regression could be fit into HMC pipeline. In a broader context, marginalization is related to data augmentation techniques that "create" conjugacy for non-normal likelihoods or non-normal effects. Those techniques were developed for Gibbs sampling, e.g. [18,55] but may also be useful for HMC.
Ey is a dense covariance method [10] or variational Bayes [39]. It will be
7.4 Integration with probabilistic programming
We have developed tool to speed up the HMC inference for LMMs. our implementation, the marginalized likelihood p(y/O, is defined special type of parametric distribution available to the user, and the recovery distribution p( pu, O u. y) is a function called after sampling. In our experiments, marginalization never hurt sampling efficiency measured by ESS/s, and usually helped. Thus, it would be desirable to always marginalize one group of random effects when the model an LMM. Future work could aim automatically apply such transformations user-specified LMMs. There are two possible high-level approaches. The first is to perform marginalization starting with model described using a high-level abstraction such as an R formula. Then, when compiling the high-level model description into concrete model (c.g., a probabilistic program). we can marginalize one or more of the effects using our methods. The second to perform marginalization starting with a user- wertter probabilistic program representing an L.M.. In this case, some compilation program tracing technique will be needed convert the user's program to a model representation suitable for manipulation. For example, Lai et al. [35] used program tracing to construct 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.
Acknowledgement
The authors thank Yuling Yao and the anonymous reviewers for comments that greatly improved the manuscript. This material based upon work supported by the National Science Foundation under Grants #1749854, #2045900.
10

---

References
[1] Alan Agresti. Foundations of linear and generalized linear models. John Wiley & Sons, 2015.
[2] Christophe Andrieu and Gareth O Roberts. The pseudo-marginal approach for efficient Monte Carlo computations. 2009.
[3] Eric Atkinson, Charles Yuan, Guillaume Baudart, Louis Mandel, and Michael Carbin. Semi- symbolic inference for efficient streaming probabillisiic programming. Proceedings of the ACM on Programming Languages, 6(00PSLA2):1668-16)))
2022.
[4] Douglas Bates, Martin Mächler, Ben Bolker, and Steve Walker. Fitting linear mixed-effects models using lme4. Journal of Statistical Software, 67(1):1-48, 2015. doi: 10.18637/jss..006 v067. i01.
[5] Eli Bingham, Jonathan P Chen, Martin Jankowiak, Fritz Obermeyer, Neeraj Pradhan, Theofanis Karaletsos, Rohit Singh, Paul Szerlip. Paul Horsfall, and Noah Goodman. Pyro: Deep universal probabilistic programming. Journal of machine learning research, 20(28):1-6, 2019.
[6] Christopher M. Bishop. Pattern Recognition and Machine Learning. Springer, 2006.
[7] Helen Brown and Robin Prescott. Applied mixed models in medicine. John Wiley & Sons, 2014.
[8] Paul-Christian Bürkner. BRMS: An R package for Bayesian multilevel models using Stan. Journal of statistical software, 80:1-25 2017.
[9] Bob Carpenter, Andrew Gelman, Matthew Hoffman, Daniel Lee, Ben Goodrich, Michael Betancourt, Marcus A Brubaker, Jiqiang Guo, Peter Li, and Allen Riddell. Stan: A probabilistic programming language. Journal of statistical software, 76, 2017.
[10] Siddhartha Chib and Edward Greenberg. Analysis of multivariate probit models. Biometrika, 85(2):477-361, 1998.
[11] Michael Creutz. Global Monte Carlo algorithms for many-fermion systems. Physical Review
D, (8(4): 1988.
[12] Brian Dillon, Alan Mishler, Shayne Sloggett, and Colin Phillips. Contrasting intrusion profiles for agreement and anaphora: Experimental and modeling evidence. Journal of Memory and Language, 69(2):85-103, 2013.
[13] Joshua V Dillon, Ian Langmore, Dustin Tran, Eugene Brevdo, Srinivas Vasudevan, Dave Moore, Brian Patton, Alex Alemi, Matt Hoffman, and Rif A Saurous. Tensorflow distributions. arXiv preprint arXiv: 17/ 10604,, 2017.
[14] Justin Domke. Generic methods for optimization- modeling. In Artificial Intelligence and Statistics, pages 318-326. PMLR, 2012.
[15] Simon Duane, Anthony D Kennedy, Brian Pendleton, and Duncan Roweth. Hybrid Monte Carlo. Physics letlers B, 195(2):216-21)) 1987.
[16] Charles R Ebersole, Olivia E Atherton, Aimee L Belanger, Hayley M Skulborstad, Jill M Allen, Jonathan B Banks, Erica Baranski, Michael J Bernstein, Diane BV Bonfiglio, Leanne Boucher, et al. Many labs 3: Evaluating participant pool quality across the academic semester via replication. Journal of Experimental Social Psychology, 67:68-88 82, 2016.
[17] Stefan L Frank. Thijs Trompenaars, and Shravan Vasishth. Cross- lingu-ling differences in pro- cessing double-embed eebedddd relative clauses: Working-memory memory constraints or language statistics? Cognitive science, 40(3):554-578, 2016.
[18] Sylvia Frühwirth-Schieee Schnatter, Rudolf Frühwirth, Leonhard Held, and Håvard Rue. Improved aux- iliary mixture sampling for hierarchical models of non-Gaussian data. Statistics and Computing, 19:479 2009.
[19] Katelyn Gao and Art B Owen. Estimation and inference for very large linear mixed effects models. Statistica Sinica, 30(4):1741-1771, 2020.
[20] Jacob Gardner, Geoff Pleiss, Kilian Q Weinberger, David Bindel, and Andrew G Wilson. GPy- Torch: Blackbox matrix-matris Gaussian process inference with GPU acceleration. Advances in newral information processing systems, 31, 2018.
[21] Timon Gehr, Sasa Misailovic, and Martin Vechev. PSI: Exact symbolic inference for proba- bilistic programs. In Computer Aided Verification: 28th International Conference, CAV 2016, Toronto, ON, Canada, July 17-23, 2016, Proceedings, Part I 28, pages 62-83. Springer, 2016.
11

---

[22]
Timon Gehr, Samuel Steffen, and Martin Vechev. XPSI: exact inference for higher-order probabilistic programs. In Proceedings of the s1st aem sigplan conference on programming language design and implementation, pages 883-897, 2020.
[23]
Andrew Gelman and Jennifer Hill. Data analysis using regression and multtieevelhhierrrchical models. Cambridge university press, 2006.
[24] Disha Ghandwani, Swarnadip Ghosh, Trevor Hastie, and Art B Owen. Scalable solution to crossed random effects model with random slopes. arXiv preprint arXiv:2307.13378, 2023.
[25] Edward Gibson and James Thomas. Memory limitations and structural forgetting: The percep- tion of complex ungrammatical sentences as grammatical. Language and Cognitive Processes, 14(3):225-244) 2448, 1999.
[26] Maria Gorinova, Dave Moore, and Matthew Hoffman. Automatic reparameterisation of proba- bilistic programs, In International Conference on Machine Learning, pages 3648-3657. PMLR, 2020.
[27] Maria Gorinova, Andrew D Gordon, Charles Sutton, and Matthijs Vákár. Conditional inde- pendence by typing. ACM Transactions on Programming Languages and Systems (TOPLAS), 44(1):1-54, 2021.
[28] Philip Greengard, Jeremy Hoskins, Charles C Margossian, Jonah Gabry, Andrew Gelman, and Aki Vehtari. Fast methods for posterior inference of two-group normal-norma models. Bayesian Analysis, 183))888-9907, 2023.
[29] Daniel Grodner and Edward Gibson. Consequences of the serial nature of linguistic input for sentenial complexity. Cognitive science, 29(2):261-290 2005.
[30] Richard Grumitt, Biwei Dai, and Uros Seljak. Deterministic Langevin Monte Carlo with normalizing flows for Bayesian inference. Advances in Neural Information Processing Systems, 2022.
[31] Xavier A Harrison, Lynda Donaldson, Maria Eugenia Correa-Cano, Julian Evans, David N Fisher, Cecily ED Goodwin, Beth S Robinson, David J Hodgson, and Richard Inger. A brief introduction to mixed effects modelling and multi-model inference in ecology. PeerJ, 6:c4794, 2018.
[32] David A Harville. Matrix algebra from statistician's 's perspective, 1998.
[33] Matthew Hoffman, Pavel Sountsov, Joshua V Dillon, Ian Langmore, Dustin Tran, and Srinivas Vasudevan. Neutra-lizing bad geometry in Hamiltonian Monte Carlo using neural transport. arXiv preprini 1903:0070000774 2019.
[34] Matthew D Hoffman, Andrew Gelman, et al. The No-U-Turn sampler: adaptively setting path lengths in Hamiltonian Monte Carlo. J. Mach. Learn. Res., 15(1):1593-1623, 2014.
[35] Jinlin Lai, Javier Burroni, Hui Guan, and Daniel Sheldon. Automatically marginalized MCMC in probabilistic programming. In International Conference on Machine Learning, pages 18301- 18318. PMLR, 2023.
[36] Daniel Lewandowski, Dorota Kurowicka, and Harry Joe. Generating random correlation matrices based on vines and extended onion method. Journal of multivariate analysis, 100(9):
1989- 200992001,2222 2009.
[37] Jun S Liu. The collapsed Gibbs sampler in Bayesian computations with applications to gene regulation problem. Journal of the American Statistical Association, 89(427):958-966 -966, 1994.
[38] Colin M MacLeod. Half a century of research on the Stroop effect: an integrative review. Psychological bulletin, 109, 1991.
[39] Stephan Mandt, Florian Wenzel, Shinichi Nakajima, John Cunningham, Christoph Lippert, and Marius Kloft. Sparse probit linear mixed model. Machine Learning. 106: 1221-1622,, 2017.
[40] Charles Margossian, Aki Vchtari, Daniel Simpson, and Raj Agrawal. Hamiltonian Monte Carlo using an adjoint-diffie difoont-iifferetttttedd Laplace approximation: Bayesian inference for latent Gaussian models and beyond. Advances in Neural Information Processing Systems, 33:9086-9000 9999669997.--9 2020.
[41] Lotte Meteyard and Robert AI Davies. Best practice guidance for linear mixed-effects models in psychological science. Journal of Memory and Language, 12:104492, 2020.
[42] Kevin Murphy and Stuart Russell. Bao-Blackwellised particle filtering for dynamic Bayesian networks. In Sequential Monte Carlo methods in practice, pages 499-515, Springer, 2001.
12

---

[43] Lawrence Murray, Daniel Lundén, Jan Kudlicka, David Broman, and Thomas Schön. Delayed sampling and automatic Bao-Blackwellizzrrr of probabilistic programs. In International Conference Artificial Intelligence and Statistics, pages 1037-1046. PMLR, 2018.
[44] Praveen Narayanan, Jacques Carette, Wren Romano, Chung-chich Shan, and Robert Zinkov. Probabilistic inference by program transformation in Hakaru (system description). In Functional and Logic Programaming 13th International Symposium, FLOPS 2016, Kochi, Japan, March 4-6, 2016, Proceedings 13, pages 62-79. Springer, 2016.
[45] Radford M Neal. Slice sampling. The annals of statistics, 31(3):705-767,2000
[46] Radford M Neal et al. MCMC using Hamiltonian dynamics. Handbook of Markov chain Monte
Carlo, 2(11):2, 2011.
[47] Bruno Nicenboim, Daniel Schad, and Shravan Vasishth. An introduction to Bayesian data analysis for cognitive science. Under contract with Chapman and Hall/CRC statistics in the social and behavioral sciences series, 2021.
[48] Mante S Nieuwland, Stephen Politzer-Ahles Evelien Heyselaar, Katrien Segaert, Emily Darley, Nina Kazanina, Sarah Von Grebmer Zu Wolfsthurn, Federica Bartolozzi, Vita Kogan, Aine Ito, et al. Large-scale replication study reveals limit on probabilistic prediction language comprehension. Eliife, 7:e33468, 2018.
[49] Omiros Papaspiliopoulos. Gareth O Roberts, and Martin Sköld. A general framework for the parametrization of hierarchical models. Statistical Science, pages 59-73, 2007.
[50] Omiros Papaspiliopouloss Gareth O Roberts, and Giacomo Zanella. Scalable inference for crossed random effects models. Biometrika, 107(():25-44) 40, 2020.
[51] Omiros Papaspiliopoulos, Timothée Stumpf-Fétizon, and Giacomo Zanella. Scalable Bayesian computation for crossed and nested hierarchical models. Electronic Journal of Statistics, 17(2):
5575-3612, 2023.
[52] Matthew D Parno and Youssef M Marzouk. Transport map accelerated Markov chain Monte Carlo. SIAM/ASA Journal on Uncertainty Quantification, 6(2):645-682, 2018.
[53] Kaare Brandt Petersen, Michael Syskind Pedersen, et al. The matrix cookbook. Technical University of Denmark, 7(15):510, 2008
[54] Du Phan, Neeraj Pradhan, and Martin Jankowiak. Composable effects for flexible and acceler- ated probabilistic programming in NumPyro. arXiv preprint 1rXiv:19/2.115550 2019.
[55] Nicholas G Polson, James G Scott, and Jesse Windle. Bayesian inference for logistic models using 76lya-Gammn latent variables. Journal of the American statistical Association, 108(504): 1339 1349, 2013.
[56] Ian Porteous, David Newman, Alexander Ihler, Arthur Asuncion, Padhraic Smyth, and Max Welling. Fast collapsed Gibbs sampling for latent Dirichlet allocation. In Proceedings of the 14th ACM SIGKDD international conference Knowledge discovery and data mining, pages 569-577, 2008.
[57] Andres Potapczynski, Mare Finzi, Geoff Pleiss, and Andrew G Wilson. CoLA:
Exploiting compositional structure for automatic and efficient numerical linear algebra. Advances in Neural Information Processing Systems, 36, 2024.
[58] Jakob Robnik, G Bruno De Luca, Eva Silverstein, and Uroš Seljak. Microcanonical Hamiltonian Monte Carlo. The Journal of Machine Learning Research, (4(1):14666-147)) 2023.
[59] Hávard Rue, Sara Martino, and Nicolas Chopin. Approximate Bayesian inference for latent Gaussian models by using integrated nested Laplace approximations. Journal of the Royal Statistical Society Series B: Statistical Methodology, 71(2):319-99) 2009.
[60] Simo Särkkä and Lennart Svensson. Bayesian filtering and smoothing, volume 17. Cambridge university press, 2023.
[61] Matthias Seeger. Gaussian processes for machine learning. International journal of neural systems, 14(02):69-106, 2004.
[62] Matthias Seeger, Asmus Hetzel. Zhenwen Dai, Eric Meissner, and Neil D Lawrence. Auto- differentiating linear algebra. arXiv preprint arXiv:1710.08717, 2017.
[63] Ahmed Shahin, Carmela Wegworth, David, Flizabeth Estes, Julia Elliott, Justin Zita, Simon- Walsh, Slepetys, and Will Cukierski. OSIC pulmonary fibrosis progression, 2020.
13

---

[64] Jun Shi, Chengming Jiang, Aman Gupta, Mingzhou Zhou, Yunbo Ouyang, Qiang Charles Xiao, Qingquan Song. Wu, Haichao Wei, and Huiji Gao. Generalized deep mixed models. In Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pages 3869-3877, 2022.
[65] Justin D Silverman, Kimberly Roche, Zachary C Holmes, Lawrence A David, and Sayan Mukherjee. Bayesian multinomial logistic normal models through marginally latent matrix-T processes. Journal of Machine Learning Research, 23(7):1-42, 2022.
[66] Giora Simchoni and Saharon Rosset. Using random effects to account for high-cardinality categorical features and repeated measures in deep neural networks. Advances in Neural Information Processing Systems, 34:25111-25122 2021.
[67] Giora Simchoni and Saharon Rosset. Integrating random effects in deep neural networks. Journal of Machine Learning Research, 24(156):1-57, 2023.
[68] Wessel N van Wieringen. Lecture notes on ridge regression. arXiv preprint arXiv:1509.09169, 2015.
[69] Shravan Vasishth, Katja Suckow, Richard I. Lewis, and Sabine Kern. Short-term forgetting in sentence comprehension: Crosslinguistic evidence from verb-final structures. Language and Cognitive Processes, 25(4):533-56) 2010.
[70] Shravan Vasishth, Zhong Chen, Qiang Li, and Gueilan Guo. Processing Chinese relative clauses: Evidence for the subject-relative ad advantage. PloS one, 8(10)xe77006, 2013.
[71] Greg Ver Steeg and Aram Galstyan. Hamiltonian dynamics with non-Newtonian momentum for rapid sampling. Advances in Neural Information Processing Systems, 34:1 1012-11025, 2021.
[72] Basil Wahn, Daniel P Ferris, W David Hairston, and Peter König. Pupil sizes scale with attentional load and task experience in a multiple object tracking task. PloS one, 11(12): c0168087, 2016
[73] Jun-Kun Wang and Andre Wibisono. Accelerating Hamiltonian Monte Carlo via Chebyshev integration time. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023, 2023.
[74] Torsten Wörtwein, Nicholas B Allen, Lisa B Sheeber, Randy P Auerbach, Jeffrey F Cohn, and Louis-Philippe Morency. Neural mixed effects for nonlinear personalized predictions. In Proceedings of the 25th International Conference on Multimodal Interaction, pages 445-454, 2023.
[75] HI Wu and E Gibson. Processing Chinese relative clauses in context. In Poster presented the 21st CUNY Conference on Sentence Processing. University of North Carolina at Chapel Hill, 2008.
[76] Yuling Yao, Aki Vehtari, Daniel Simpson, and Andrew Gelman. Yes, but did it work?: Evaluat- ing variational inference. In International Conference on Machine Learning, pages 5581-5590. PMLR, 2018.
[77] Zhaoxia Yu, Michele Guindani, Steven F Grieco, Lujia Chen, Todd C Holmes, and Xiangmin Xu. Beyond test and ANOVA: applications mixed-effects models for more rigorous statistical analysis in neuroscience research. Neuron, 110(1):21-355 2022.
14

---

A
Notation table
We summarize the important symbols used in the paper.
Symbols
Description
N
Number of observations, and dimension y Dimension for all effects one class mixed effects Number of classes of mixed effects Dimension for all mixed effects Dimension for effects of group in class
M, M₁
L
D
d
k. ki
Number of groups in a class
Intercept for linear regression Slope for linear regression Standard deviation Random effects
B
u. u
Concatenated random effects Observations
V.Y
c,t
Covariates, or treatments
T
A prior variable sampled from half-normal distributions prior variable sampled from I.K distributions Grouping variables
g
Global variables, including priors and fixed effects Mean of random effects Design matrix for random effects Concatenated design matrices
A
B
Intercept term in the canonical form for LMMs Covariance matrix for class of random effects Covariance matrix for the observations Covariance matrix for all random effects Scale for En with the scaled identity assumption Scale for Ey with the scaled identity assumption A shared matrix in the reversed model
Σu
Ey
Σ
Tv
Ty
M
Difference between observation and mean of the marginalized likelihood A dense N x N matrix that is difficult to directly compute The core matrix after applying the two linear algebra lemmas An intermediate matrix in the implementation An intermediate vector in the implementation
E
F
G
A row of A column of A
C
A block of d columns of A
The eigenvector matrix for eigendecompsition of B
The eigenvalue matrix for eigendecomposition of B
15

---

B Proofs and details
B.1 Proof of Theorem
We first review the tree structure of the matrix A. A is an N M matrix where every block of d columns corresponds to the effects for one group (e.g., an individual subject, age, school, or gender). For example, if N 3, k 2 and d 2, one possible graphical model is as below.
U1
U2
1/1
y2
y/
Figure 5: A tree-structured model conditioned on O.
Each Uj E IR2. If the coefficients are all 1s, then
A
1
1
0
0
0
To generalize, if for Yis the grouping variable is Ji, then the ith row of A, only A.jas can be nonzero for j = (0: - 1)d + 1 and h = ged. We consider three representations of the matrix A. By rows,
F1
A
r2
TN
by columns,
A (C1
Ckd),
and by blocks of columns,
A (A.,1.d A.
A (.(
-C1 C2 C₁)
where each C, (i= 1,2,
k) is N d. Now we restate and prove Theorem
Theorem 1. If Ey is diagonal, E. is block- dickodiaaonnll with blocks of size d d, then F = Σ A ATE," A is also block- diagonaa with d x d blocks and computing A AT A takes O(Nd)) time.
Proof. The theorem has two parts: (a) the property of F, and (b) the computation F. We address them with the three representations of A.
(a) F is block-diagonal. Because Σ, is block diagoni,, Σ, is also block diagonaa with the same sizes. Also, Ey is diagonal, the block block-diagonality of A Ey A is the same A A. We consider column representation of A, then
ATA
(C)
C2
C+)
CT
CTC
CT
C.
C. Cf
C²
C
C[C, clC,/
16

---

Forl C, For sisks
c
1)+41
C
1)d42
(Cj-11dt41
C(-1)d+2 Cyd)
Cid
The following lemma shows that C. Cj =
Lemma 1. For any j and S 8, 1. S d, it holds that of (1-1)d+ CU-1)d++ 0.
T Proof. The lemma can be proved by contradiction. Suppose (3- 1)-1)dd+ (i-)d=s(--)ddd 1)d+t * 0. Then there exists index such that C(=-1)d+a(r) A 0 and ((--)dd++ [n] f O. This means that in the nih row of A, both (n,(f-1)d+s and An.(j-1ddd are non-zero. This contradicts with the tree-structure where only group of d elements can be non-zero
With the lemma, we have that A A is block-diagonal, thus E + A My A. is also block-diagonal and each block
d.
(b) The computation of A Σy A is O(Nd)). Since Ey is diagonal, A' = = AA has the same pattern of zeros and nonzeros A. We consider the row representations such that
r
A'
T'
Then
ATA' (1f
T
TN)
I2
Σ
rTr,
note that each of r; and r has d non-zero elements. So computing A A' is O(Nd)).
B.2 Pseudocode for recovery after marginalizing one group of random effects
Algorithm 2 Sampling from puu, O 1,y,uu. u.;)
==y - Ejfs Aju, Aip,
Sparse matrix multiplication in O(NLd) time
G A A
Block diagonal computation in O(Nd)) time
F Σ,1+G
Block diagonal computation in (MMd) time
= μ4; (I GF AA'-
>> Sparse matrix multiplication in O(M + N)d) time
L = En Cholesky(Ep
G GF-IG) Cholesky of block diagonal matrix in (Md)) time
return u, Normal LLT).
B.3 Details of scaled identity covariance matrices
With the assumptions of scaled identity covariance matrices, effects can be marginalized with preprocessing of the eigendecomposition of B B.
(a) Preprocessing before HMC. We compute
BTB - QAQ
In LMMs, the computation of B 2 and the eigendecomposition of is O(D)). So the overall complexity for preprocessing is O(D NL"d²,. Compared with the HMC sampling
Each AT Aj is (NNd), corollary of Theorem

---

loop that takes thousands of steps and visits the model hundreds of times each step, the cost of preprocessing not expensive. In our attempt to marginalize all random effects for the instructor evaluation model in Section 6.1 this step takes less than 10 seconds.
(b) Marginalized likelihood during HMC. During HMC sampling, the log density log p(y((/ ()) would be calculated, which is
1 ply ( det (E)
1 2² z
where z = y By - and E BB B + By. The computation of z takes (N time. With the two lemmas, have
det E = det D B T Σ
B)
det (Ev) det
E 2= 2 z Σ B E B B B)
'B" My z.
shared matrix in the formulas is
Ey + B" Σy
LB. Then
F= F= + B" E Q
-
Q²:
Ty
With the trick, evaluating (et reduced to O(D) time det (Q
1. Also z Z=E== E-12 becomes
- I+ Ty AA)
z²E DE''z=- "z=z==z"
z² Ey BQ
Q Q'B"E, Ey''
Note that B z E RRD. can be computed in ONN time, but its multiplication with Q" takes O(D)) time. Given that I + 1 A is diagonal, the complexity of evaluating log boy((y/) once is then O + NLd).
(c) Ancestral sampling after HMC. the recovery step, we perform ancestral sampling from p(v y). To efficiently generate samples, we give the following theorem.
Theorem 2. Ev = IvI, Ey = and B = QAQ", then
v/O.y N(H Evjo.y);
where
Hvjoy ++
Tv
B
QA
1
A
-1
Q²
z.
Ty
Tv
Evjo,y = Q
TvI
Tv
A+
Tv
A
-1
A
QT
Ty
Tv
Ty
In Theorem 2 from z, we can apply matrix multiplications from right to left to get Hvje.y The whole computation takes O(D2 N.N Ld). To generate normal samples a Cholseky factorization for
2 diagonal, so can be obtained in -1
Evje.y is required. But T.I TV
O(D)) time as well. Now we prove Theorem
Proof: Pvle.yy and Σvje.y can both be derived algebraically.
Mvjo, μ Mz
F + ETB
BT +
Ey)
= + E.B
Σ
E B
+ BBT
BTT B'' z
+++ Tv (B"
B B B
B" B" (B)B"'
Ty
(B
QAQTQ
1 It 1 A
Q B
4
Ty
Tv
Ty
1 A)
+
(BT
QA
-
I
jz.
Ty
Ty
Ty
18

---

Evjey = - E,
(I E,B² (BE, 2 EB B)E,
.TI TV BBB BBE, + My 'B
= Tv - Tv 1,,
Tvl Tv B My
B + Tv B
B B
B B
B E
+B EB BB"B B
2
1 A
Tyl
QAQT=
+
QAQQO
Q
- 1+
Q Q'QAQT
Tv
Ty
Tv Ty
Tv
Tvl
QAQT+
+
QA
----A
- A
AQT
Ty
Tv
Ty
222 2
2
=Q
TvI
A
+
2
Q".
Ty
Ty
Tv
Ty
19

---

C Additional experimental results
manginalization
Marginalize
Manginalize
Marginafzo
Marginalize
10020
5250
9300
10000
5002
9250
30000
9200
9300 Marginalize uy
20002 13000
5220
No marginatication
Marginalizes
Marginalize MJ
Marginalize
1.0
-0.3
9250
2750
20010
2000
9250
10002
0002
$250 5750 10020
0000
$253
9738 10200
2200 9230
1500
9750
10002
4.5
Marginalize
Marginalize
Marginalioe
4.0
2.5
2003
10500
12000
10d0o
stop
Figure 6: Trace plots for M2,1, 13,1 and a of an interval of 1,000 sampling steps after warmup on the ETH instructor evaluation model, using the same data Figure 2 in the paper.
Table 5: Number of parameters (out of 4117) whose R exceed a threshold for 1,000 samples from HMC, with without marginalization. Mean and standard deviation over independent reported.
Threshold
No marginalization
Marginalize uj
Marginalize u2
Marginalize ug
Marginalize u
> 1.01
186.80 (37.26)
295.60 (134.57)
11.80 (6.05)
59.80 (37.35)
5.20 (1.72)
> 1.02
99.40 (18.91)
153.20 (99.90)
6.20 (7.19)
9.80 (10.48)
0.00 (0.00)
1.05
13.40 (6.83)
54.00 (51.99)
0.00 (0.00)
0.00 (0.00)
0.00 (0.00)
1.10 0.00 (0.00)
23.20 (24.51)
0.00 (0.00)
0.00 (0.00)
0.00 (0.00)
Table Divergence (mean and standard deviation) out of 10,000 samples with different strategies on the grouseticks model across 5 random seeds under different target probabilities. We use MI to represent marginalizing u₁, M2 to represent marginalizing u2 RI to represent reparamcterizing u, R2 to represent reparameterizing u2.
Transformation
Number of divergence
No marginalization
42.60 (25.76)
MI
14.60 (13.85)
M2
0.00 (0.00)
M1, R2
22.60 (18.13)
M2, R1
0.00 (0.00)
R1, R2
431.60 (507.37)
20

---

D Models and example probabilistic programs
We provide the details of the nine cognitive science datasets and their corresponding models and probabilistic programs. We follow [47] and use maximal models with correlated varying intercept and slopes for each the datasets. The model for the pupil dataset described in Section 2
D.1 Agreement attraction in comprehension
The dataset (dillonE1) studies the effect of the agreement attraction phenomenon when reading a noun with the auxiliary verb [1]]. The predictor is
log(y:) a + U1 I + U2.92..1 44 + 21. 2 + 22.92..22) + E,E (²)
Each experiment result yi is from subject 91, on sentence J2,s, with I, being the interference level (4, E {0,1}) Bayesian hierarchical modeling assigns prior to the variables.
Ti ~ / (0, diag(5", ,,)), L1 (2JCholesky(2.1), T2 ~ N (0, diag(5?, 53)), L2 LKJCholesky(2..) ~ (O, 8 ~ (0 5²), ~ (0, Un ~ (O, LTT,, U2 ~ (O, T2L,LTT,))
The probabilistic program in NumPyro is then
def model mn, ttste, nLobe, gl. g treatnent, oba):
alpbs numpyro. aalpha'
Normal 10))
muxpyro cample
signa numpyro exppla( ("signa', dist. Hal filornlllll (5))
cigna_n numprro, eaxple( sigra_e", dist (2)Choleckyy()) eepple"ttm.ol "ta,,', stgna_y nungyro. eawppe uigiga__r'
eemple
(bltiiommlll)) (5), sempl_ebape-(2,
(2)
tauyy nurpyro samplel ttau_', dist Hhlffiorma (6), sampl_ )) e_n natmal (jnp diag(tan,u), sigta,u)
sat.ul (Jup. diag(tau_ sigra_v) sumpyro cample
hul.tivaaasrrillmmllleee (2,)) scale, tril-s_u), chaple.,caeenn__ee
zumpyro. cample(ty', diet Nultsrerriareloommlloee (jop (jap. 2eron((2))) sealee compl_.ttaper(__tten
numgyzo. sample("y', dist LogHorrel (SIpha (..i))...,.) v[e2) [...,]]
treatment
(beta u(gi]
r(g2).... sigra), obsrobe)
We use and v in the codes represent the two random effects. The probabilistic program with marginalization is similar. Suppose we marginalize our probabilistic program becomes
def model (n, rub, tter, nLobe, gl, g2, treatnent, obe):
alpha
'alphn' sarpyyro..splll ('ataa,
Normal (9, 103)
signa numpyro cample eapale'stgnna dist. Inlifioomal (5))
aigm__ numpyro maplla siggaa"" diat. (X)Chol
numpyro. sample "ttn,j'',
(biffiorml (6), sample_cappe"2,
signs_x nungyro.
san_y muppyro expple( (iigraa"
(XJCholeeky)))
diat. fllormal (5), sampla_ e_u satmul (jop diag(tan, sigra_u)
matmal (jop. diag(tau,y),, sigts_7) jop_zerea((n__noh
y zumpyro. cample('y', dist. hultirerdsareeiomalllaao (sap 2eros((2))) seale seale, tril-s_v), cample_chae-raatteen
nunpyro. mumpyro.cample( (nyp
wwwproooooooooooooooeeeeeeooop (alph sigil...... a vigh [...,0]
(beta ulgi] v(g2] (1),), eigas, gt,
obsnobe)
To marginalize v, the probabilistic program is
model au,, Huis, aLoba, gl, g2, Leoatewet, aba):
alpha mimpyro casple( Calpha" ' dist, Norral (o, 10))
beta sumpyro.xuplee ('bta', dist. Morzal (0, 5))
signs muspyro. muppro..asple comple(Cstggeer "sggraus" dies. (e)))))
signa_n nungyro earple
tai_s muppyro assple( (ommppl("ttaaaa taaja", dist. Hol flformal (5), sampla_ahape-(33, })
signs_v oumplee sumple( "iggaay", dist, 1XICholesky())) (2))
mimpyro. amsple( (ttnp"",
sunpla_ahage=(0,,
*_a jup satnul (jup. diag(tan__)) signa_n)
e_7 nateul (jnp diag(tan,y), sigtayy) (en.-erraa(nn.. itwe,
satppyo..arpll""" numpyro.
saattivarittttionaalll (jtpppppt..((..))) scale_trilossu)) suppl wigll.. (.. o0 vIgh
exmple('y", eoperro..appeee.. traatrant. (beta [lg!].....1] vig2] [.... a.y. aignn, g2, trastrent, nuitem, nLoba, v), ohe-nba
The probabilistic programs for the other models will be similar and omit them for simplicity.
21

---

D.2 English and Dutch Grammaticality illusion
The datasets (english 69 dutch [7] study the forgeetgg hypothesis 23 for different languages. They use the same predictor and priors. The predictor is
y/ a + 1 11.... + 12. + ( + 11.01..22 + 22.02.. + N((,, where & is the treatment variable and i, E (-1,1). And the prior is
T1 ~ (0, diag(12, 12)) L1 ~ LKICholesky(2,1), T2 ~ ~ (O,diaa(1",,1) Lg ~ LKJCholesky(2,1) ~ (0, 3 ~ (0,, ~ N U1,y ~ (O, TT,,, U2 ~ NO T2L2L,T2)
D.3 Electrophysiological responses with N400 effect
In the study of language, the electreencephalography (EGG) responses with N400 effect is studied Experimental results of subjects from the Edinburgh lab are collected 488. The predictor is
y/ "1 1 + 12 + [ u1,91,2 (2.92..2) + 6,6 (00,
where & is the treatment variable and 1 [0, 1]. And the prior is
T1 No (0, diag(202, 20²)); L I.KICholesk((2, 1), T2 ~ (0, diag(20² 20")) L2 ~ LKICholesky(22) 1), 103), (O, 10?), a (0,,50))") u1, ~ LL U2 ~ N(0, T2L2L, T2).
D.4 Subjective and objective relatives
Grodner and Gibson [29] (gg05) studies the processing time difference between object relative clause and subject relative clause sentences. The predictor is
log(u:) = a + 1 + U2.02...1 + 13. + (2 + u1.01:.2 + 22.92.122 + U3.93...1) + 6, (V0,
and the treatment variable 4, E [-, The third effect u3 is related to different repeats of the experiment and has only two groups. We consider the first two effects for marginalization match the other experiments. The prior for the variables is
T ~ ~ (0, diag(a" ,3))) T2 ~ N (0, diag(52, (5)), T3 ~ N diag(5?, L LKICholesky(2,, 12, 1), L2 AKICholesky(2,1) L3 LKICholesky(21 1), N00. 10²), B (0,05 (00,6")) U1 (O, T,L,,,,, U2 ~ (O,, T2L T2), Uz, ~ NoO0, T3L Ts).
D.5 Relative clause processing in Mandarin Chinese
The datasets (mandarin 75 mandarin2 (70) are collected from experiments to study the effect of
relative clause type reading time of Mandarin Chinese. In model, the predictor is
log(ys) = a + M1.0...1 + 12.02..11 + 1,(3 + .1.01..2 + u2.22..2) + e, ~ (O(,
where 4 is the treatment variable and I, E (-0.5, 0.5). And the prior is
T, ~ N (0, d))) L, ~ LKJCholesky(2, 1), To ~ N (0, diag(52,55) L2 ~ LKJCholesky(2, 1), N (0, B (3,, a~ ~ uny ~ ~O, (O.,TTL,TTT.)) u2 NO, T2L,LTT) T2).
D.6 The Stroop effect
The Stroop effect describes the change of response time between congruent and incongruent stimuli 38). The dataset is from Ebersole et al. [16]. Different from the other models, the noise scale for each observation is also grouped. In our model, the predictor is
log(us) = + U + (,3 + Mg..2) + e~ (O of), o, = exp(a. + 89,,1 + ti(as + sg.2)),
and the treatment variable is t; E {-1, 13. Priors for the model are
T, ~ N (0, diag(1, 1)), Lg ~ LKJCholesky((..1) T. ~ ~N* (o, diag(1, 1)), Lg ~ LKJCholesky(2,1) 1), a ~~~N N66 1.5²), B ~ (0000.013) oa ~N(00,) 1), 03 ~ 1), u ~ N(0, L. T.L.L 8; ~ N(0. T. T.L.L.TT).
22

---

NeurIPS Paper Checklist
The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility. transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and precede the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No], or [NA].
[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justifiiaton right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation.
While "[Yes] " is generally preferable to "[No]", ", it is perfectly acceptable to answer "[No] provided proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] or "[NA] is not grounds for rejection. While the questions are phrased in binary way, we acknowledge that the true answer often more nuanced, so please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS paper checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Our major claim is the faster algorithm for marginalization and applications in various models, which are supported by the theories and experiments.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]
23

---

Justification:
The discussion section covers three future directions to improve over.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approxinations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification:
The details and proofs are included in the appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: The algorithm we propose can be directly implemented and we use publicly available HMC.
Guidelines:
24

---

The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open- source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome deseribe the particular way they provide for reproducibilityy In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
We provide the data and code in the supplemental material.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips... cc/
wwt.:/uwwissssooossoooooooo for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https CodeeSumiiii www.iisssooooioii
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
25

---

Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
The hyperparameters for all models are included in the appendix.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer:
[Yes]
Justification:
We reproduce all experiments five times and report error bar or std if possible.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (c.g. negative error rates).
show in tables
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We report both the type of machines and the time in our experiments. The memory usage is below 32GB in all experiments.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
26

---

The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethies
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics http::///oauupp...ppoooooo//tctt
Answer:
[Yes]
Justification: We have reviewed and verified that it confirms to the CoE of NeurIPS.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g.. if there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: The paper is methodological in nature.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Answer: [NA]
Justification: The data and model are all publicly available.
Guidelines:
The answer NA means that the paper poses such risks.
27

---

Released models that have high risk for misuse or dual-use use should be released with necessary safeguards to allow for controlled use of the model. for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We have cited all the frameworks and datasets use.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
We have documented our submitted codes.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: We public datasets and do not conduct research with human subjects.
28

---

Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
We use public datasets and do not conduct research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
29


Reviews:
[{"review_id": "YiajExIS8T", "reviewer": "Reviewer_u9Xy", "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.", "questions": "What sort of prior structures allow for this to be applied?\n\nAre there methods that can be developed that perform approximate marginalization in some way?\n\nHow hard is it to automate marginalization in practice?", "limitations": "Yes.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "```\nWhat sort of prior structures allow for this to be applied?\n```\nIn the context of LMMs, the distributions of the effects are assumed to be normal. We do not restrict the priors of other variables.\n\n```\nAre there methods that can be developed that perform approximate marginalization in some way?\n```\nIn 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.\n\n```\nHow hard is it to automate marginalization in practice?\n```\nThere 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. \n\n[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.\n\n[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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "435VI3aMd5", "reviewer": "Reviewer_UJHx", "summary": "The paper studies marginalization in linear mixed-effects models with a Gaussian likelihood.\nMarginalizing out variables is a common trick to improve the convergence of MCMC.\nAs 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.\nThe paper introduces two tricks using the matrix inversion lemma to speed them up.\nSome experiments are performed to demonstrate the effectiveness and to compare against the reparameterization alternatives.", "questions": "Is the method specific to HMC?", "limitations": "Limitations are discussed in section 7 however I do think those limitations make the paper a bit weak (not enough contributions) for acceptance.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 1, "strengths": "## originality\nThe paper presents a simple improvement to a known problem.\n\n## quality\nThe technique presented in the paper is relatively simple and straightforward. Such techniques are widely used in Gaussian processes community.\n\n## clarity\nThe 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.\n\n## significance\nThe 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\nOnly models with Gaussian or log-Gaussian likelihood is studied.\n\n## Experiments are not enough\n1. In the experiments there is only naive baselines. Other alternative automatic marginalization methods are discussed in the related work section but not compared against.\n2. Only 1 dataset (section 6.1) and 1 toy target (section 6.2) are studied.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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. \n\n```\nSuch techniques are widely used in Gaussian processes community.\n```\nAlthough 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. \n```\nOnly models with Gaussian or log-Gaussian likelihood is studied.\n```\nWe 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. \n\n```\nIn the experiments there is only naive baselines. Other alternative automatic marginalization methods are discussed in the related work section but not compared against.\n```\nAutomatic 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. \n\n```\nOnly 1 dataset (section 6.1) and 1 toy target (section 6.2) are studied.\n```\nReviewer UJHx counts 2 datasets but in our experiments, there are in total 13 datasets from Section 6.1 to Section 6.4. \n\n```\nIs the method specific to HMC?\n```\nOur method works for any MCMC algorithm that accepts a black-box model. We highlight using HMC from practical perspectives. See the general responses. \n\n[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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "winC3DZYlj", "reviewer": "Reviewer_jeEo", "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.", "questions": "1. Can you provide more detailed convergence diagnostics for your experiments, such as trace plots, for example?\n2. Are divergences grouped in specific regions of the posterior distribution?\n3. 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?\n4. Could you elaborate on the choice of priors in your experiments? Why were these particular priors selected?\n\nEdit: the rebuttal has answered all my questions.", "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?", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "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. \n\nTo 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. \n\nI 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.\n\n### Experimental Details: \n\nOverall, I find that the experiments lack sufficient details and metrics:\n* 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?\n* 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.\n* 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? \n* The choice of priors for various models seems to me quite broad and lacks justification. \n\n### Significance of Results\n\nThe 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.\n\nEdit: the rebuttal clarified on all of these points.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n```\nCan you provide more detailed convergence diagnostics for your experiments, such as trace plots, for example?\n```\nIn 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. \n\n\n```\nAre divergences grouped in specific regions of the posterior distribution?\n```\nWe notice the divergences of R1,R2 are related to specific combinations of values for $\\sigma_2$ and some of the parameters in $\\textbf{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.\n\n```\nDoes 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?\n```\nWe 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 $\\textbf{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). \n\n```\nCould you elaborate on the choice of priors in your experiments? Why were these particular priors selected?\n```\nFor 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. \n\n```\nThe paper mentions Rao-Blackwellization, but doesn’t go into details.\n```\n\nThanks 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.\n\n  \n\nMore formally, if one is interested in some expectation $\\mathbb{E}[f(X,Y)]$ one can view Rao-Blackwellization as transforming the estimator as\n\n  \n$$\n\\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],\n$$\n  \n\nwhere $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).\n\n\nIn 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\n\n$$\n\\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),\n$$\n  \n\nwhere $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.\n\n  \n\nFinally since the distribution of $Y$ given $X$ is tractable, one could imagine applying both our technique and Rao-Blackwellization, with the transformation\n\n  \n$$\n\\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],\n$$\n  \n\nwhere 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.\n\n\n[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.\n\n[2] Andrew Gelman et al. (2024) Prior Choice Recommendations.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Obznd4BA0O", "reviewer": "Reviewer_aGrs", "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.", "questions": "1. Can this be extended to multivariate observations?\n2. 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?", "limitations": "I don't foresee any negative societal impact.", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "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.\n2. The presentation is good and the writing is clear.\n3. There is sufficient empirical demonstration, in my view.", "weaknesses": "**Title**\n\nI 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). \n\n\n**Presentation**\n\nOverall, I think the presentation is good. Just a few minor comments:\n\n1. 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\".\n2. 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).\n3. 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.\n4. 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.\n5. I couldn't quite understand Section 6.3. What does \"marginalization in a vectorized way\" mean?\n6. Check the bibliography; it has many mistakes, e.g., missing capital letters in proper nouns and journal names.\n\n\n**Other**\n* L241: document -> documentation", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "```\nI don't think the title (and the emphasis on HMC in the abstract and introduction) is entirely adequate.\n```\nThanks 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..\n\n```\n... 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\n```\nWe appreciate your meticulous efforts in reviewing our submission. We will make a complete revision of this manuscript to incorporate your suggestions.\n\n```\nWhat does \"marginalization in a vectorized way\" mean?\n```\nWe 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.\n\n```\nCan this be extended to multivariate observations?\n```\nFor multivariate observations, sparsity still exists in the model structure and our method should work with no issue. \n\n```\n... \"no marginalization\" ... slower than ... marginalization. ... Do the authors have an explanation for this?\n```\nThere 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. \n\n[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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
# General response

## Title of the work

It’s a good point that the method is not limited to HMC. We will adjust the title to better reflect our contributions. We propose the alternative "Hamiltonian Monte Carlo Inference of Marginalized Linear Mixed-Effects Models". In our revision, we will also rework the introduction to distinguish between the contribution of marginalization and the application to HMC.

While marginalization could be applied with any MCMC algorithms that target a log density, we chose HMC as the kernel and in the title with the following considerations:
- HMC is the working algorithm of the influential package BRMS, and has shown great success in solving real world linear mixed-effects models (LMMs).
- With HMC, the benefits of marginalization can be justified as reducing the state dimension $H$, thereby reducing the computation time of reaching an independent sample with the Leapfrog integrator, whose complexity is approximately $\mathcal{O}(H^{5/4})$. Such a bound may not exist in the literature for all MCMC algorithms, so we were conservative in defining the scope of our method.

## Normal / log-normal likelihoods

Normal and log-normal likelihoods in the paper cover a lot of studies with continuous response variables, and it is direct to generalize our method to other continuous domains with proper transform functions. The package lme4 [1] has over 79,000 citations and a substantial number of those studies have continuous observations. We have included 9 studies in cognitive science in our experiments as demonstration. Also in the Bayesian workflow, normal models can act as template models [2] when developing more complicated models. We believe the scope of this work has broad interest and will lead to exciting future works.

Exact marginalization in the context of discrete likelihoods (eg. Bernoulli, Binomial, Poisson) often involves data augmentation tricks such as [3] and other related works. The generalization of our methods to those likelihoods with the tricks may also require a novel MCMC algorithm, which is an interesting avenue for future work but beyond the scope of this paper.

[1] Bates, D., Mächler, M., Bolker, B., & Walker, S. (2014). Fitting linear mixed-effects models using lme4. _arXiv preprint arXiv:1406.5823_.

[2] Greengard, P., Hoskins, J., Margossian, C. C., Gabry, J., Gelman, A., & Vehtari, A. (2023). Fast methods for posterior inference of two-group normal-normal models. _Bayesian Analysis_, _18_(3), 889-907.

[3] Frühwirth-Schnatter, S., Frühwirth, R., Held, L., & Rue, H. (2009). Improved auxiliary mixture sampling for hierarchical models of non-Gaussian data. _Statistics and Computing_, _19_, 479-492.

Meta-review: This manuscript considers Bayesian inference in a popular class of models, linear mixed-effects models. The authors discuss strategies to improve the performance of Hamiltonian Monte Carlo sampling in this class of models, including how to use ideas from numerical linear algebra to accelerate the running time dramatically.

During the initial review period, the reviewers generally praised the work, identifying only some minor concerns, which were adequately addressed during the author-reviewer discussion period. Following the discussion period, there was consensus in favor of accepting the work.
Decision: Accept (poster)