cs.CLSep 30, 2026
SaveBayesian Fine-tuning Yields Language Models that are as Bayesian as their Beliefs Allow
Organizations: University of Tübingen · Massachusetts Institute of Technology · New York University
Abstract
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 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 tuning on a or an (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.
Figures & tables
Figure 1: Overview of our framework for analyzing to which extent LMs instantiate Bayesian inference . We test the framework on a flight recommendation task (left), which is optimally solved through Bayesian decision making (middle). We operationalize the consistency with Bayes-optimal decision making in LMs through four increasingly fine-grained requirements capturing different layers of Bayesian inference (right), showing that LMs are less consistent with increasing requirements.
Figure 4: The quantities of Bayes’ rule emerge in the middle layers. Layer-wise probing for the prior, the likelihood update, and the posterior, each scored against the BayesAssist . Dashed lines mark the reference KL between a uniform distribution and the assistant; the grey area marks the layers – with the main effect; the bands show the deviation across the probes per layer.
Figure 5: Path of the average decoded choice policy ( ) of BayesLM and OracleLM and the optimal one under the BayesAssist ( top corner).
| R1 | R2 | R3 | R4 | |
|---|---|---|---|---|
| StartLM | ||||
| OracleLM | ✓ | ✓ | ||
| BayesLM | ✓ | ✓ | (✓) | (✓) |
Table 2: Overview of results. Fine-tuning with Bayesian signal, but not oracle one, leads to stronger compliance with the requirements.
Appendix figures & tables25 assets
Supplementary material from the paper’s appendix.
Appendix
| term | 95% CI of the OR | OR | ||
|---|---|---|---|---|
| BayesLM | 2.4147 | |||
| StartLM | 0.2403 | |||
| Shuffled | 0.9705 | |||
| Low-Noise | 0.8105 | |||
| High-Noise | 0.6561 | |||
| BayesLM:Shuffled | 0.9938 |
Table 3: Logistic regression analysis of Bayes accuracy in § 3 , regressing Bayes correctness against the LM, the test set and their interaction: correct C(model) * C(set) . The predictors are dummy coded. OR refers to odds ratio.
| term | 95% CI of the | OR per SD | ||
|---|---|---|---|---|
| BayesLM | 3.2007 | |||
| StartLM | 0.2286 | |||
| H(policy) | 0.6096 | |||
| H(policy):BayesLM | 0.5359 | |||
| H(policy):StartLM | 1.5357 |
Table 4: Logistic regression analysis of policy divergence in § 3 , regressing Bayes correctness against the LM, the entropy of the Bayesian policy and their interaction: correct C(model) * H(Bayesian policy) . The model is dummy coded. OR refers to odds ratio.
| term | 95% CI of the | ||
| div(behavioral policy) C(model) * div(prior) ( , AIC=473.7) | |||
| Intercept | — | ||
| BayesLM | |||
| div(prior) | |||
| div(prior):BayesLM | |||
| div(behavioral policy) C(model) * div(probed policy) ( , AIC=-276.4) | |||
Table 5: Linear regression analysis of Bayes-consistency of decoded beliefs in § 4 , regressing the LM’s token-based behavioral policy against the LM, one of the two predictors (probed policy divergence or prior belief divergence) and their interaction. The model is dummy coded.
Figure 7: Prompt presenting the first round of the recommendation task, marking the probed token positions (reported probing was performed on round 5).
Figure 8: Visualized search over the softmax rationality parameter for scaling the BayesAssist ’s policy: is a uniform distribution, corresponds to the original BayesAssist , is near-deterministic, and is the best match to the LMs’ choice policies.
Figure 9: The BayesLM ’s uncertainty tracks the assistant’s. Entropy of the LM’s behavioral choice policy against the entropy of the Bayesian policy, one point per sample of the original set, colored by whether the recommendation agrees with the BayesAssist . The BayesLM ’s entropy correlates with the assistant’s ( ), the OracleLM ’s only weakly ( ). For both models, the differing instances are concentrated where the BayesAssist itself is most uncertain, but spread more for the OracleLM .
Figure 10: Mean decoded choice policy by layer on the simplex (corners ordered by Bayesian policy), split by whether the LM’s recommendation agrees with the BayesAssist , i.e., is Bayes-correct, (green) or differs (red). Triangles mark the mean decoded choice policy for each layer, and circles mark the mean Bayesian policy of the BayesAssist for the same instances.
Figure 11: Per-layer probing across the four evaluation sets (detail for Figure 4 ). Each row is an evaluation set (original, shuffled, low-noise, high-noise); columns first present probed quantities of Bayes’ rule, and then the probed policy and resulting prediction (argmax of probed policy). On the original and shuffled sets, the belief and choice KL curves fall in the mid-stack and the BayesLM outperforms OracleLM and StartLM in the late layers. On the noisy sets the prior/posterior KL instead rises in the upper layers: the probes are trained on the original set and learn to represent sharper beliefs in late layers, which diverges from the flatter beliefs induced by the noisy sets. Therefore, the KL grows even as choice top-1 stays high — likely a calibration gap of the probe, not a representational failure of the model.
Figure 12: Per-layer probing across token positions (detail for Figure 4 ). Each row is a read-out position (feedback token, last prompt token, digit-generation token); columns first present probed quantities of Bayes’ rule, and then the probed policy and resulting prediction (argmax of probed policy). The choice representation (flight-choice KL, top-1) develops in the mid-to-late layers at the last-prompt and digit tokens — where the BayesLM and OracleLM models rise well above StartLM — but stays near chance at the feedback token, showing the choice is not yet committed there. Belief targets (prior/posterior KL) are decodable across all positions and depths.
Figure 13: Prior probing per feature. KL of the prior probe per layer at the last-input token, on the original set. Dashed: a uniform probe for that feature; grey band: layers 15–25.
| setting | flip accuracy | stay accuracy | balanced accuracy | noise effect |
|---|---|---|---|---|
| Varied dose , at 32 tokens and 8 layers | ||||
| 0.439 | 0.953 | 0.696 | 0.037 | |
| 0.817 | 0.847 | 0.832 | 0.037 | |
| 0.939 | 0.800 | 0.870 | 0.037 | |
| 0.963 | 0.765 | 0.864 | 0.037 | |
| Varied dose , at 4 tokens and 8 layers | ||||
Table 6: Hyperparameter search for prototype patching. We select the best hyperparameter for patching using an iterative search, varying one parameter at a time while keeping the rest fixed. We used a separate set of instances where the recommendation should change and where it should remain the same, generated by the BayesAssist . We report for every setting, accuracy on instances that should flip and those that should stay, as well as the effect of injecting a norm-matched random vector (noise effect), We use tokens, and layers. Based on these results, we choose to patch with on tokens across layers, which offers the best trade-off.
Figure 14: Intervening on the encoded belief changes the recommendation in BayesLM . Balanced accuracy results of prototype belief patching.
Figure 15: Prototype patching versus CAA steering. Balanced accuracy of the BayesLM by layer window on the original test set, averaged over the eight feature-direction settings. Prototype patching uses tokens and ; additive steering uses at the tokens after the feedback token or at the digit token. The dotted line marks a model that never changes recommendations; the dashed lines mark the layers where § 4 decodes the belief.
| neuro-symbolic | over LM | ||
|---|---|---|---|
| Llama-3.1-8B | BayesLM | 0.86 | |
| OracleLM | 0.82 | ||
| Qwen-2.5-7B | BayesLM | 0.80 | |
| OracleLM | 0.74 |
Table 7: Neuro-symbolic predictions. We read each model’s prior beliefs from its hidden states with linear probes, turn the four decoded distributions into per-attribute rewards, and pick the flight the BayesAssist ’s own rule would pick from them. neuro-symbolic is how often that flight matches the BayesAssist ’s recommendation, at the layer where each model scores best. over LM is the gain over the same model answering on its own.
| BayesLM | |||||||||
|---|---|---|---|---|---|---|---|---|---|
| feature | direction | original | shuffled | low noise | high noise | OracleLM | BayesLM OracleLM | BayesLM StartLM | steering |
| costs | negative | 0.775 | 0.799 | 0.817 | 0.818 | 0.566 | 0.582 | – | 0.549 |
| costs | positive | 0.769 | 0.765 | 0.766 | 0.766 | 0.530 | 0.612 | – | 0.589 |
| departure time | negative | 0.760 | 0.763 | 0.751 | 0.782 | 0.558 | 0.662 | – | 0.565 |
| departure time | positive | 0.756 | 0.773 | 0.787 | 0.773 | 0.533 | 0.590 | – | 0.628 |
| duration | negative | 0.766 | 0.770 | 0.785 | 0.786 | 0.505 | 0.606 | 0.525 | 0.613 |
Table 8: Balanced accuracy per feature and direction. Prototype patching at layers – with tokens and . The first four columns patch the BayesLM with its own prior on each test set. The next three use the original test set and patch the OracleLM with its own prior, and the BayesLM ’s prior into the OracleLM and the StartLM . The last column is CAA steering at the tokens after the feedback token ( ) on the original test set. The StartLM recipient was only run for duration and stops.
Figure 16: The choice policy on the simplex, with corners labeled by the BayesAssist ’s ranking of the three options under the updated , split by whether the BayesAssist ’s recommendation changes after the intervention. Open circles are the BayesLM ’s mean policy before the intervention, filled circles after the intervention. The diamond marks the BayesAssist ’s own policy before (open) and after (filled).
Figure 17: Overview of the attention masks applied for investigating how posterior beliefs might be formed.
Figure 18: Overview of the attention masking accuracies (y-axis) of the BayesLM and OracleLM , across conversational rounds (x-axis).
Figure 19: Distribution of the choice policy of the BayesLM under masking.
Figure 20: Results of neuro-symbolic predictions across layers. The development of the average Bayes accuracy of the three neuro-symbolic models when the probes stem from different model layers.
Figure 21: Bayes accuracy across backbones (R1, § 3 ). Bayes accuracy of the BayesLM and OracleLM on the four test sets, with 95% confidence intervals. Chance is (dashed line). The BayesLM is more Bayes-accurate than the OracleLM on every backbone, and is invariant to evidence order, and behaves as expected under noise.
Figure 22: Oracle accuracy across backbones (R1, § 3 ). Oracle accuracy shows the proportion of recommendations matching the user’s true preference. All models lose accuracy in the presence of noise, as the BayesAssist does, because the evidence is less informative of the user’s preference. The BayesLM exceeds the OracleLM on the original and shuffled sets for every backbone.
Figure 23: Decoded prior across backbones (R2, § 4 ). Belief divergence between the decoded and the BayesAssist ’s prior by layer, on the original set. The dashed line is the divergence of a uniform prior. On every backbone, the prior becomes decodable in the middle layers and more strongly for the BayesLM than for the OracleLM .
Figure 24: Prior patching across backbones (R3, § 5 ). Balanced accuracy of the BayesLM upon intervening on the prior, by the layer window of the edit. The dashed line is a model that never changes its recommendation. The marker shows the best window. The effect is largest for Gemma-2-9B, present for Llama-3.1-8B, and near chance for Qwen-2.5-7B.
Figure 25: Attention masking across backbones (R3, § 5 ). Bayes accuracy without masking and under the one-swoop and sequential masks. Chance is (dashed line). On every backbone, both masks reduce accuracy, and the sequential mask reduces it more, though the gap between the two masks is smaller for Llama-3.1-8B and Qwen-2.5-7B than for Gemma-2-9B.