AR-BED: Amortizing Sequential Experimental Design in Reasoning LLMs

Large language models (LLMs) are not good at determining "which question to ask to clarify the situation as quickly as possible." The authors use information gain as a reward to directly train the model's reasoning process before selecting the next question; the generative model is set independently of the LLM's own beliefs, allowing it to optimize expected information gain according to its chosen model. The paper also systematically compares design choices such as generative model construction, reward estimation, hypothesis space, and multi-step objectives.

Large language models (LLMs) are not good at determining "which question to ask to clarify the situation as quickly as possible." The authors use information gain as a reward to directly train the model's reasoning process before selecting the next question; the generative model is set independently of the LLM's own beliefs, allowing it to optimize expected information gain according to its chosen model. The paper also systematically compares design choices such as generative model construction, reward estimation, hypothesis space, and multi-step objectives.

Likes 5 · Code Stars 2

Rationale

  • The authors claim improvements in both information gathering and task performance on the Twenty Questions game, outperforming inference-time and training-time Bayesian experimental design baselines with significantly lower computational costs, and demonstrating generalization to unseen categories.
  • Transfer from Twenty Questions training to the real-world clinical consultation benchmark Medi-Q shows positive effects, indicating that "active information gathering" can be learned as a general reasoning skill.
  • The code and training data for the animal / Medi-Q tasks are available in the repository under data/tq/, with one-line commands provided for training, evaluation, and BED-LLM baselines.

Usage

  • Code: github.com/anitakriz/ar-bed, pip install -r requirements.txt.
  • Training the animal version: python scripts/run.py --config-name k1_expected; the Medi-Q three-domain version: --config-name tq_mediq_3domain_k1_expected.
  • Evaluation: python -m scripts.evaluate_tq_checkpoints k1_expected --checkpoint-root … --batched-full-game; add --evaluate-base-model to test the base model; use scripts.evaluate_bed_llm_baseline for baselines.

Risks

  • The README is brief: it does not specify which base model is used, the required VRAM, or provide trained weights; reconstructing data or filling in missing likelihoods requires API credentials from utils/api_models.py.
  • Experiments are focused on game-like tasks such as Twenty Questions, with clinical transfer validated on only one benchmark; the README does not specify a license.
  • The results are reported by the authors, and the abstract does not include specific numbers.

Open · Code · Paper Page

Hugging Face Daily Papers

Explore the latest AI Radar updates