Step-DAD: Semi-Amortized Policy-Based Bayesian Experimental Design

Marcel Hedman, Desi R. Ivanova, Cong Guan, Tom Rainforth

TL;DR: We develop a semi-amortized, policy-based, approach to Bayesian experimental design (BED) called Stepwise Deep Adaptive Design (Step-DAD), demonstrating superior decision-making and robustness compared with current state-of-the-art BED methods.

Abstract: We develop a semi-amortized, policy-based, approach to Bayesian experimental design (BED) called Stepwise Deep Adaptive Design (Step-DAD). Like existing, fully amortized, policy-based BED approaches, Step-DAD trains a design policy upfront before the experiment. However, rather than keeping this policy fixed, Step-DAD periodically updates it as data is gathered, refining it to the particular experimental instance. This test-time adaptation improves both the flexibility and the robustness of the design strategy compared with existing approaches. Empirically, Step-DAD consistently demonstrates superior decision-making and robustness compared with current state-of-the-art BED methods.

Lay Summary:

The objective of Bayesian experimental design is to collect data that is informative as possible under a given generative model. In this setting, it’s often critical to adapt your future collection strategy as new data becomes available. Traditional Bayesian methods handle this by continually updating beliefs and selecting the next best experiment—but they are computationally expensive and myopic, focusing only one step ahead. More scalable approaches pre-train a neural network to map histories to future design choices, but these fixed policy networks can fail when real data deviates from training conditions.

We introduce Step-DAD, a hybrid method that starts by pre-training a policy but then iteratively updates it during deployment. This refinement enables decisions tailored to the observed data, improving robustness and effectiveness without incurring the full cost of traditional methods. We find that Step-DAD empirically achieves state-of-the-art performance across multiple problems.

Primary Area: Probabilistic Methods->Bayesian Models and Methods

Keywords: Bayesian experimental design, Bayesian optimal design, Bayesian adaptive design, adaptive design optimization, information maximization

Submission Number: 3635