Amortizing Sequential Experimental Design in Reasoning LLMs

AR-BED trains a reasoning LLM with information-gain rewards under an externally specified finite Bayesian experimental-design model, improving adaptive questioning and transferring to a clinical dialogue task.

LLMs struggle to adaptively gather information in multi-turn experimental design and can be insensitive to experimental feedback, relying on prior knowledge. Sequential Bayesian experimental design selects informative actions but requires costly posterior and expected-information-gain estimation at every step, and one-step EIG can be myopic. The paper asks how to amortize sequential BED into an LLM's reasoning process and which generative-model and reward choices make this work.

The policy outputs a reasoning trace followed by a design and is trained with GRPO using information-gain rewards. The experimental generative model is specified independently of the policy: a fixed uniform prior over a finite contrastive hypothesis set and a separate larger LLM as semantic likelihood evaluator over a finite outcome space. The reward is the expected one-step EIG, marginalizing over hypotheses and outcomes as a Rao-Blackwellized realized reward; returns sum discounted per-step EIG over K-step rollouts. Training uses Twenty Questions with a 200-animal contrastive set, 100 targets, 250 GPT-5 games yielding 1,200 partial trajectories; AR-BED_K is a Gemma-4-E4B-it policy with rollout horizon K and gamma=0.5 for multi-step training. The paper ablates direct non-reasoning policies, misspecified priors, realized versus expected rewards, and multi-step objectives.

• AR-BED1 achieves the best final solve rate across all categories and the highest average EIG per turn on Animals, Plants, and Countries. §5.1
• AR-BED1 is much faster than inference-time BED-LLM per game and than training-time ASIG per step. §5.2
• After Twenty Questions training, medical AR-BED1 improves diagnostic accuracy on held-out dermatology cases and on other specialties. §5.4
• Restricting the training generative model to 100 mammals substantially reduces acquisition on animal targets relative to 100 diverse animals. §B.1
• Expected information rewards generally improve acquisition over realized rewards, except on Celebrities. §B.2
• In Cube-NM, the two-step policy learns to select an initially uninformative context test and overtakes the one-step policy after three tests. §B.3.1

• The paper studies relatively short rollout horizons and leaves extending to longer horizons with a reliable training signal as future work (stated in section 6).
• Performance is sensitive to the coverage of the contrastive hypothesis set when it is misspecified relative to the true hypothesis space; size, composition, and prior weighting need further study (stated in section 6).
• Evaluation repeats use different game target order or decoding seeds, not independent training seeds, so variability across training runs is not estimated (stated in section C.1).
• Cube-NM results are matched single-training-seed results, not estimates of variability across training runs (stated in section B.3.1).

• Training-time amortization decouples the experimental generative model from the policy, so the information objective stays fixed while the policy updates. §4.1
• The generative model uses a uniform prior over a fixed finite contrastive set and a separate larger LLM as semantic likelihood evaluator over a finite outcome space. §4.1
• The expected information reward marginalizes over hypotheses and outcomes, equals one-step EIG under the finite contrastive model, and is accumulated over discounted K-step rollouts. §4.2
• Adding reasoning at inference alone reduces solve rate across categories and increases average EIG only on Countries; AR-BED1 instead achieves the best final solve rate across all categories. §5.1
• AR-BED1 outperforms BED-LLM and ASIG on most domains, attributed to amortizing the information-gain objective into reasoning and using an external generative model. §5.2
• A mammals-only contrastive set gives a flawed signal, for example Is it a mammal receives EIG approximately 0, and hurts performance on both non-mammals and mammals. §5.3.1
• Multi-step rewards learn non-myopic context acquisition in Cube-NM, while Twenty Questions gains are smaller and concentrated at longer evaluation horizons. §5.3.3
• A policy trained on Twenty Questions transfers to MediQ clinical dialogue, improving accuracy on held-out dermatology diagnoses and on other specialties. §5.4