Learning What to Ask and What to Believe: Distilling a Variational Bayesian Agent into Language Models

Baichuan Li; Yanling Chang; Alfredo Garcia

cs.LGOct 31, 2026

Learning What to Ask and What to Believe: Distilling a Variational Bayesian Agent into Language Models

Authors: Baichuan Li, Yanling Chang, Alfredo Garcia

Organizations: Southern Methodist University · University of Virginia

Abstract

Interactive tasks such as eliciting a user's preferences or diagnosing a hidden mechanism require an agent to track its uncertainty over a latent state, ask informative questions and act on the answers. Language models have no explicit mechanism for this, whereas Bayesian information seeking provides one through posterior inference and expected information gain (EIG), and we ask whether it can be distilled into an 8B language model. Our teacher, a lightweight Bayesian agent with a sequential variational belief and EIG query selection, matches or approaches exact and particle references on three task families. We generate synthetic interactions with the teacher on episodes disjoint from the benchmark and fine-tune \Student{} to reproduce its queries and its belief state as a structured scratchpad. On three Multi-Turn Puzzles tasks the student raises Circuit Decoding from 0.00 to 0.93 success, lowers the normalized rank of the recommended film from 0.50 to 0.10 and cuts word guessing from 32 to 3.4 attempts, exceeding the strongest reported language models on the two comparable tasks. Exposing the belief state closes much of the gap where the decision is an implicit computation rather than a read-out of the interaction; in matched teacher contrasts, EIG supervision consistently outperforms random querying, whereas greater inference fidelity does not uniformly improve the distilled student; a 70B language-model teacher does not match the Bayesian teacher in downstream student performance, and the learned skill does not transfer to an unseen task family. Variational Bayesian information seeking can thus be amortized into a language model by transferring the belief state that drives its questions and decisions.

Figures & tables

Panel of Figure 1 — Figure 1: Overview. A Bayesian teacher interacts with synthetic episodes disjoint from the benchmark; its interactions are serialized as actions only or as belief scratchpad plus action and distilled into Llama-3.1-8B-Instruct, which is then evaluated interactively on the complete benchmark test splits. The targets are shown for two turns of one movie-recommendation dialogue of the teacher; the scratchpad adds the <belief> and <scores> lines (blue; formats for all tasks in Appendix D.4 ).

	Movie Recommendation	Circuit Decoding	Word Guess
Latent state $\theta$	weights $\in[0,0.9]^{8}$ (continuous)	$24$ -bit truth table (structured binary)	one of $40$ words (finite symbolic)
Prior $p(\theta)$	uniform on the grid $\{0,0.1,\dots,0.9\}^{8}$	uniform over the tables realizable under the inventory	uniform over the $40$ words
Query $a$	pair of seen films ( $190$ )	(circuit, input) ( $24$ )	vocabulary word ( $40$ )
Observation $y$	Yes / No / No Preference	one bit	$5$ colors ( $\leq 243$ patterns)
Budget $T$	$10$	$18$	$40$ (minimize attempts)
Decision $d$	one of $40$ unseen films	all $24$ bits	the secret word

Table 1: The three tasks and the teacher’s variational family for each (

1{,}000

300

400

test episodes).

Task	Metric	Base (8B)	Student	Teacher	Student	Teacher	Reported LM
			ELBO+EIG		Exact / Particle Bayes+EIG
Circuit	Success@18 $\uparrow$	$0.000$	$\mathbf{0.927}$	$1.000$	$\mathbf{0.840}$	$1.000$	–
Circuit	benchmark score (%) $\uparrow$	$0.7$	$\mathbf{97.4}$	$100.0$	$\mathbf{92.8}$	$100.0$	$39.8$
Movie	Success@10 $\uparrow$	$0.032$	$\mathbf{0.345}$	$0.410$	$\mathbf{0.409}$	$0.562$	–
Movie	benchmark score (%) $\uparrow$	$50.4$	$\mathbf{90.0}$	$92.9$	$\mathbf{93.1}$	$96.5$	$84.3$
Word	mean attempts $\downarrow$	$31.97$	$\mathbf{3.37}$	$2.29$	$\mathbf{3.10}$	$2.29$	–

Table 2: Main result on the complete test splits: base model, the belief-scratchpad students of the ELBO+EIG teacher and of the exact or particle reference, their teachers, and the strongest language model reported by the benchmark (Gemini 2.5 Pro; none for Word Guess: the normalization is unpublished and its two turn counts disagree). Benchmark score: circuit-wise accuracy (circuit) and

100\times(1-\text{normalized rank})

(movie). All student–base McNemar or Wilcoxon

p<10^{-60}

Panel of Figure 2 — Figure 2: Action-only against belief-scratchpad distillation. (a) Movie Recommendation: success of the action-only students under their own and their teacher’s questions, with their own and a Bayes-optimal read-out, against the belief-scratchpad students and the teachers. (b) Word Guess: success within $k$ attempts for the base model and the two ELBO+EIG students. Labels in (a) give questions / read-out.

	Circuit, Success@18 $\uparrow$		Movie, Success@10 $\uparrow$		Word, attempts $\downarrow$
Teacher	A / B	B $-$ A	A / B	B $-$ A	A / B	B $-$ A
ELBO+EIG	$0.920$ / $0.927$	$+0.007$ $[-0.037,+0.050]$	$0.085$ / $0.345$	$+0.260$ $[+0.228,+0.292]$	$19.2$ / $3.4$	$-15.8$ $[-17.5,-14.0]$
ELBO+Random	$0.500$ / $0.590$	$+0.090$ $[+0.010,+0.170]$	$0.105$ / $0.176$	$+0.071$ $[+0.044,+0.098]$	$15.2$ / $3.1$	$-12.1$ $[-13.4,-10.8]$
ELBO+Uncertainty	–	–	$0.064$ / $0.194$	$+0.130$ $[+0.103,+0.157]$	–	–
Exact / Particle+EIG	$0.643$ / $0.840$	$+0.197$ $[+0.133,+0.257]$	$0.079$ / $0.409$	$+0.330$ $[+0.296,+0.362]$	$15.4$ / $3.1$	$-12.3$ $[-13.8,-10.8]$

Table 3: Action-only (A) against belief-scratchpad (B) students of the same teacher on identical test episodes: score under each format and paired difference B

-

A with

95\%

bootstrap CI.

Contrast	Circuit	Movie	Word
Query rule: random $\to$ EIG	$+0.420$ $[+0.350,+0.487]$	$+0.169$ $[+0.137,+0.202]$	$+0.205$ ( $103/21$ )
Query rule: uncertainty $\to$ EIG	identical rules	$+0.151$ $[+0.117,+0.186]$	–
Inference: variational $\to$ exact or particle	$-0.277$ $[-0.333,-0.217]$	$+0.064$ $[+0.029,+0.098]$	$0.000$ ( $23/23$ )
Teacher type: Bayesian $\to$ LLM teacher	$-0.920$ $[-0.950,-0.887]$	$-0.033$ $[-0.055,-0.011]$	–

Table 4: Teacher-component contrasts: paired differences between students on identical test episodes (second condition minus first) with

95\%

bootstrap CIs; word: Success@3 with discordant counts. Circuit uses action-only students, movie and word belief-scratchpad students (the LLM teacher row is action-only; belief-scratchpad circuit contrasts in Appendix F.1 ).

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Condition	Circuit Decoding	Movie Recommendation	Word Guess
Base model	yes	yes	yes
ELBO+EIG student	A, B	A, B	A, B
ELBO+Random student	A, B	A, B	A, B
ELBO+Uncertainty student	(= EIG)	A, B	–
Exact / Particle Bayes+EIG student	A, B	A, B	A, B
LLM teacher student	A	A	–

Appendix

Table 5: Conditions evaluated on the complete test splits. A = action-only targets, B = belief-scratchpad targets. Uncertainty sampling coincides with EIG on Circuit Decoding; the LLM teacher has no belief and therefore no scratchpad condition.

Task	Metric	ELBO+EIG	ELBO+Rand.	ELBO+Unc.	Exact / Part.
Circuit	Success@18 $\uparrow$	$\mathbf{1.00}$	$0.70$	–	$\mathbf{1.00}$
	turns to identify $\downarrow$	$11.0$	$17.1$	–	$11.0$
	EIG-opt. / regret (bits)	$0.98$ / $0.001$	$0.30$ / $0.33$	–	$1.00$ / $0$
Movie	Success@10 $\uparrow$	$0.410$	$0.373$	$0.428$	$\mathbf{0.562}$
Movie	benchmark score (%) $\uparrow$	$92.9$	$91.0$	$93.8$	$\mathbf{96.5}$
Word	mean attempts $\downarrow$	$\mathbf{2.29}$	$3.09$	–	$\mathbf{2.29}$

Appendix

Table 6: Teacher validation on the complete released test splits (

300

1{,}000

400

episodes; five paired seeds on Circuit Decoding and for ELBO+Random on Word Guess, one run on Movie Recommendation). EIG-optimal rate and regret are measured under the exact posterior. ELBO+Uncertainty asks the pair whose predicted answer is closest to even odds (Eq. 3 , Appendix C.5 ) and is run only on Movie Recommendation, where it differs from EIG.

Figure 3: Inference dynamics of the teachers. (a) Circuit Decoding: mean entropy of the belief after each query on the synthetic training episodes (dotted line: the

18

-query budget). (b) One variational update on a training episode: distance of the ELBO from its analytic optimum over the

80

Adam steps, for the first three updates. (c) Movie Recommendation: cosine between the posterior mean and the true weights after each question (synthetic training episodes). (d) Word Guess: mean belief entropy before each attempt on the complete test split (episodes still running; the two EIG teachers coincide).

Figure 4: Validation loss per epoch (synthetic validation split, assistant tokens) of every distilled student, one panel per task. Solid lines with filled markers: action-only targets; dashed lines with open markers: belief-scratchpad targets. Validation loss measures how predictable a teacher’s interactions are, not how good they are: the LLM teacher students have the lowest loss and the worst benchmark scores, the action-only exact-inference circuit student has twice the loss of the corresponding variational student, and the word action-only students’ losses are more than an order of magnitude above the scratchpad students’, whose targets contain the deterministic bookkeeping. Losses are not comparable across target formats within a task (different supervised tokens). The checkpoint with the lowest validation loss is the one evaluated.

Panel of Figure 5 — Figure 5: Students inherit their teacher’s information-seeking policy on Circuit Decoding and only partly on Movie Recommendation and Word Guess: the teacher’s EIG-optimal rate on the training episodes against its student’s on the complete test split, one panel per task, one labeled point per teacher (belief-scratchpad students; the LLM teacher’s students are action-only).

Language Model Distillation

Language model distillation transfers a larger model’s capabilities to smaller or specialized models using outputs, probabilities, or internal representations.

See all 548 papers

Sep 30, 2026

Bayesian Fine-tuning Yields Language Models that are as Bayesian as their Beliefs Allow

Language models (LMs) are increasingly used for tasks that require reasoning about hidden variables from a few observations, for which Bayesian inference is the normatively correct solution. While supervised fine-tuning of an LM on the outputs of an optimal

\textit{Bayesian}

model leads to near-Bayesian behavior, standard supervised fine-tuning (SFT) on the true answers for the task falls short of it. But behavior alone does not tell us

\textit{why}

tuning on a

\textit{Bayesian}

or an

\textit{oracle}

(true answers) signal differs: whether the resulting LM represents Bayesian beliefs, acts on them, or turns them into a choice the way Bayes' rule does. To compare them, we formulate increasingly demanding requirements for an LM to count as a Bayesian decision maker, spanning its behavior, representations, and computations, and test them on a flight recommendation task. The Bayes-trained LM acts Bayesian, encodes quantities of Bayes' rule in its middle layers, and uses the encoded belief for the recommendation to a certain extent. The oracle-trained LM differs from it both in the beliefs it holds and whether it reads beliefs out into recommendations. Exchanging beliefs between the LMs transfers a part of the Bayesian advantage. Bayes fine-tuning thus installs usable Bayesian beliefs in an LM for reasoning under uncertainty in a way standard SFT on oracle answers cannot, highlighting the advantage of nuanced supervision.

Polina Tsvilodub, Andreas Waldis, Linlu Qiu +2

University of Tübingen · Massachusetts Institute of Technology · New York University

Fine-Tuning LLM Fine-Tuning LLM Uncertainty Estimation Belief Updating in LLMs Fine-Tuning Language Models

Ask PDF

Jul 20, 2026

Uncovering Latent Reasoning Strategies in Language Models

A language model

p_θ(y \mid x)

trained on reasoning tasks learns to solve problems via multiple distinct strategies, yet these strategies are implicit and entangled within the model's response distribution. We study the problem of decomposing the response distribution of a given pretrained language model into a structured, strategy-conditioned representation. Specifically, we learn a latent-variable factorization

p_θ(y \mid x) \leadsto (r_φ(z \mid x), g_φ(y \mid x,z))

, where a router

r

maps each input to a distribution over latent strategies

z

and a generator

g

produces the response conditioned on that strategy. A key challenge is that the generator, initialized from the base model, already represents

p_θ(y \mid x)

without using

z

. Standard variational inference therefore gives the model no incentive to route information through

z

and can yield a severe form of posterior collapse. To address this, we propose a variational objective that measures fractional information gain relative to the base model's response loss and concentrates reconstruction pressure on tokens with high base model surprisal, encouraging

z

to encode strategy-relevant response variation. We introduce a benchmark of multi-strategy algorithmic tasks and show that this objective recovers latent codes aligned with distinct reference strategies while preserving the base model's response distribution.

Awni Altabaa, John Lafferty

Yale University

Posterior Collapse Implicit Reasoning in Language Models Variational Inference Structured Reasoning in Language Models

Ask PDF

Jul 3, 2026

Amortising Bayesian Experimental Design for Sequential Information Gathering in LLMs

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.

Jakob Hartmann, James Harvey, Jhonathan Navott +5

University of Oxford, Oxford, United Kingdom · 2Ellison Institute of Technology, Oxford, United Kingdom

RL for Language Models Amortized Inference Active Information Gathering Bayesian Experimental Design

Ask PDF

Abstract

Language Model Distillation

Explore similar work

Bayesian Fine-tuning Yields Language Models that are as Bayesian as their Beliefs Allow

Uncovering Latent Reasoning Strategies in Language Models

Amortising Bayesian Experimental Design for Sequential Information Gathering in LLMs