Amortized Bayesian Experimental Design with In-Context Knowledge Conditioning

Amortized Bayesian experimental design (BED) enables real-time design strategies by shifting acquisition cost offline, but existing methods remain tied to the prior and task distribution they were trained on and cannot exploit additional information available at deployment. We introduce in-context amortized BED with unified knowledge encoding (IMBUE), an amortized BED framework that accepts two forms of such information without retraining, namely user-specified prior knowledge and observations from previous instances of the experiment. Both are encoded as additional input tokens and processed in-context with the accumulated observations by a shared Transformer, and a learned reliability filter scores each auxiliary token against the accumulated observations and excludes inconsistent ones. Across four standard BED benchmarks, IMBUE accelerates early-stage information acquisition when the auxiliary knowledge is reliable, and remains close to the no-auxiliary baseline when it is not.