JADAI: Jointly Amortizing Adaptive Design and Bayesian Inference

TL;DR

JADAI is a framework that combines adaptive experimental design and Bayesian inference into a single trainable system, enabling efficient parameter estimation with high-dimensional posteriors.

Contribution

It introduces a joint amortization approach that trains a policy, history network, and inference network end-to-end for adaptive design and inference.

Findings

Achieves superior or competitive performance on adaptive design benchmarks.

Utilizes diffusion-based posterior estimators for high-dimensional, multimodal posteriors.

Effectively reduces posterior error across experimental sequences.

Abstract

We consider problems of parameter estimation where design variables can be actively optimized to maximize information gain. To this end, we introduce JADAI, a framework that jointly amortizes Bayesian adaptive design and inference by training a policy, a history network, and an inference network end-to-end. The networks minimize a generic loss that aggregates incremental reductions in posterior error along experimental sequences. Inference networks are instantiated with diffusion-based posterior estimators that can approximate high-dimensional and multimodal posteriors at every experimental step. Across standard adaptive design benchmarks, JADAI achieves superior or competitive performance.