metabeta

metabeta is an amortized model for Bayesian generalized linear mixed-effects models (GLMMs). Given a grouped dataset, a model formula, and an optional prior specification, it returns posterior samples for fixed effects, random effects, variance components, and correlations with the speed and batching capabilities of a PyTorch forward pass.

• Supports Normal, Bernoulli, and Poisson outcomes.
• Accepts convenient lme4-style formulas.
• Conditions on prior family and hyperparameters at inference time.
• Supports batched prior-sensitivity analysis in one sample() call.
• Includes diagnostics for posterior contraction and prediction accuracy.
• Ships the pretrained inference API, plotting helpers, diagnostics, and demo notebooks.
• Keeps simulation, training, reference-method evaluation, and experiment scripts in the source repository for research workflows.

Pretrained checkpoints are hosted on Hugging Face and are downloaded automatically on first use.

Quick start

Install the pretrained-model inference package from PyPI:

pip install metabeta

For local development, install from source with uv:

git clone https://github.com/adkipnis/metabeta.git
cd metabeta
uv sync
uv pip install -e .

The default PyPI install is for pretrained-model inference. To reproduce training runs, generate synthetic datasets, benchmark against reference methods, or extend the package, clone the repository and install the optional research dependencies:

git clone https://github.com/adkipnis/metabeta.git
cd metabeta
uv sync --extra research --group simulation --dev
uv pip install -e ".[research]"

Run posterior inference on a grouped dataframe:

import statsmodels.api as sm
from metabeta.models.api import Api

mb = Api.from_pretrained("normal")
df = sm.datasets.get_rdataset("sleepstudy", "lme4").data
result = mb.sample(df, formula="Reaction ~ Days + (Days | Subject)", n_samples=1000)
print(mb.posteriorSummary(result))

With n_samples >= 64, sample also refines the flow posterior by Independence Metropolis-Hastings (the default, refine=True). The chain uses the flow as its proposal, so the refined posterior does not hinge on network accuracy alone; per-dataset acceptance rates and a suggested pool size are reported in result.safeguards. If the suggestion exceeds the n_samples you requested, rerun with the larger value (or pass refine=False to keep the raw flow posterior).

See demos/intro.ipynb for the full sleepstudy walkthrough and demos/priors.ipynb for an exemplary prior-sensitivity analysis.

From simulation to deployment

Each pretrained model is built through the same pipeline, from a dataset simulator to a checkpoint that can be loaded by the public API.

1. Define the model family. Choose the likelihood, GLMM dimensions, group structure, covariate styles, and prior families covered by a checkpoint.
2. Generate training data. Simulate hierarchical datasets and posterior reference targets across the configured design space.
3. Format inputs. Convert grouped tabular data into padded tensors, masks, and prior encodings shared by training and inference.
4. Train amortized posteriors. Use set-transformers (to learn low-dimensional permutation-invariant summaries of the datasets) and conditional coupling flows to approximate posteriors for fixed effects, variance parameters, correlations, and group-wise random effects.
5. Validate behavior. Check parameter recovery, credible interval coverage, simulation-based calibration, posterior predictive accuracy and run comparisons against reference methods.
6. Package checkpoints. Bundle weights, configs, routing metadata, and preprocessing expectations into a joint checkpoint.
7. Deploy through the API. Load joint checkpoints with Api.from_pretrained(...) for convenient posterior estimation and diagnostics.

Repository map