Amortizing Sequential Experimental Design in Reasoning LLMs

AI summary

Overview

Research area: Machine learning — reinforcement learning for large language models, sequential Bayesian experimental design (BED), and adaptive information gathering in multi-turn agents.

Technical level: Advanced. The paper assumes familiarity with expected information gain, posterior predictive distributions, contrastive bounds on mutual information (sPCE), and policy-gradient fine-tuning (GRPO).

Scope (one sentence): The paper trains a reasoning LLM to internalize sequential Bayesian experimental design — rewarding the reasoning traces that precede each question with expected information gain computed under a generative model specified independently of the policy — and shows this improves information acquisition, beats inference- and training-time BED baselines at lower cost, and transfers to a clinical dialogue benchmark.

What This Paper Is About

Large language models bring broad pretraining knowledge and strong reasoning to multi-turn tasks, but prior work cited in the paper (Gupta et al., 2025) finds they are insensitive to experimental feedback — their questioning behavior leans on possibly inaccurate priors rather than adapting to new evidence. The authors ask whether the ability to gather information adaptively can be learned as a reasoning skill rather than reconstructed by expensive search at every interaction. Their answer is to amortize sequential Bayesian experimental design into the model's reasoning process: the LLM writes a deliberation trace and then a question, and that joint output is rewarded by how much expected information the question would yield under a generative model the researchers control.

Key Contributions

1. A formulation of amortized sequential BED for reasoning LLMs. The policy is factorized as a reasoning trace followed by a design, π_φ(z_t, x_t | h_{t−1}) = π_φ(z_t | h_{t−1}) · π_φ(x_t | h_{t−1}, z_t), and the joint output is trained with information-gain rewards. The experimental generative model is decoupled from the policy LLM, so information is valued under a model of the authors' choosing rather than the model's own (possibly poor) predictive distribution.

2. A demonstration that information-gathering behavior can be amortized and transfers. Policies trained on Twenty Questions improve information acquisition on held-out and unseen categories, and positive transfer is shown from Twenty Questions training to MediQ, a real-world clinical dialogue benchmark, including specialties not used in fine-tuning.

3. A systematic design study of amortized BED. The paper examines alternative generative-model constructions, expected versus realized information-gain rewards, hypothesis-space specification (misspecified and smaller contrastive sets), multi-step (non-myopic) objectives, and computational cost against both training-time and inference-time BED baselines (Table 1 summarizes these choices across prior methods and the ablations).

Main Findings

• Reasoning alone does not help, and can hurt. Adding inference-time reasoning to the Base model reduced final solve rate across all evaluation categories and increased average EIG only on Countries. The authors note a similar effect appears for some models in Table 1 of Choudhury et al. (2025).

• Training a non-reasoning policy transfers inconsistently. The direct policy (AR-BED₁^Direct, which samples a question without a reasoning trace) improved in-distribution behavior but was worse than Base on Plants while improving on Celebrities and Countries — useful in-distribution question sampling without robust cross-domain transfer.

• Amortizing BED into reasoning gives the best overall information acquisition. AR-BED₁ achieved the best final solve rate across all categories and the highest average EIG per turn on Animals, Plants, and Countries, outperforming both the Base model and the directly trained policy.

• It beats BED-style baselines on most domains at much lower cost. AR-BED₁ outperformed inference-time BED-LLM and the training-time ASIG baseline on most domains; BED-LLM improved only modestly over the reasoning Base model and was worse than the non-reasoning Base model on some domains. BED-LLM was about 61× slower per game at inference (2376 s vs 38.9 s), and ASIG was about 2.9× slower per training step (136.3 s vs 47.3 s).

• A misspecified hypothesis space produces a flawed information signal. A model trained with a 100-mammal contrastive set (prompted to guess in Animals) performed significantly worse than both the animals-trained model and Base on held-out non-mammals, and was also worse than the animals-trained model on mammal targets — the question "Is it a mammal?" receives EIG ≈ 0 under that prior. Training on the broader animal hypothesis space appeared beneficial even within the shared mammal subset.

• Expected information gain beats realized information gain as a reward. Marginalizing over hypotheses and answers improved both EIG and solve rate on Animals, Plants, and Countries (though not on Celebrities), consistent with the variance-reduction (Rao–Blackwellization) motivation.

• Larger contrastive sets help out-of-distribution transfer. Appendix B.1 reports that a 200-sample contrastive set improves OOD transfer relative to a 100-sample set.

• Multi-step rewards help where non-myopic reasoning matters. In a natural-language adaptation of Cube-NM (Schütz et al., 2026), the two-step policy learned to acquire initially uninformative context for downstream benefit and substantially outperformed the one-step policy in later disease identification; on the TwentyQuestions target-prediction task of Wang et al. (2025), multi-step policies improved attribute prediction, particularly on medium and hard questions. In open-ended Twenty Questions, immediate and downstream EIG rankings can differ, with some gains from K=2 and K=3 over K=1 at two- and three-question horizons.

• Transfer to clinical dialogue. On MediQ, after training only on Twenty Questions, AR-BED₁ raised accuracy from 90.5% to 94.7% on 95 held-out dermatology cases (same domain, different task, held-out labels) and from 84.8% to 90.2% on 92 cases from other specialties (different domain and task). It also increased the likelihood assigned to the correct answer (76.0% → 79.2% and 72.8% → 78.1%) and reduced Brier score (0.165 → 0.120 and 0.227 → 0.157).

Methodology in Plain English

The researchers treat information gathering as a game with a hidden answer. In their training environment, Twenty Questions, a target (e.g., "horse") is hidden, the agent asks an open-ended natural-language question, and the answer comes from a small fixed set: yes, no, or maybe.

A good question is one that is expected to shrink uncertainty about the target. Computing that quantity requires a model of the world: a prior over which targets are possible, and a likelihood saying how each target would answer each question. The key move is that this world model is not the model being trained. The prior is a fixed uniform distribution over a finite list of 200 animals, and the likelihood comes from a separate, larger LLM used as an oracle — the paper uses GPT-5 — which can be queried and cached during training.

Training proceeds in three steps. First, the policy model (Gemma-4-E4B-it) is given the history of questions and answers and produces a reasoning trace followed by its next question. Second, the information value of that question is scored: instead of waiting for one sampled answer, the method marginalizes over all hypotheses and all possible answers, giving the expected information gain directly (the Rao–Blackwellized version of the realized sPCE reward). Third, that score is summed over a short rollout of K steps with discount γ (set to 0.5 when K > 1) to form a return, and the policy is updated with GRPO.

The authors also run a stripped-down variant, AR-BED_K^Direct, which skips the reasoning trace and is rewarded only for the question, to isolate what the reasoning contributes. Training used a batch size of 128 with 8 rollouts per state for 250 steps. Evaluation uses two metrics: solve rate within the 20-question budget and average EIG per turn (to account for games of different lengths). Each model plays 50 games across three seeds.

For the transfer study, the authors built three synthetic medical Twenty Questions environments — dermatology, neurology, and gastroenterology — each with a 50-disease contrastive set drawn from MediQ's answer inventory, trained a separate medical policy on 1,200 starting states, and had each policy ask five questions of a fixed GPT-5 "Patient," with a separate GPT-5 decider predicting the answer from the resulting history.

Why This Matters

The work reframes adaptive information gathering as a trainable reasoning skill rather than a search procedure bolted on at deployment. If a policy can internalize expected-information-gain reasoning, the expensive posterior estimation and design optimization of BED disappear at inference time — the paper measures that difference as roughly 61× per game against BED-LLM. It also shows that the world model used to define "informative" need not be the policy's own beliefs, which matters because LLM self-estimates of likelihoods can be unreliable.

Real-world applications motivated by the paper's framing:

• Clinical diagnosis support. The paper opens with a clinician choosing which symptoms to ask about and which tests to order, and validates transfer on MediQ, where policies question a simulated patient to reach a diagnosis.
• Diagnostic test selection under cost constraints. BED is described as useful precisely when experiments carry cost or time limits, so the chosen action should be as informative as possible.
• Automated customer or technical troubleshooting. Sequential yes/no/maybe-style questioning to localize a fault is structurally the Twenty Questions task the policy is trained on.
• Task disambiguation in agentic systems. The related work cites adaptive questioning for clarifying underspecified user requests, which is the same information-acquisition loop.
• Educational or interviewing agents. Systems that must elicit information from a person through a limited number of open-ended questions.

Industry relevance: The headline for practitioners is cost and deployability — a fine-tuned policy that answers in a single forward pass replaces repeated belief updating and expected-information-gain estimation at serving time. The explicit separation between the reward's world model and the policy also gives teams a controllable knob: they can encode domain knowledge or constraints in the prior and likelihood without waiting for the policy to learn them.

Future Directions

• Longer horizons. The authors state they study relatively short rollout horizons and flag extending to longer horizons while keeping a reliable training signal as an important direction.
• Characterizing the contrastive set. Performance is sensitive to the coverage of the contrastive hypothesis set when it is misspecified relative to the true hypothesis space; how size, composition, and prior weighting affect learning and transfer is left open.
• Closing the gap between controlled and open-ended non-myopia. The paper distinguishes a controlled non-myopic mechanism (Cube-NM) from the smaller lookahead benefits observed in open-ended questioning, implying more work is needed on when multi-step rewards pay off in natural-language design spaces.
• Broadening task transfer. Only one downstream benchmark (MediQ) is used for task transfer, leaving open how far the learned behavior generalizes to other information-gathering tasks and domains.

Target Audience

Researchers and practitioners working on LLM agents, multi-turn dialogue, and reinforcement learning with verifiable rewards, especially those interested in adaptive questioning, task disambiguation, and information-theoretic objectives. It will also appeal to the Bayesian experimental design community curious about amortizing design policies into language models. Readers need comfort with expected information gain, posterior predictive distributions, and policy-gradient methods; those without that background can still follow the empirical story, since the central claim — that training the reasoning trace, not just the action, is what produces transferable information gathering — is stated and tested plainly.

Authors’ abstract

Despite their extensive pretraining knowledge and strong reasoning capabilities, large language models (LLMs) struggle to adaptively gather information in experimental design tasks. We study how to amortize sequential Bayesian experimental design (BED) into an LLM by explicitly training the reasoning that guides action selection using information-gain rewards. Our approach specifies the experimental generative model independently of the LLM's own beliefs, allowing us to optimize expected information gain under a model of our choosing. We investigate the design choices for amortization, including alternative generative-model constructions, reward estimation, hypothesis-space specification, multi-step objectives, and computational costs relative to training- and inference-time BED baselines. On Twenty Questions, the resulting policies improve information acquisition and task performance, outperform inference- and training-time BED baselines at substantially lower computational cost, and generalize to unseen categories. We further demonstrate positive transfer from Twenty Questions training to Medi-Q, a real-world clinical dialogue benchmark. These results suggest that information gathering can be learned as a generalizable reasoning skill in LLMs.