Amortising Bayesian Experimental Design for Sequential Information Gathering in LLMs
2026-07-03 • Machine Learning
Machine LearningArtificial Intelligence
AI summaryⓘ
The authors created a method called ASIG that helps large language models (LLMs) gather information better when they need to ask questions step-by-step. They trained the models to plan their questions more smartly using a technique based on Bayesian Experimental Design. Testing showed that ASIG more than doubled the success of a base model on a guessing game and worked much faster than other methods. It also helped the models do better on a new medical diagnosis task, even though they hadn’t seen that type of data before. This suggests their approach makes LLMs better and more efficient at gathering information over many turns.
Large Language ModelsSequential Decision-MakingBayesian Experimental DesignInformation GatheringPolicy OptimizationExpected Information GainMulti-Turn InteractionsFine-TuningInference CostOut-of-Distribution Generalization
Authors
Jakob Hartmann, James Harvey, Jhonathan Navott, Erik Y. Wang, Luckeciano C. Melo, Flaviu Cipcigan, Cheng Zhang, Alessandro Abate
Abstract
Large language models (LLMs) exhibit strong reasoning and world-knowledge capabilities, yet often struggle to gather information effectively across the multi-turn interactions required in sequential decision-making settings. We introduce Amortised Sequential Information Gathering (ASIG), a fine-tuning approach that amortises Bayesian Experimental Design (BED) into LLM policies via a multi-turn extension of Group Relative Policy Optimisation with an Expected Information Gain reward. Evaluated on the 20 Questions task, ASIG more than doubles the success rate of the 7B base model and reduces inference cost by over $25\times$ relative to BED-LLM, a competitive inference-time baseline. Applied to MediQ, a medical diagnosis benchmark unseen during training, ASIG improves information-seeking performance at the 7B scale, suggesting that the learned strategies can transfer out of distribution. Our findings show that amortising BED into LLM policies provides an effective and computationally efficient approach to sequential information gathering.