JADAI: Jointly Amortizing Adaptive Design and Bayesian Inference

TL;DR

JADAI addresses the problem of jointly optimizing data acquisition and Bayesian inference in simulation-based experiments. It proposes an end-to-end, amortized framework that learns a design policy $\pi_\phi$, a history encoder $\eta_\omega$, and a diffusion/flow-based posterior estimator $q_\psi(\boldsymbol{\theta} \mid \mathbf{h}_t)$ to update beliefs at every step. By optimizing a surrogate incremental posterior loss that approximates the expected information gain ($\mathrm{EIG}$), JADAI delivers scalable, multimodal posteriors and competitive performance across LF, CES, and ID benchmarks. The approach enables fast rollouts and adaptable inference, with potential for black-box simulators and broader scientific applications.

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.

Figures (8)

• Figure 1: Overview of amortized SBI, BAD, and our proposed JADAI framework. Left: Amortized SBI, where a neural posterior estimator $q_\psi(\boldsymbol{\theta} \mid \mathbf{x})$ is trained on simulator pairs $(\mathbf{x},\boldsymbol{\theta})$ under a fixed design $\boldsymbol{\xi}$. Middle: Amortized BAD, where a designer runs a full $T$-step rollout for a fixed parameter $\boldsymbol{\theta}$: a policy $\pi_\phi$ maps the history state $\mathbf{h}_{t-1} = \eta_\omega(\{(\boldsymbol{\xi}_k,\mathbf{x}_k)\}_{k=1}^{t-1})$ to a new design $\boldsymbol{\xi}_t$, and the simulator returns $\mathbf{x}_t$, iteratively forming the history ${\mathcal{H}}_T = \{(\boldsymbol{\xi}_t,\mathbf{x}_t)\}_{t=1}^T$. Right: JADAI embeds this designer (middle block) into the amortized SBI loop and jointly trains $\pi_\phi$, $\eta_\omega$, and $q_\psi$ end-to-end, amortizing both experimental design and posterior inference across rollouts and experiments.
• Figure 2: Rollout process for Location Finding. Panels: posterior samples and chosen designs over time $t$, with crosses marking the true source locations. The second posterior mode is typically uncovered around $t=10$ measurements. Bottom right: corner plot of the learned posterior over the two sources at $t=10$ shows nearly identical densities at $(\theta_{11},\theta_{12})$ and $(\theta_{21},\theta_{22})$, indicating that the model correctly captures exchangeability of the two modes, that is, $p([\theta_{11},\theta_{12}], [\theta_{21},\theta_{22}] \mid \mathbf{h}_{10}) = p([\theta_{21},\theta_{22}], [\theta_{11},\theta_{12}] \mid \mathbf{h}_{10})$.
• Figure 3: Image discovery rollout. Each column is one measurement step. From top to bottom: ground-truth digit, cumulative measurements with the newly chosen design highlighted in blue, one posterior sample, and the posterior mean over 100 samples. The correct digit class is typically identified after 1–2 measurements, after which the policy mainly refines local structure in uncertain regions.
• Figure 4: Validation SSIM ($\uparrow$) and NRMSE ($\downarrow$) as a function of the number of measurements for CoDiff and our methods (using diffusion models). Shaded bands indicate the interquartile range over the validation set. Our learned methods achieve better SSIM and NRMSE than CoDiff and the random baselines at all steps and remain robust to additive measurement noise.
• Figure 5: Additional rollout examples for location finding (DAD setting, cf. \ref{['tab:lf-settings']}). Each row shows one simulated instance, and columns correspond to $t \in \{1,3,5,7,10,15,20,30\}$. The Gaussian prior over source locations biases early posteriors toward the domain center. By around $t=10$, posterior mass is typically concentrated near the true locations. Even when a location is identified early, the policy continues to explore rather than collapsing onto a single point; later measurements refine the inferred posteriors, with designs concentrating around the two source locations instead of being scattered across the domain.