JADAI: Jointly Amortizing Adaptive Design and Bayesian Inference

2025-12-28

AI Summary

Accurate parameter estimation in active experimentation remains challenging due to sequential decision-making under uncertainty, high-dimensional and multimodal posteriors, and computational inefficiency in Bayesian inference. Method: This paper proposes an end-to-end learnable adaptive experimental design framework that jointly optimizes three components: a policy network (for selecting design variables), a history encoding network (to model temporal dependencies across experiment sequences), and a diffusion-based inference network (to approximate complex, potentially multimodal posteriors). The framework is trained via incremental posterior error minimization, enabling efficient sequential posterior modeling. Contribution/Results: It introduces the first amortized learning approach that unifies adaptive experimental design and Bayesian inference. Evaluated on multiple standard benchmarks, the method achieves state-of-the-art or superior performance in both posterior estimation accuracy and experimental efficiency—reducing the number of required experiments while improving posterior fidelity.

Technology Category

• Machine Learning: Bayesian Learning
• Search and Optimization: Sampling/Simulation-based Search
• Reasoning under Uncertainty: Stochastic Optimization

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Economics, Online Markets and Human Computation: Incentives in network design for Web infrastructures and ecosystems
• User Modeling, Personalization and Recommendation: Practical large-scale studies of user experience

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.

Problem

• Optimizes design variables for maximum information gain
• Jointly trains networks for adaptive design and Bayesian inference
• Approximates high-dimensional multimodal posteriors using diffusion estimators

Innovation

• Jointly amortizes Bayesian adaptive design and inference
• Uses diffusion-based posterior estimators for high-dimensional posteriors
• Trains policy, history, and inference networks end-to-end