ADAVI: Automatic Dual Amortized Variational Inference Applied To Pyramidal Bayesian Models

TL;DR

This paper introduces ADAVI, a novel neural network-based variational inference method for large-scale hierarchical Bayesian models, enabling scalable and efficient posterior inference in high-dimensional neuroimaging data.

Contribution

ADAVI automatically constructs a dual variational family using attention-based encoders and normalizing flows, reducing complexity while maintaining expressivity for large hierarchical models.

Findings

Successfully applied to simulated data demonstrating scalability.

Effective in high-dimensional brain parcellation experiments.

Reduces parameterization complexity significantly.

Abstract

Frequently, population studies feature pyramidally-organized data represented using Hierarchical Bayesian Models (HBM) enriched with plates. These models can become prohibitively large in settings such as neuroimaging, where a sample is composed of a functional MRI signal measured on 300 brain locations, across 4 measurement sessions, and 30 subjects, resulting in around 1 million latent parameters.Such high dimensionality hampers the usage of modern, expressive flow-based techniques.To infer parameter posterior distributions in this challenging class of problems, we designed a novel methodology that automatically produces a variational family dual to a target HBM. This variational family, represented as a neural network, consists in the combination of an attention-based hierarchical encoder feeding summary statistics to a set of normalizing flows. Our automatically-derived neural network exploits exchangeability in the plate-enriched HBM and factorizes its parameter space. The resulting architecture reduces by orders of magnitude its parameterization with respect to that of a typical flow-based representation, while maintaining expressivity.Our method performs inference on the specified HBM in an amortized setup: once trained, it can readily be applied to a new data sample to compute the parameters' full posterior.We demonstrate the capability and scalability of our method on simulated data, as well as a challenging high-dimensional brain parcellation experiment. We also open up several questions that lie at the intersection between normalizing flows, SBI, structured Variational Inference, and inference amortization.