Research map / Probabilistic and scientific modeling / Variational and generative sampling
Hamiltonian Monte Carlo Inference of Marginalized Linear Mixed-Effects Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Shows how to marginalize random effects in linear mixed-effects models inside HMC via fast linear algebra (matrix inversion/determinant lemmas with tree-structured design), reducing per-step cost to linear in N for one random-effect class and speeding up NUTS in NumPyro.
Problem
Efficient Bayesian inference with HMC in linear mixed-effects models is hampered by funnel geometry and high dimension of random effects, and manual marginalization is difficult and naively cubic in cost.
Approach
Vectorized analytic marginalization of random effects in LMMs, where the tree structure of the design matrix makes the core matrix block-diagonal so the log-density and recovery steps run in time linear in the number of observations for one random-effect class; with scaled-identity covariances, all classes can be marginalized using an eigendecomposition of B^T B.
Key results
- On the ETH instructor evaluation model, running time dropped from 13417 s without marginalization to 631 s when marginalizing all random effects.
- On the grouse ticks model, only samplers that marginalize u2 reported zero divergent transitions after warm-up, and marginalizing u2 sped sampling by about 20%.
- Vectorized marginalization reduced compilation time from 552 s to 7 s on the electric company model and running time from 1249 s to 252 s relative to scalar marginalized MCMC.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Marginalize u1, u2, or u3 instead of sampling all random effects | NUTS · ETH instructor evaluations | running time: 13417 (98) -> 5004 (1468); 2607 (3); 3071 (4) | Helped |
| Marginalize all random effects instead of sampling them | NUTS · ETH instructor evaluations | running time (s): 13417 (98) -> 631 (12) | Helped |
| Use vectorized rather than scalar marginalization | Marginalized MCMC · Electric company | running time (s): 1249 (95) -> 252 (23) | Helped |
| Use vectorized rather than scalar marginalization | Marginalized MCMC · Pulmonary fibrosis | running time (s): 2208 (80) -> 178 (3) | Helped |
| Use vectorized rather than scalar marginalization | Marginalized MCMC · Electric company | compilation time (s): 552 (4) -> 7 (0) | Helped |
| Use vectorized rather than scalar marginalization | Marginalized MCMC · Pulmonary fibrosis | compilation time (s): 727 (11) -> 10 (1) | Helped |
| Marginalize u2 rather than leave random effects untransformed (ablation) | HMC · Grouse ticks | divergent transitions: 42.60 (25.76) -> 0.00 (0.00) | Helped |
| Marginalize u1 rather than leave random effects untransformed (ablation) | HMC · Grouse ticks | divergent transitions: 42.60 (25.76) -> 14.60 (13.85) | Helped |
| Reparameterize both u1 and u2 instead of leaving them untransformed (ablation) | HMC · Grouse ticks | divergent transitions: 42.60 (25.76) -> 431.60 (507.37) | Hurt |
| Reparameterize rather than marginalize u2 | HMC · Grouse ticks | running time | Helped |
| Marginalize random effects | HMC · Grouse ticks | running time: about 20% | Helped |
| Marginalize random effects | HMC · Cognitive science LMMs | ESS per iteration | Helped |
Surprising findings
- Running time was not predicted by the dimension of the marginalized variable: marginalizing u3 (14 dims) was faster than marginalizing u1 (2972 dims) on the ETH model.
- Reparameterization did not improve HMC running time on the grouse ticks model, whereas marginalization did.
Limitations
- Applies only to linear Gaussian structures (normal or log-normal likelihoods and transformed normals).
- Marginalizing all classes of effects requires scaled-identity covariance assumptions; general multi-class case left to future work.
- Quadratic complexity in D for the all-effects approach.
Experimental setup
Datasets: ETH instructor evaluations grouse ticks electric company pulmonary fibrosis pupil size dillonE1 dutch eeg english gg05 mandarin mandarin2
Benchmarks: ETH instructor evaluations grouse ticks electric company pulmonary fibrosis 9 cognitive science LMM datasets
Metrics: wall-clock running time effective sample size ESS per iteration divergent transitions R-hat
Baselines: No marginalization HMC Non-centered reparameterization Naive marginalization Marginalized MCMC (automatic scalar marginalization)
Builds on: NUTS NumPyro Stan BRMS marginalized MCMC
Closest papers by meaning
- Automatically marginalized MCMC in probabilistic programming (ICML 2023 · similarity 0.90)
- Mixed Hamiltonian Monte Carlo for Mixed Discrete and Continuous Variables (NeurIPS 2020 · similarity 0.86)
- Parameter elimination in particle Gibbs sampling (NeurIPS 2019 · similarity 0.85)
- Hamiltonian Monte Carlo using an adjoint-differentiated Laplace approximation: Bayesian inference for latent Gaussian models and beyond (NeurIPS 2020 · similarity 0.85)
- Learning Likelihood-Free Reference Priors (ICML 2025 · similarity 0.83)
- Causal Effect Identification in lvLiNGAM from Higher-Order Cumulants (ICML 2025 · similarity 0.82)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.