🤖 AI

JADAI: Jointly Amortizing Adaptive Design and Bayesian Inference

JADAI is an end-to-end trained framework that jointly amortizes Bayesian adaptive design and inference using diffusion-based posterior estimators to maximize information gain and achieve superior performance in parameter estimation tasks.

Original authors: Niels Bracher*, Lars Kühmichel*, Desi R. Ivanova, Xavier Intes, Paul-Christian Bürkner, Stefan T. Radev, * equal contribution

Published 2026-07-07✓ Author reviewed ⓘ

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Niels Bracher*, Lars Kühmichel*, Desi R. Ivanova, Xavier Intes, Paul-Christian Bürkner, Stefan T. Radev, * equal contribution

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a detective trying to solve a mystery, but you can't just look at the crime scene. Instead, you have a magical machine that can simulate what would happen if you asked a specific question or looked in a specific direction. However, the machine is slow, and you only have a limited number of questions you can ask before you run out of time.

Your goal is twofold:

The Detective Work (Inference): Figure out the truth (the hidden parameters) based on the clues you get.
The Strategy (Design): Decide where to look next to get the most useful clue possible.

Usually, scientists treat these as two separate jobs. One team figures out the best place to look, and another team tries to solve the puzzle once all the clues are gathered. This paper introduces JADAI, a new framework that combines these two jobs into one super-efficient team.

Here is how JADAI works, using simple analogies:

1. The "Amortized" Team (The One-Stop Shop)

In the old way, every time you wanted to solve a new mystery, you had to start from scratch: "Okay, where should I look? Let me calculate that. Now, let me analyze the clues. Okay, now I know the answer." This is slow.

JADAI is like a seasoned detective agency that has trained on millions of fake mysteries. Instead of calculating the strategy and the solution from scratch every time, they have "amortized" (spread out) the learning. They have trained three specific AI "agents" to work together instantly:

The Strategist (Policy): Decides where to look next.
The Scribe (History Network): Remembers everything that has happened so far and summarizes it into a neat note.
The Solver (Inference Network): Looks at the note and immediately guesses the answer.

2. The "Diffusion" Solver (The Sculptor)

To solve the puzzle, the "Solver" agent uses something called a Diffusion Model. Imagine a sculptor who starts with a block of marble covered in noise (static). They slowly chip away the noise, step by step, until a clear statue emerges.

In JADAI, the "statue" is the answer to the mystery.
The "chipping away" is the AI refining its guess based on the clues.
This is powerful because it can handle messy, complex answers (like a statue with two heads or a twisted shape) that simpler methods would miss.

3. The "Telescoping" Training (Learning from the Journey)

How do you teach these three agents to work together? You don't just tell them, "Get the right answer at the very end." That's too hard.

Instead, JADAI uses a telescoping loss. Imagine you are walking up a mountain.

Old Way: You only get a reward when you reach the very peak. If you take a wrong turn early on, you don't know until you get stuck at the top.
JADAI Way: You get a tiny "pat on the back" (a reward signal) at every single step of the hike. If you take a step that makes the view clearer (reduces uncertainty), you get a reward. If you take a step that makes things confusing, you get a penalty.

This teaches the Strategist to pick the best path and teaches the Solver to get better at guessing along the way. They learn together, end-to-end, without needing to check the "true answer" until the very end of training.

4. The Results: Smarter and Faster

The paper tested this on three different "mysteries":

Finding Hidden Sources: Like trying to locate two hidden speakers in a room by listening to the sound at different spots. JADAI figured out where to listen to find both speakers faster and more accurately than previous methods.
Consumer Choices: Figuring out what a person prefers when choosing between baskets of goods. The "sweet spot" for asking questions is very narrow, but JADAI learned to navigate it perfectly.
Reconstructing Images: Imagine trying to guess a digit (like a "7") by only seeing small, blurry patches of it one by one. JADAI learned to pick the most revealing patches first, reconstructing the image much better than other methods, even when there was "noise" (static) in the signal.

The Bottom Line

JADAI is a framework that teaches an AI to ask the right questions and solve the puzzle at the same time. By training them together to improve at every single step of the process, it can solve complex, high-dimensional problems (like finding hidden objects or reconstructing images) much faster and more accurately than methods that treat "asking" and "answering" as separate tasks.

It's like upgrading from a detective who has to re-investigate the whole case every time they get a new clue, to a detective who has a perfect memory, a perfect strategy, and a perfect intuition, all working in perfect sync.

Technical Summary: JADAI – Jointly Amortizing Adaptive Design and Bayesian Inference

Problem Statement
Many scientific and engineering domains involve inferring unobserved parameters $\theta$ from observations $x$ generated by complex simulators. While Simulation-Based Inference (SBI) has advanced the ability to approximate posterior distributions $p(\theta|x)$ without explicit likelihood evaluations, these methods typically assume a fixed experimental design. Conversely, Bayesian Experimental Design (BED) and Adaptive Design (BAD) focus on selecting design variables $\xi$ to maximize information gain (e.g., Expected Information Gain, EIG) but often treat inference as a separate, computationally expensive step performed after data collection. Existing approaches frequently decouple these tasks: BAD methods often rely on slow, non-amortized inference or simple parametric families, while SBI methods do not optimize data acquisition. This separation leads to suboptimal information gain and computational inefficiency, particularly in high-dimensional or multimodal settings.

Methodology
The authors introduce JADAI (Jointly Amortizing Adaptive Design and Bayesian Inference), a framework that unifies adaptive design and posterior inference into a single, end-to-end trainable system. The core innovation is the joint amortization of three components:

A Policy Network ( $\pi_\phi$ ): Proposes the next design $\xi_t$ based on the current experimental history.
A History/Summary Network ( $\eta_\omega$ ): Encodes the sequence of past design-observation pairs $\{(\xi_k, x_k)\}_{k=1}^t$ into a fixed-dimensional summary statistic $h_t$ .
An Inference Network ( $q_\psi$ ): Approximates the posterior distribution $p(\theta|h_t)$ at every step of the experimental sequence.

Optimization Objective
JADAI minimizes a generic loss function that aggregates incremental reductions in posterior error along experimental sequences. Instead of maximizing the Expected Information Gain (EIG) directly—which requires density evaluations—the framework utilizes a telescoping sum of posterior losses.

The training objective is derived from a variational lower bound on the Total Expected Information Gain (TEIG).
The loss at step $t$ , $\ell_t$ , is defined by the posterior estimator (e.g., diffusion or flow matching loss).
The utility function $u_T$ is constructed as a sum of differences between consecutive losses: $u_T = \sum (\ell^*_{t-1} - \ell_t)$ , where $\ell^*$ denotes a detached (stop-gradient) version of the previous loss.
This formulation ensures that the policy and summary networks are trained to maximize the reduction in posterior loss at every step, effectively maximizing information gain without requiring explicit density evaluations of the posterior.

Key Architectural Choices

Diffusion Models: The inference network $q_\psi$ is instantiated using diffusion models (or flow matching), allowing the framework to approximate high-dimensional and multimodal posteriors at every experimental step, a capability often lacking in previous BAD methods restricted to normalizing flows.
Training Stability: To avoid nested backpropagation through time (BPTT) and ensure stable training, the framework employs:
- Detached History States: Designs are generated from a detached summary state to prevent cyclic gradient paths between the policy and history network.
- Curriculum Learning: The rollout length is gradually increased during training, and a "design prior mix-in" strategy is used early in training to expose the networks to diverse designs before relying solely on the learned policy.

Results
JADAI was evaluated on three benchmarks representing increasing complexity in policy and posterior structure:

Location Finding (LF): A task requiring the inference of source locations from noisy intensity measurements. JADAI achieved superior or competitive sPCE (sequential Prior Contrastive Estimation) scores compared to state-of-the-art baselines (e.g., DAD, iDAD, RL-BOED), particularly for longer horizons ( $T > 10$ ). The learned policy effectively balanced exploration and exploitation, and the diffusion-based posterior correctly captured the multimodal and exchangeable nature of the sources.
Constant Elasticity of Substitution (CES): A challenging design problem involving preference inference where informative designs lie in a narrow "sweet spot." JADAI outperformed all existing baselines at the standard horizon $T=10$ , demonstrating robustness in scenarios where random designs are uninformative.
MNIST Image Discovery: A high-dimensional task where the goal is to reconstruct an image from partial measurements. JADAI achieved state-of-the-art Structural Similarity Index Measure (SSIM) and Normalized Root Mean Square Error (NRMSE) scores, outperforming the CoDiff baseline. The method successfully identified digit classes early and refined local details, maintaining performance even under additive measurement noise.

Significance and Claims
The paper claims that JADAI represents a significant step forward by unifying design and inference, thereby eliminating the need for separate, costly inference steps during the adaptive process. Key contributions include:

Joint Amortization: It is the first framework to jointly amortize adaptive design and posterior inference using a general utility that aggregates incremental posterior error reductions, avoiding the need for density evaluations.
Scalability to Complex Posteriors: By leveraging diffusion models, JADAI can handle high-dimensional and multimodal posteriors at every experimental step, overcoming the expressiveness limitations of previous methods that relied on normalizing flows or Gaussian mixtures.
Performance: The framework achieves superior or competitive performance across diverse benchmarks, including high-dimensional image reconstruction, while maintaining computational efficiency during deployment (rollouts completed in milliseconds).

The authors note that while JADAI currently relies on differentiable simulators, the approach is well-suited for the growing landscape of differentiable simulation frameworks. They acknowledge that posterior sampling remains a bottleneck for sub-second deployment in high-dimensional cases but suggest post-hoc distillation as a potential future direction. The work does not claim to solve black-box simulator optimization without gradients but positions itself as a robust solution for differentiable settings where classical pipelines are computationally impractical.