JADAI: Jointly Amortizing Adaptive Design and Bayesian Inference

This repository contains the code for running and reproducing the experiments from the paper JADAI: Jointly Amortizing Adaptive Design and Bayesian Inference, published at ICML 2026 (arXiv version).

Installation

For improved reproducibility, we recommend installing the exact package versions using the provided lockfile with uv:

uv sync --locked

For a more relaxed installation, you can also permit resyncing some dependencies by omitting the --locked flag:

uv sync

The pip equivalent is:

python -m venv .venv/
source .venv/bin/activate
python -m ensurepip
python -m pip install .

Location Finding

To train the LF model at a time horizon $Q = 10$, with $K = 2$ source locations, evaluating using $L = 5 \cdot 10^5$ contrastive samples, use:

uv run --locked -m jadai --task lf -Q 10 --num-sources 2

CES

To train the CES model at a time horizon $Q = 10$, with $K = 3$ goods, evaluating using $L = 10^7$ contrastive samples, use:

uv run --locked -m jadai --task ces -Q 10 --num-goods 3

Image Discovery

To train the ID model at a time horizon $Q = 6$, use:

uv run --locked -m jadai --task id -Q 6

Cite

The article can be cited as:

@misc{bracher2025jadaijointlyamortizingadaptive,
      title={JADAI: Jointly Amortizing Adaptive Design and Bayesian Inference}, 
      author={Niels Bracher and Lars Kühmichel and Desi R. Ivanova and Xavier Intes and Paul-Christian Bürkner and Stefan T. Radev},
      year={2025},
      eprint={2512.22999},
      archivePrefix={arXiv},
      primaryClass={stat.ML},
      url={https://arxiv.org/abs/2512.22999}, 
}