Learning Dynamic Belief Graphs for Theory-of-mind Reasoning

Abstract

Theory of Mind (ToM) reasoning with Large Language Models (LLMs) requires inferring how people's implicit, evolving beliefs shape what they seek and how they act under uncertainty -- especially in high-stakes settings such as disaster response, emergency medicine, and human-in-the-loop autonomy. Prior approaches either prompt LLMs directly or use latent-state models that treat beliefs as static and independent, often producing incoherent mental models over time and weak reasoning in dynamic contexts. We introduce a structured cognitive trajectory model for LLM-based ToM that represents mental state as a dynamic belief graph, jointly inferring latent beliefs, learning their time-varying dependencies, and linking belief evolution to information seeking and decisions. Our model contributes (i) a novel projection from textualized probabilistic statements to consistent probabilistic graphical model updates, (ii) an energy-based factor graph representation of belief interdependencies, and (iii) an ELBO-based objective that captures belief accumulation and delayed decisions. Across multiple real-world disaster evacuation datasets, our model significantly improves action prediction and recovers interpretable belief trajectories consistent with human reasoning, providing a principled module for augmenting LLMs with ToM in high-uncertainty environment. https://anonymous.4open.science/r/ICML_submission-6373/

Figures (8)

• Figure 1: Belief trajectories evolve with high-stakes observations, triggering actions upon threshold crossing.
• Figure 2: Overview of the Structured Cognitive Trajectory ToM framework. Here, $s_t$ denotes the latent environment state, $o_t$ the agent's observation, $a_t$ the observed action, $b_t$ the latent belief state, and $e_t$ the semantic embedding extracted by a LLM.
• Figure 3: Action prediction accuracy over training epochs for intermediate actions and final evacuation decisions.
• Figure 4: ELBO component dynamics during training. Evolution of the action likelihood term and the KL divergence between the inference posterior and the belief transition prior.
• Figure 5: (a) Spearman correlation between model-predicted beliefs and human ratings for individual beliefs. (b) Spearman correlation for pairwise belief structure learning.