JADAI: Jointly Amortizing Adaptive Design and Bayesian Inference

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 without density evaluations. Inference networks are instantiated with diffusion models that can approximate high-dimensional and multimodal posteriors at every experimental step. JADAI achieves superior or competitive performance across adaptive design benchmarks.

Many scientific and engineering problems involve quantities that cannot be measured directly. Researchers often work with computational models to link these quantities to empirical observations. In doing so, they face two hard questions: which measurement should be taken next, and what does the data reveal? We introduce JADAI, a simulation-based framework that addresses the two questions simultaneously. After each measurement, it chooses the next one and updates the range of plausible values for the hidden quantities explicitly. Crucially, JADAI learns from simulated measurements and uses flexible generative models to handle cases where several explanations fit the data. On standard test problems, JADAI matches or improves upon earlier automatic measurement methods. This also holds for larger problems, where earlier methods were either much slower or too restrictive. This makes JADAI a step toward faster, more reliable and accurate estimation when researchers can choose what data to collect next.