metabeta -- A fast neural model for Bayesian mixed-effects regression

Hierarchical data with multiple observations per group is ubiquitous in empirical sciences and is often analyzed using mixed-effects regression. In such models, Bayesian inference gives an estimate of uncertainty but is analytically intractable and requires costly approximation using Markov Chain Monte Carlo (MCMC) methods. Neural posterior estimation shifts the bulk of computation from inference time to pre-training time, amortizing over simulated datasets with known ground truth targets. We propose metabeta, a neural network model for Bayesian mixed-effects regression. Using simulated and real data, we show that it reaches stable and comparable performance to MCMC-based parameter estimation at a fraction of the usually required time, enabling new use cases for Bayesian mixed-effects modeling.

💡 Research Summary

The paper introduces “metabeta,” a neural network–based amortized inference framework for Bayesian mixed‑effects regression. Mixed‑effects models are essential for hierarchical data because they capture both population‑level (fixed) effects and group‑specific (random) effects. Traditional Bayesian inference for such models relies on Markov Chain Monte Carlo (MCMC), especially Hamiltonian Monte Carlo (HMC), which provides accurate posterior estimates but is computationally intensive, often requiring long runtimes and careful tuning of hyper‑parameters.
Metabeta addresses these limitations by employing Neural Posterior Estimation (NPE), a simulation‑based approach that shifts most of the computational burden to a pre‑training phase. The authors generate millions of synthetic hierarchical datasets by sampling priors for fixed effects β, random‑effect covariance S, noise variance σ², and group‑specific random effects α_i, then drawing predictors X from both synthetic distributions and real benchmark datasets (PMLB, SRM). For each simulated dataset, outcomes y are produced via the standard linear mixed‑effects equation y_i = X_iβ + Z_iα_i + ε_i.
The model architecture consists of two main components: (1) a summary network and (2) posterior networks. The summary network uses a Set‑Transformer to produce permutation‑invariant local summaries for each group and a global summary across groups. This hierarchical summarization respects the exchangeable structure of mixed‑effects data. The posterior networks are conditional normalizing flows that take the summaries and the user‑specified priors as inputs. Separate flows are trained for global parameters (β, S, σ²) and for each group’s random effects α_i. The flows employ conditional affine coupling layers and a multivariate t‑distribution base, enabling flexible, tractable density estimation.
Training minimizes the forward Kullback‑Leibler divergence between the true posterior (available because the data are simulated) and the flow‑based approximation. The loss is back‑propagated through both the summary and posterior networks using a Schedule‑Free AdamW optimizer. Separate models are trained for different numbers of fixed and random effects, requiring between 10⁵ and 10⁶ simulated datasets to converge.
Because a finite network cannot perfectly represent the true posterior, the authors add a post‑hoc refinement step. Importance sampling re‑weights samples from the learned posterior using the exact likelihood p(D|θ) and the prior, improving point estimates and credible‑interval calibration. To address the tendency of normalizing flows to produce overly wide credible intervals, conformal prediction is applied on a calibration set, adjusting interval widths without retraining the network.
Empirical evaluation covers three scenarios: (a) toy data with simple, uncorrelated normal predictors, (b) in‑distribution data where predictors are drawn from real datasets, and (c) out‑of‑distribution data where the true parameters are unknown and outcomes are left untouched. Across all settings, metabeta’s posterior predictive performance (measured by negative log‑likelihood) closely matches that of HMC, while inference time is reduced by one to two orders of magnitude. Variational inference (VI) is also used as a baseline; metabeta outperforms VI in both accuracy and uncertainty calibration. Coverage error analyses show that the importance‑sampling and conformal‑prediction steps bring empirical coverage close to nominal levels.
The paper acknowledges limitations: the prior distribution is fixed at training time, so changing priors requires retraining; the current implementation assumes linear relationships for both fixed and random effects, limiting applicability to nonlinear mixed‑effects models; and the quality of the flow approximation may degrade for highly complex or high‑dimensional hierarchical structures. Nonetheless, the open‑source PyTorch implementation and the demonstrated speed‑accuracy trade‑off make metabeta a compelling tool for researchers in ecology, psychology, education, pharmacology, and other fields that routinely employ Bayesian mixed‑effects regression. By amortizing inference, metabeta enables rapid, scalable Bayesian analysis that was previously prohibitive with traditional MCMC methods.