Simulated Annealing Enhances Theory-of-Mind Reasoning in Autoregressive Language Models

TL;DR

The paper addresses the difficulty of maintaining global latent-state reasoning in autoregressive language models for Theory-of-Mind tasks. It introduces a test-time optimization method based on simulated annealing of a power-sharpened sequence-level distribution, implemented via Metropolis–Hastings with an autoregressive LM as the proposal and an exponential temperature schedule $\tau_k = \tau_{\text{start}} (\tau_{\text{end}}/\tau_{\text{start}})^{k/K}$, where $\tau_k$ gradually cools from $0.90$ to $0.25$, and where the temperature $\tau$ relates to the power parameter by $\tau = 1/\alpha$. By evaluating on small models with the BigToM benchmark, the approach demonstrates that strong ToM performance can be recovered without training or verification, outperforming direct decoding, Chain-of-Thought, and fixed-temperature power sampling. This reveals substantial latent ToM capabilities in pretrained autoregressive models and proposes test-time optimization as a unifying framework for decoding strategies, albeit with higher inference cost and some limitations on deeply nested reasoning. The work has practical implications for eliciting complex reasoning from compact models and motivates extending sequence-level optimization to other domains beyond ToM.

Abstract

Autoregressive language models are next-token predictors and have been criticized for only optimizing surface plausibility (i.e., local coherence) rather than maintaining correct latent-state representations (i.e., global coherence). Because Theory of Mind (ToM) tasks crucially depend on reasoning about latent mental states of oneself and others, such models are therefore often thought to fail at ToM. While post-training methods can improve ToM performance, we show that strong ToM capability can be recovered directly from the base model without any additional weight updates or verifications. Our approach builds on recent power-sampling methods (Karan & Du, 2025) that use Markov chain Monte Carlo (MCMC) to sample from sharpened sequence-level (rather than token-level) probability distributions of autoregressive language models. We further find that incorporating annealing, where the tempered distribution is gradually shifted from high to low temperature, substantially improves ToM performance over fixed-temperature power sampling. Together, these results suggest that sampling-based optimization provides a powerful way to extract latent capabilities from language models without retraining.

Figures (3)

• Figure 1: Illustration of the simulated annealing approach for optimizing reasoning in autoregressive language models.(a) Simulated annealing controls the temperature (i.e., the inverse of the power parameter $\alpha$) across successive MCMC iterations. At each iteration $k$, a token position $t$ is selected uniformly at random from the current sequence. The autoregressive model then resamples the sequence suffix conditioned on $\textbf{x}_{<t}$ to form a proposal sequence. The current and proposed sequences are subsequently evaluated using a Metropolis–Hastings acceptance step. (b) The temperature schedule used in our experiments, along with the corresponding fixed-temperature control conditions.
• Figure 2: (a) An example ToM narrative from the BigToM benchmark gandhi2023understanding. The same narrative supports both forward inference (predicting an agent’s action from its beliefs) and backward inference (inferring an agent’s beliefs from its action). Within each inference type, scenarios are further divided into True Belief (the agent observes the critical event) and False Belief (the agent does not observe the event) conditions. (b) Representative examples of model-generated reasoning under different decoding strategies. These examples highlight distinct modes of reasoning corresponding to globally consistent versus locally plausible but inconsistent belief interpretations.
• Figure 3: Accuracy across small language models and decoding strategies. Panels (left to right) correspond to Phi-3.5-Mini-Instruct, LLaMA-3.2-3B-Instruct, and Qwen3-1.7B. Within each panel, bars (left to right) represent five decoding strategies: direct decoding, Chain-of-Thought (CoT), low-temperature power sampling, high-temperature power sampling, and simulated annealing. Results are shown separately for True Belief (TB) and False Belief (FB) conditions in backward inference. Blue bars indicate TB accuracy, while orange bars indicate FB accuracy.