Scalable Expectation Propagation for Mixed-Effects Regression

TL;DR

This paper introduces a scalable, numerically stable expectation propagation framework for mixed-effects regression that efficiently handles large, grouped datasets by ensuring linear scaling and distributed computation.

Contribution

It develops a novel EP method with sparse reparameterization and moment propagation for mixed-effects models, addressing scalability and stability in distributed Bayesian inference.

Findings

Achieves linear scaling with the number of groups

Maintains accuracy comparable to existing scalable methods

Enables distributed inference for large datasets

Abstract

Mixed-effects regression models represent a useful subclass of regression models for grouped data; the introduction of random effects allows for the correlation between observations within each group to be conveniently captured when inferring the fixed effects. At a time where such regression models are being fit to increasingly large datasets with many groups, it is ideal if (a) the time it takes to make the inferences scales linearly with the number of groups and (b) the inference workload can be distributed across multiple computational nodes in a numerically stable way, if the dataset cannot be stored in one location. Current Bayesian inference approaches for mixed-effects regression models do not seem to account for both challenges simultaneously. To address this, we develop an expectation propagation (EP) framework in this setting that is both scalable and numerically stable when distributed for the case where there is only one grouping factor. The main technical innovations lie in the sparse reparameterisation of the EP algorithm, and a moment propagation (MP) based refinement for multivariate random effect factor approximations. Experiments are conducted to show that this EP framework achieves linear scaling, while having comparable accuracy to other scalable approximate Bayesian inference (ABI) approaches.