论文

metabeta开源:用神经网络做GLMM贝叶斯推断,比NUTS快百到千倍

Prior-Amortized In-Context Bayesian Inference for Generalized Linear Mixed-Effects Models

精选理由

跑GLMM贝叶斯推断的试试metabeta,开源开放权重,比NUTS快两三个数量级,后验基本一致。

metabeta是一个预训练神经网络,对分层模型GLMM做先验摊销的in-context贝叶斯推断,先验族和超参数在测试时作为输入,实现零样本泛化。架构上用两个set transformer加条件归一化流,分别对应跨组共享的全局参数和每组各自的局部参数,训练数据为数百万个覆盖连续、二分类和计数结果的模拟数据集。默认流程由Independence Metropolis-Hastings对非归一化后验做精修,正确性依赖采样器而非网络本身,无需调参,速度比NUTS快2到3个数量级;该归一化流也可为NUTS热启动,换取更高速度和稳定性。在带真值参数的受控基准上,metabeta在参数恢复、校准和样本外预测上与NUTS持平;在分布外真实数据集上后验与NUTS基本吻合,并在似然误设、共线设计和数据稀少情形下保持稳健。模型开源且开放权重,可直接部署。

原文 · arXiv cs.LG

Prior-Amortized In-Context Bayesian Inference for Generalized Linear Mixed-Effects Models

Hierarchical data is ubiquitous in the empirical sciences and is most commonly analyzed with generalized linear mixed-effects models (GLMMs). Bayesian inference for GLMMs yields calibrated uncertainty but requires MCMC; the No-U-Turn Sampler (NUTS) is the gold standard but is slow and must restart from scratch for every new dataset, model and prior. We introduce metabeta, a pretrained neural network for prior-amortized in-context Bayesian inference over GLMMs. Unlike previous neural posterior estimators that fix the prior at training time, metabeta accepts prior families and hyperparameters as inputs at test time, enabling zero-shot generalization. Two set transformers and conditional normalizing flows mirror the posterior's two-level structure (global parameters shared across groups, local parameters per group). The model is trained on millions of realistic simulated datasets spanning continuous, binary, and count outcomes. By default, the flow posterior is refined by Independence Metropolis-Hastings against the unnormalized posterior, so its correctness rests on the sampler rather than the network; this yields tuning-free inference two to three orders of magnitude faster than NUTS. Alternatively, the flow can warm-start NUTS, giving nearly identical inference with substantially increased speed and stability. On controlled benchmarks with ground-truth parameters, metabeta matches NUTS in parameter recovery, calibration and out-of-sample prediction. On out-of-distribution real datasets, its posteriors closely match those of NUTS across all parameter types, and they remain faithful under misspecified likelihoods and priors, out-of-distribution predictors, collinear designs, and data-poor regimes. The model is open-source and open-weights and thus immediately deployable.