Google researchers propose Bayesian teaching method for large-scale language models

Machine Learning


Google researchers have proposed a training method that teaches large language models to approximate Bayesian inference by learning from the predictions of optimal Bayesian systems. This approach focuses on improving how models update their beliefs as they receive new information during multi-step interactions.

This research investigates how language models update their beliefs as they interact with users over time. In many real-world applications, such as recommendation systems, models must incrementally infer a user’s preferences based on new information. Bayesian inference provides a mathematical framework for updating probabilities when new evidence becomes available. The researchers investigated whether language models behave in a manner consistent with Bayesian belief updating and considered training methods to improve their behavior.

To evaluate this, the team created a simulated flight recommendation task. In the experiment, the model interacted with a simulated user over five rounds. In each round, the assistant and user were presented with three flight options defined by departure time, duration, number of stops, and price. Each simulated user had hidden settings for these attributes. After each recommendation, the user indicated whether the assistant had selected the correct option and revealed the desired flight. The assistant was expected to use this feedback to improve future recommendations.

The researchers compared several language models to a Bayesian assistant that maintains a probability distribution over possible user preferences and updates them using Bayesian rules after each interaction. In experiments, the Bayesian assistant reached approximately 81% accuracy in selecting the correct option. Language model performance was often poor, with limited improvement after the first interaction. This suggests that the language model is not effectively updating its internal estimates of the user’s preferences.

The study then tested a training approach called Bayesian education. The model was trained to mimic the Bayesian assistant’s predictions during mock interactions, rather than learning only from the correct answers. In early rounds, the Bayesian assistant sometimes made incorrect recommendations due to uncertainty about the user’s preferences, but its decisions reflected probabilistic reasoning based on available evidence.

The image below shows Gemma and Qwen’s recommendation accuracy after fine-tuning the user interaction with the Bayesian Assistant or oracle.

The training data for supervised fine-tuning consists of simulated conversations between the user and the Bayesian assistant. For comparison, the researchers tested how the model learns from an assistant that always chooses the correct option because it fully knows the user’s preferences.

Both fine-tuning approaches improved model performance, but Bayesian learning produced better results. Models trained in this way made predictions that more closely matched those of the Bayesian assistant and showed strong improvement over multiple interaction rounds. The trained model also showed higher agreement with the Bayesian system in evaluating user choices.

Community reaction to the Google Research post was largely positive, with commenters highlighting improvements to probabilistic inference and multiturn adaptation in LLM.

Software developer Jan Kronberg comments:

People talk about inference benchmarks, but this is basically about belief updating. We know that most LLMs do not modify their internal assumptions much as new information arrives, so teaching @GoogleResearch to approximate Bayesian inference could be very important for long-term agents.

Some also question the use of supervised fine-tuning instead of reinforcement learning to approximate Bayesian inference.

Researcher Aidan Lee said:

Why did the authors use SFT rather than RL to train the model to approximate probabilistic reasoning? There is a lot of research related to RL and probabilistic reasoning, even in LLM. Maybe I’m missing something, but RL seems like the obvious choice.

The researchers describe the method as a type of model distillation in which a neural network learns to approximate the behavior of a symbolic system that implements Bayesian inference. This result suggests that language models can acquire probabilistic reasoning skills through post-training that demonstrates optimal decision-making strategies during sequential interactions.





Source link