🤖 AI

Amortising Bayesian Experimental Design for Sequential Information Gathering in LLMs

This paper introduces Amortised Sequential Information Gathering (ASIG), a fine-tuning method that integrates Bayesian Experimental Design into LLM policies to significantly improve sequential information gathering efficiency and success rates in tasks like 20 Questions and medical diagnosis while drastically reducing inference costs.

Original authors: Jakob Hartmann, James Harvey, Jhonathan Navott, Erik Y. Wang, Luckeciano C. Melo, Flaviu Cipcigan, Cheng Zhang, Alessandro Abate

Published 2026-07-07

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jakob Hartmann, James Harvey, Jhonathan Navott, Erik Y. Wang, Luckeciano C. Melo, Flaviu Cipcigan, Cheng Zhang, Alessandro Abate

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Lost in Conversation" AI

Imagine you are playing a game of 20 Questions with a very smart friend (a Large Language Model, or LLM). You are thinking of an object, like "a toaster," and your friend has to guess what it is by asking yes-or-no questions.

Ideally, your friend should ask smart questions that cut the possibilities in half every time (e.g., "Is it used in the kitchen?"). However, the paper points out that current AI models often get "lost in conversation." They might ask repetitive questions, forget what they already know, or fail to realize when they have enough information to make a guess. They are good at knowing facts, but bad at figuring out what to ask next to learn something new.

The Old Way: The "Over-Thinker" (Inference-Time Optimization)

Before this paper, researchers tried to fix this by making the AI "think harder" every time it needed to ask a question.

The Analogy: Imagine your friend stops after every question to pull out a giant whiteboard, write down every possible object in the world, calculate the probability of each one being the answer, and mathematically figure out which question would give the most information.
The Result: This works well, but it is incredibly slow and expensive. It's like hiring a team of mathematicians to help your friend ask a single question. The paper calls this Bayesian Experimental Design (BED). While effective, it's too heavy for real-time use.

The New Solution: ASIG (The "Muscle Memory" Approach)

The authors introduce a new method called ASIG (Amortised Sequential Information Gathering). Instead of making the AI do complex math every time it speaks, they teach the AI to learn the skill so it becomes second nature.

The Analogy: Instead of pulling out a whiteboard every time, you train your friend for weeks. You play thousands of rounds of 20 Questions with them. When they ask a bad question, you say, "No, that doesn't help much." When they ask a great question that narrows down the options, you say, "Great! That's exactly what we need."
The Magic: Over time, your friend stops needing the whiteboard. They develop "muscle memory." They instinctively know which questions are the most informative. They have "amortized" (spread out) the heavy thinking into their training, so now they can play the game fast and efficiently without stopping to calculate.

How They Trained the AI

The paper describes a specific training loop (Figure 1 in the paper):

The Proposer: The AI being trained (the questioner).
The Oracle: A super-smart, frozen AI that knows the secret answer and acts as the game master.
The Reward System: The Proposer gets points in two ways:
- The "Curiosity" Score (EIG): Did the question split the possibilities evenly? (e.g., "Is it alive?" is better than "Is it a toaster?").
- The "Success" Score: Did the AI eventually guess the right answer?
The Training: They used a method called GRPO (Group Relative Policy Optimization). Think of this as running a race where the AI tries 5 different questions at once. The one that gets the best score gets a "high five," and the AI learns to do that one more often.

What They Found (The Results)

The researchers tested this on the 20 Questions game and a medical diagnosis simulation (MediQ).

Much Better at Guessing: On the 20 Questions game, the trained AI (ASIG) more than doubled its success rate compared to the untrained base model. It got better at asking the right questions to narrow things down quickly.
Super Fast: Because the AI doesn't need to do the heavy math during the game, it is 25 to 36 times faster than the "Over-Thinker" method (BED-LLM). It's the difference between a sprinter and a marathon runner carrying a backpack.
It Works on New Stuff (Generalization): They tested the AI on a medical diagnosis task (MediQ) that it had never seen before. The 7-billion-parameter model (a smaller, efficient model) showed it could transfer its "asking skills" to medicine, becoming better at knowing when to ask a follow-up question and when to make a diagnosis.
No Brain Damage: Crucially, training the AI to be a better questioner didn't make it worse at other things. It didn't forget how to write poems or solve math problems; it just got better at gathering information.

The Takeaway

The paper shows that you don't need to force an AI to do complex math in real-time to make it a good investigator. Instead, you can train it to internalize the strategy of "asking the right questions." This makes the AI faster, cheaper to run, and just as good (or better) at finding the truth in a conversation.

Technical Summary: Amortising Bayesian Experimental Design for Sequential Information Gathering in LLMs

Problem Statement

Large Language Models (LLMs) possess extensive world knowledge and reasoning capabilities but frequently struggle with effective information gathering in sequential, multi-turn decision-making settings. While LLMs can reason, they often fail to recognize when additional information is required or to acquire that information efficiently during interactions. This limitation is critical in domains like healthcare and scientific experimental design, where maintaining coherent beliefs over extended interactions is essential.

Existing approaches to improve LLM information gathering fall into two categories:

Inference-time optimisation: Methods like Bayesian Experimental Design (BED) explicitly maximize Expected Information Gain (EIG) during generation. While principled, these methods suffer from high computational costs due to repeated sampling and belief maintenance, are often myopic, and require task-specific tailoring.
Weight-space optimisation: Fine-tuning methods that adapt model parameters. While these offer fast inference, previous approaches have struggled to generalize beyond specific training tasks or have relied on external retrievers to maintain belief states, limiting their self-contained applicability.

The core challenge is to develop a method that amortizes the computational cost of Bayesian Experimental Design into the model weights, enabling efficient, generalizable sequential information gathering without the heavy inference-time overhead of explicit BED.

Methodology: Amortised Sequential Information Gathering (ASIG)

The authors introduce ASIG, a fine-tuning framework that embeds Bayesian Experimental Design (BED) directly into an LLM's policy. The approach treats the LLM as a sequential BED agent that iteratively selects experiments (e.g., questions) to reduce uncertainty over a latent target.

Core Components

Training Environment: The method is instantiated using the 20 Questions game. Here, the LLM (Proposer) asks questions (experimental designs, $\xi_t$ ) to an Oracle LLM (with access to the hidden target $\hat{\theta}$ ) to identify the target within a budget of 20 turns. The binary nature of the outcomes ("Yes"/"No") allows for efficient EIG estimation.
Reward Signal: The training utilizes a multi-component reward signal within a Group Relative Policy Optimisation (GRPO) framework:
- Expected Information Gain (EIG): Calculated using a Rao-Blackwellised estimator. The Proposer LLM samples a set of plausible target beliefs ( $\Theta_t$ ) consistent with the history, computes likelihoods $p(y|\theta, \xi)$ for each, and estimates the reduction in entropy. This provides dense, turn-level feedback on question quality.
- Outcome Reward: A sparse reward based on task success (correctly guessing the target) and penalties for non-terminal states. This drives goal-directed behavior.
Multi-Turn GRPO Extension: The authors extend standard GRPO to handle multiple reward streams with distinct temporal horizons.
- The total reward $R_t$ for a rollout combines the EIG reward (weighted by $\alpha$ , horizon $E=1$ ) and the outcome reward (weighted by $\beta$ , horizon $O=3$ ).
- This structure encourages maximizing information gain in the immediate turn while optimizing for the probability of a correct guess within a short horizon (3 turns).
Learning Curriculum: To prevent the model from avoiding specific guesses, a curriculum is employed where the model is first trained without the EIG reward ( $\alpha=0$ ) to encourage hypothesis-seeking behavior, followed by the introduction of the EIG signal.

Key Contributions

ASIG Framework: A novel fine-tuning method that amortizes BED into LLM weights, enabling sequential information gathering without additional inference-time optimisation.
Extended GRPO: An adaptation of GRPO to support joint optimization of turn-level EIG and multi-horizon task success, handling distinct reward streams with different temporal scopes.
Performance and Efficiency: Demonstration that ASIG substantially improves performance on the 20 Questions task (doubling success rates for 7B models) while reducing inference costs by over 25× compared to competitive inference-time baselines (BED-LLM).
Out-of-Distribution (OOD) Generalization: Evidence that the learned information-seeking strategies transfer to the MediQ medical diagnosis benchmark at the 7B scale without in-domain fine-tuning, improving clinical reasoning accuracy and reducing overconfidence.

Results

20 Questions Task

Success Rate: On the 7B model (Qwen2.5-7B-Instruct), ASIG more than doubles the success rate of the baseline and outperforms BED-LLM in the "Animals" category. On the 14B model, it improves over the baseline by 4.2 and 1.8 percentage points on "Animals" and "Plants," respectively.
Behavioral Shift: ASIG-trained models exhibit a higher frequency of "hypothesis-seeking" questions (specific guesses) earlier in the interaction compared to baselines, particularly in the 7B model.
Efficiency: ASIG matches the inference runtime of the base LLM, whereas BED-LLM is 25× to 36× slower due to the need for repeated belief sampling and EIG maximization at every turn.

Generalization to MediQ

Applied to the MediQ benchmark (medical diagnosis), the 7B ASIG model improved accuracy by 0.9% and voluntary accuracy by 1.7%, while reducing the overconfidence rate from 36.1% to 33.4%.
The 14B model showed mixed results, with minimal changes in accuracy and a slight decrease in voluntary accuracy, suggesting that larger models with stronger priors may require less behavioral adaptation from fine-tuning.

Capability Retention

Evaluation on standard benchmarks (ARC-Challenge, HellaSwag, MMLU) confirmed that the fine-tuned models retained their general language capabilities, achieving near-identical scores to the base models.

Significance and Claims

The paper claims that amortising Bayesian Experimental Design into LLM policies provides a practical and scalable approach to test-time information acquisition. By shifting the computational burden of information gathering from inference-time optimization to model weights, ASIG achieves:

Computational Efficiency: Drastic reduction in inference latency and cost, making sequential information gathering feasible for real-world applications where latency is constrained.
Generalizability: The ability to transfer learned information-seeking strategies to unseen domains (e.g., from 20 Questions to medical diagnosis) without retraining or task-specific inference-time tuning.
Robustness: The method improves performance across different model scales and maintains general language capabilities, addressing the trade-off between specialized information gathering and general reasoning.

The authors position ASIG as a step toward making LLMs more effective agents in sequential decision-making environments, moving beyond "lost in conversation" phenomena by equipping models with learned strategies for efficient belief updating and information acquisition.