Learning Dynamic Belief Graphs for Theory-of-mind Reasoning
| Authors | Ruxiao Chen et al. |
| Year | 2026 |
| Field | AI / Agents |
| arXiv | 2603.20170 |
| Download | |
| Categories | cs.AI |
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/
Engineering Breakdown
Plain English
This paper addresses a critical gap in how large language models perform Theory of Mind (ToM) reasoning—understanding what people believe and how those beliefs drive their actions, especially in high-stakes domains like emergency medicine and disaster response. The authors propose a dynamic belief graph model that represents mental states as evolving networks of interconnected beliefs, rather than static, independent variables. Unlike prior approaches that either use raw LLM prompting or treat beliefs as fixed, this method jointly infers latent beliefs, learns how their dependencies change over time, and connects belief evolution directly to information-seeking behavior and decision-making. The core innovation is structuring ToM as a temporal reasoning problem where belief relationships are dynamic, enabling the model to maintain coherent mental models across multiple reasoning steps.
Core Technical Contribution
The technical novelty centers on representing Theory of Mind as a dynamic graph structure rather than treating beliefs as isolated latent variables. The model jointly optimizes three coupled objectives: inferring what beliefs an agent holds at each timestep, learning how the causal dependencies between beliefs shift over time, and linking observed information-seeking and decision behavior back to underlying belief states. This differs fundamentally from prior latent-state models that assume belief structure is static—here, the graph topology itself evolves, and the model must infer both the node values (what beliefs are held) and the edge structure (how beliefs depend on each other). The integration of belief dynamics with LLM-based reasoning creates a structured cognitive trajectory that produces temporally coherent mental models, addressing the persistent problem that standard prompting produces fragmented or contradictory reasoning across timesteps.
How It Works
The model takes as input a sequence of observations about an agent's behavior, information exposure, and decisions, then infers a latent belief graph at each timestep. Each node in the graph represents a belief dimension (e.g., risk assessment, resource availability, person's competence), and edges encode causal or dependency relationships between beliefs. The system uses the LLM as a reasoning engine to map observations to belief states, but constrains the output through a structured graph representation rather than allowing freeform generation. A key mechanism is time-varying edge inference—the model learns which belief relationships were active or relevant in the past and predicts how the graph structure changes as new information arrives. Finally, the model includes a linking function that connects belief graph states to observable decisions and information-seeking actions, creating a fully grounded cognitive trajectory that can be evaluated against real agent behavior. The training likely uses supervised or semi-supervised learning on sequences where you have both the behavior trajectory and some ground-truth or proxy signal for beliefs.
Production Impact
For teams building human-in-the-loop AI systems or autonomous agents that must model human operators, this approach directly improves decision quality and safety. Instead of treating human user intent as a black box, you can now maintain an explicit, interpretable model of their evolving mental state, which reduces miscommunication and enables the system to proactively provide information when beliefs are likely misaligned. In emergency response systems (the paper's motivating domain), this means the AI can detect when a human responder's risk assessment is diverging from ground truth and trigger targeted clarification prompts rather than silent misalignment. The production trade-off is that you must annotate training data with ground-truth belief states or use weak supervision from behavior logs, increasing data labeling cost and requiring domain expertise to define the belief graph schema. Latency increases modestly since the model must perform inference over a graph structure at each timestep, but this is typically acceptable for non-real-time advisory systems; real-time applications would need careful optimization or belief state caching.
Limitations and When Not to Use This
The paper assumes beliefs can be meaningfully decomposed into a graph structure with interpretable nodes—this works well for technical domains (medicine, robotics) but may fail for reasoning about abstract, entangled social beliefs where cause-effect structure is ill-defined or culturally variable. The approach requires either labeled training data with ground-truth beliefs or high-quality behavior logs where belief states can be reliably inferred from actions; obtaining this at scale is expensive and may introduce systematic bias toward observable, rational beliefs while missing intuitive or culturally-specific reasoning. The model also inherits LLM limitations: it can still hallucinate or confabulate belief states that sound plausible but are incorrect, and the graph structure itself (which belief dimensions matter, which relationships to model) must be manually designed or data-driven but not learned end-to-end. Performance likely degrades in truly novel or adversarial settings where an agent's beliefs shift in ways not seen in training data, and the paper does not address how to detect when the learned graph schema itself becomes invalid.
Research Context
This work builds on a long tradition of computational Theory of Mind, from rational agent models to recent neural approaches that use LLMs for belief inference, but adds temporal and structural rigor that prior work lacked. It advances beyond static latent variable models (like standard VAEs applied to agent modeling) by introducing time-varying causal structure, connecting it to recent work on causal inference and temporal graphical models in the ML literature. The paper likely evaluates on established ToM benchmarks (such as false-belief tasks or complex reasoning scenarios) and possibly on new real-world datasets in emergency response, though the abstract does not specify exact benchmarks. This opens a research direction toward LLM-grounded cognitive science: using structured belief graphs as a bridge between neural language models and interpretable cognitive models, potentially enabling better integration of AI with human teams and richer explanations of agent reasoning.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::