Hamiltonian Monte Carlo Inference of Marginalized Linear Mixed-Effects Models

2024-10-31

Neural Information Processing Systems

AI Summary

In Bayesian inference, linear mixed-effects models (LMMs) typically rely on Markov chain Monte Carlo (MCMC) sampling; explicit modeling of random effects induces cubic computational complexity—O(n³)—in Hamiltonian Monte Carlo (HMC). While manual marginalization of random effects reduces complexity to O(n), it is error-prone and burdensome for users. This paper introduces the first fully automated analytical marginalization algorithm for LMMs, integrating HMC with the matrix inversion lemma, Woodbury identity, and block-wise linear algebra to enable automatic, exact, and lossless marginalization of random effects. The method reduces the core HMC computational cost to O(n) and provides a theoretical guarantee of strict efficiency improvement across all applicable settings. Empirical evaluation on real-world cognitive science models demonstrates substantially faster convergence, markedly increased effective sample size, and zero requirement for user intervention or reparameterization.

Technology Category

Machine Learning: Probabilistic Circuits and Graphical ModelsReasoning under Uncertainty: Relational Probabilistic ModelsCognitive Modeling & Cognitive Systems: Conceptual Inference and Reasoning

Application Category

Economics, Online Markets and Human Computation: Cost models of using LLMs in production systemsGraph Algorithms and Modeling for the Web: Algorithms and analysis for incomplete, noisy, or partially observed Web-related graphsUser Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation

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.

Problem

• Improving Bayesian inference efficiency in linear mixed-effects models
• Automating marginalization of random effects to enhance HMC performance
• Reducing computational complexity from cubic to linear time

Innovation

• Algorithm for marginalizing random effects in LMMs
• Linear time using fast linear algebra
• Improves inference efficiency in cognitive science models