Given observed data and a probabilistic generative model, Bayesian inference searches for the distribution of the model's parameters that could have yielded the data. Inference is challenging for large population studies where millions of measurements are performed over a cohort of hundreds of subjects, resulting in a massive parameter space. This large cardinality renders off-the-shelf Variational Inference (VI) computationally impractical. In this work, we design structured VI families that efficiently tackle large population studies. Our main idea is to share the parameterization and learning across the different i.i.d. variables in a generative model, symbolized by the model's \textit{plates}. We name this concept \textit{plate amortization}. Contrary to off-the-shelf stochastic VI, which slows down inference, plate amortization results in orders of magnitude faster to train variational distributions. Applied to large-scale hierarchical problems, PAVI yields expressive, parsimoniously parameterized VI with an affordable training time. This faster convergence effectively unlocks inference in those large regimes. We illustrate the practical utility of PAVI through a challenging Neuroimaging example featuring 400 million latent parameters, demonstrating a significant step towards scalable and expressive Variational Inference.
翻译:给定观测数据与概率生成模型,贝叶斯推断旨在寻找能产生该数据的模型参数分布。对于涉及数百名受试者、执行数百万次测量的大规模群体研究,推断面临参数空间庞大的挑战,这使得现成变分推断在计算上不可行。本文设计了能高效处理大规模群体研究的结构化变分族。核心思想是通过模型中的“板”符号,在生成模型中不同独立同分布变量间共享参数化与学习过程,我们将其称为“板式摊销”。与降低推断效率的传统随机变分推断不同,板式摊销可将变分分布的训练速度提升数个数量级。应用于大规模层次化问题时,PAVI能以可负担的训练时间生成表达力强且参数简洁的变分分布,这种快速收敛特性有效解锁了大规模场景下的推断能力。通过一个包含4亿隐参数的神经影像学挑战性案例,我们展示了PAVI的实用价值,标志着向可扩展且表达力强的变分推断迈出重要一步。