Distilling a Bayesian Agent into an 8B Model: Major Gains in Proactive Question-Asking

Language models often lack explicit uncertainty-tracking mechanisms in interactive tasks, making it difficult for them to efficiently gather information through proactive questioning the way humans do. A new study attempts to distill the information-seeking capabilities of a variational Bayesian agent into an 8-billion-parameter language model, achieving significant breakthroughs on several multi-turn puzzle tasks.

Problem Background: Language Models Lack an Explicit Belief State

Interactive tasks—such as eliciting user preferences or diagnosing a hidden mechanism—require an agent to maintain uncertainty over latent states, ask informative questions, and act on the answers. Bayesian information seeking provides a mathematical framework for this process through posterior inference and expected information gain (EIG), but language models themselves have no such explicit mechanism.

The core question the research team posed is: can Bayesian information-seeking capabilities be distilled into an 8B-scale language model? They designed a lightweight Bayesian teacher agent that uses a sequential variational belief and an EIG-based query selection strategy, matching or approaching exact inference and particle-filter reference methods on three task families.

Expected Information Gain (EIG) is a core metric in Bayesian experimental design, used to quantify how much a candidate question or observation can reduce uncertainty about a hidden state in expectation. For a candidate query (q), EIG is defined as the expected entropy reduction in the belief distribution before and after observing an answer: (\text{EIG}(q) = H(\text{prior}) - \mathbb{E}_{a \sim p(a|q)}[H(\text{posterior}|a)]). Intuitively, the queries with the highest EIG are those most likely to redistribute probability mass and thus most effectively narrow the candidate set. In interactive tasks, an agent can achieve near-optimal information seeking by computing EIG over all feasible queries at each turn and selecting the maximum. This metric is in the same lineage as the optimal bisection strategy in the game of "20 Questions," and it is also the basis on which the teacher agent in this paper selects queries.

Distillation Method: Belief State as a Structured Scratchpad

The key design of the study is that the student model imitates not only the teacher's queries but also reproduces its belief state—output in the form of a structured scratchpad. Training data comes from dialogues generated by the teacher on synthetic interaction episodes that do not overlap with the benchmarks, and the student model learns through fine-tuning to predict the teacher's queries and belief states.

This design exposes the belief state explicitly, making what would otherwise be implicit decision computations observable. The paper points out that when decisions depend on implicit computation rather than directly reading the interaction content, exposing the belief state can close most of the gap.

Variational Bayesian methods approximate an intractable posterior (p(z|\text{history})) with a parameterized approximate distribution (q_\phi(z)), trained by optimizing the evidence lower bound (ELBO). In sequential interaction settings, the belief state is recursively updated with each turn's observation, and variational methods allow the agent to perform inference efficiently in an amortized manner, without having to rerun expensive particle filtering for every new dialogue. The teacher agent in this paper adopts exactly such a sequential variational belief: a neural network encodes the interaction history and outputs an approximate posterior over latent states, on which the EIG for each candidate query is then computed. This design enables the teacher to generate belief-annotated dialogues at scale on synthetic data, providing supervision signals for subsequent distillation.

Experimental Results: Performance Leaps on Multiple Tasks

On the three Multi-Turn Puzzles tasks, the distilled student model performed remarkably:

• Circuit Decoding: success rate rose from 0.00 to 0.93
• Movie Recommendation: normalized rank of recommended movies dropped from 0.50 to 0.10
• Word Guessing: number of attempts cut from 32 to 3.4

On the two comparable tasks, the student model surpassed the strongest language models previously reported.

Multi-Turn Puzzles is a set of benchmark tasks for evaluating interactive information-seeking ability, requiring an agent to identify a hidden target through proactive questioning over multiple turns of Q&A. Circuit Decoding requires the agent to infer an unknown logic circuit by probing input-output pairs; in the movie recommendation task, the agent must narrow down the user's preferred movies through questioning; and the word guessing task resembles the game of "20 Questions," in which the agent must guess a hidden word within a limited number of turns. What these tasks share is a large answer space, a cost to each query, and an optimal strategy that depends on precise maintenance of the posterior. Success rate, normalized rank, and number of attempts measure, respectively, final correctness, recommendation quality, and information-seeking efficiency.

Ablations and Comparisons: EIG Supervision Outperforms Random Queries

In matched teacher comparison experiments, EIG supervision consistently outperformed random query strategies, indicating that information-gain-guided questioning indeed yields transferable benefits. However, higher inference fidelity did not uniformly improve the performance of the distilled student model, suggesting that structured output of the belief state may be more critical than exact posterior matching.

One notable finding is that a 70B language model teacher did not match the Bayesian teacher in downstream student performance. This suggests that scaling model size alone cannot replace the principled advantages of Bayesian information seeking. In addition, the skills learned by the student model did not transfer to unseen task families, indicating that the generalization ability of current distillation methods remains limited.

Knowledge distillation typically refers to having a student model imitate a teacher model's output distribution or intermediate representations. The distillation in this paper is sequence-level behavior cloning: the training data are complete dialogue trajectories generated by the teacher in synthetic interactions, and the student learns not only the teacher's final answers but also the belief-state scratchpad it outputs at each turn. Unlike traditional logit matching, this "structured scratchpad" supervision externalizes the teacher's implicit computation process as a learnable text sequence, thereby compressing the algorithmic logic of Bayesian inference into the language model's parameters. This explains why higher inference fidelity does not uniformly improve student performance—what matters is whether the scratchpad format provides a generalizable computational scaffold, rather than exact matching of posterior values.

Significance and Limitations

This work demonstrates the feasibility of amortizing variational Bayesian information seeking into a language model—achieving capability transfer by passing on the belief states that drive questioning and decision-making. For AI applications that require multi-turn interaction, information gathering, and uncertainty management, this approach offers a technical path that does not rely on larger model scale.

The limitations are equally clear: cross-task generalization has not yet been achieved, and the student model cannot reuse its learned skills on entirely new task types. Future improvements may include richer training task distributions or the introduction of meta-learning mechanisms to enhance transferability.

One-line Summary

Distilling variational Bayesian information seeking into an 8B language model yields major performance leaps on multi-turn puzzle tasks.

Paragraph Summary

Addressing the lack of explicit uncertainty tracking in language models, this study proposes distilling the information-seeking capabilities of a variational Bayesian agent into an 8-billion-parameter model. The teacher agent uses a sequential variational belief and expected information gain to select queries, and the student model not only imitates the queries but also reproduces the belief state as a structured scratchpad. On three multi-turn puzzle tasks, the student model's success rate rose from 0.00 to 0.93, the normalized rank in movie recommendation dropped from 0.50 to 0.10, and word-guessing attempts fell from 32 to 3.4. Ablations show that EIG supervision outperforms random queries, but higher inference fidelity does not uniformly improve performance; a 70B language model teacher yields worse downstream student performance than the Bayesian teacher, and skills do not transfer to unseen task families.