CoSToM:Causal-oriented Steering for Intrinsic Theory-of-Mind Alignment in Large Language Models

Abstract: Theory of Mind (ToM), the ability to attribute mental states to others, is a hallmark of social intelligence. While LLMs demonstrate promising performance on standard ToM benchmarks, we observe that they often fail to generalize to complex task-specific scenarios, relying heavily on prompt scaffolding to mimic reasoning. The critical misalignment between the internal knowledge and external behavior raises a fundamental question: Do LLMs truly possess intrinsic cognition, and can they externalize this internal knowledge into stable, high-quality behaviors? To answer this, we introduce CoSToM (Causal-oriented Steering for ToM alignment), a framework that transitions from mechanistic interpretation to active intervention. First, we employ causal tracing to map the internal distribution of ToM features, empirically uncovering the internal layers' characteristics in encoding fundamental ToM semantics. Building on this insight, we implement a lightweight alignment framework via targeted activation steering within these ToM-critical layers. Experiments demonstrate that CoSToM significantly enhances human-like social reasoning capabilities and downstream dialogue quality.

• The paper presents a causal tracing method to diagnose and align ToM representations in LLMs using targeted activation steering.
• It demonstrates early-layer primacy in encoding ToM, significantly amplifying desire signals and stabilizing deep-layer representations.
• It shows that improved ToM alignment leads to robust downstream dialogue performance and coherent, human-like social reasoning.

Causal-Oriented Steering for Theory-of-Mind Alignment in LLMs: An Authoritative Analysis of CoSToM

Introduction

The advancement in LLMs has spurred research into their emergent social reasoning capacities, notably Theory of Mind (ToM)—the human-like ability to attribute beliefs, desires, and intentions to others in communication. While numerous benchmarks indicate promising ToM reasoning by LLMs in controlled, prompt-based settings, a pronounced gap remains in the translation of these internal inferences to robust, contextually grounded behavior in unconstrained, task-driven dialogues. "CoSToM: Causal-oriented Steering for Intrinsic Theory-of-Mind Alignment in LLMs" (2604.10031) directly addresses this mechanistic-to-behavioral misalignment by moving beyond black-box prompt engineering, instead introducing a principled, interventionist framework that inspects, attributes, and modulates ToM representations within LLMs to attain both interpretability and stable, human-like social reasoning.

Intrinsic ToM in LLMs: Diagnosis by Causal Tracing

CoSToM begins with a methodical diagnosis of ToM encoding within LLMs via causal tracing, operationalized by extracting and probing hidden activations at each layer to quantify the recoverable BDI (belief, desire, intention) state information. Using the NegotiationToM and PersuasiveToM benchmarks, causal tracing pinpoints where ToM-relevant information emerges and decays.

Results unambiguously demonstrate early-layer primacy in ToM encoding. Contrary to standard assumptions in LLM literature suggesting higher-order reasoning is isolated to deeper layers, causal decodability of ToM states is maximal in shallow regions (Layers 0–3), with rapid representational collapse in later layers focused on token generation.

Figure 2: Layer-wise probing on NegotiationToM with Qwen2.5 highlights ToM information concentration in shallow layers.

Additionally, causal tracing reveals semantic stratification: desire states are more readily encoded and decodable than intention, with intention consistently showing lower reconstruction accuracy—mirroring developmental and modular theories of ToM in human cognition.

Causal-Oriented Steering: From Interpretation to Intervention

Building on causal localization, CoSToM implements an activation steering paradigm. A dual-model architecture is employed, comprising a context encoder and a frozen, probe decoder. The decoder acts as both a verifier (during training) and a standard generator (at inference). The steering objective supervises encoder activations, aligning them to ground-truth BDI states via a cross-entropy loss, while a gradient bridge mechanism restricts updates to only the ToM-critical layers; all other parameters, including the decoder, remain fixed.

Figure 3: Overview of the CoSToM framework’s steps: causal tracing locates ToM, then targeted steering aligns representations as a PEFT plug-in.

CoSToM’s intervention yields two central outcomes:

• Stability: It robustly mitigates ToM feature decay, maintaining high-fidelity BDI representations across the transformer stack—even in deep layers that would otherwise experience near-total information loss.
• Magnitude: On both Llama-3 and Qwen2.5, CoSToM amplifies ToM signal recovery, e.g., boosting agent desire accuracy from sub-3% to over 50% in deep layers—a marked contrast to standard full-model fine-tuning or prompt scaffolding.

  Figure 1: CoSToM stabilizes and amplifies ToM signals (Llama-3 NegotiationToM), contrasting rapid baseline collapse in higher layers.

  

  Figure 4: On PersuasiveToM, CoSToM increases deep-layer recoverability for persuadee intent and sends persuader desire tracking to near-perfect.

Downstream Behavioral Impact: ToM-Aligned Dialogue Generation

A critical empirical question is whether intrinsic ToM alignment (as realized through latent intervention) translates into improved social agency in dialogue. CoSToM is evaluated against zero-shot, prompt-based, and parameter-efficient fine-tuning baselines, with downstream generation judged using both LLM-based and human expert metrics on ToM reasoning, contextual coherence, and strategic effectiveness.

CoSToM consistently outperforms all baselines, delivering both substantially higher ToM scores (e.g., Llama-3 NegotiationToM: 0.499 vs. 0.081 for zero-shot, with structure-preserving partial-layer tuning even doubling global LoRA baselines) and greater strategic flexibility in producing contextually grounded, socially appropriate responses.

Figure 5: Example negotiation: the vanilla model infers the correct user desire (firewood) but fails to act on it in unconstrained setting; only explicit ToM reasoning, or CoSToM intervention, yields correct behavior.

Performance gains generalize across tasks (negotiation, persuasion), architectures, and evaluation paradigms. Most notably, CoSToM achieves high ToM-alignment without global parameter updates, supporting the hypothesis that social cognition is a localized cognitive function in transformer-based architectures, and that excessive layer-level tuning may attenuate strategically critical modules via overwriting and noise.

Figure 6: Qualitative generation contrast in persuasion: CoSToM resolves reasoning pitfalls, yielding coherent, empathetic, and ToM-grounded dialogue.

Theoretical and Practical Implications

The adoption of a mechanistic, causal-intervention perspective for ToM grounding yields both empirical and conceptual advances:

• Mechanistic Interpretability: Causal tracing and activation-level probing clarify where ToM is represented, offering an actionable decomposition of LLM social reasoning and opening pathways for fine-grained interpretability research.
• Localized Social Cognition: The results validate theories positing distributed yet locally concentrated social reasoning in neural architectures, aligning with similar findings in concept bottlenecking and factual model editing.
• Plug-and-Play Alignment: The PEFT-compatible, semi-localized alignment is scalable and modular; it can be readily applied to open-weight LLMs for deployment in negotiation, recommendation, and safety-critical cooperative agents.
• Advancing Dialogue Robustness: Practical downstream gains—elevated ToM-consistency and strategic coherence—suggest clear utility for RLHF augments, controllable dialog agents, and long-horizon planning in real-world, multi-agent applications.

Limitations and Future Directions

CoSToM’s reliance on open source model weights is a structural constraint; closed LLMs with proprietary internals are inaccessible for this form of layerwise intervention. Furthermore, while efficacy on negotiation and persuasion tasks is established, extension to broader interactional contexts (e.g., non-cooperative, multi-party, or multimodal ToM) remains an open frontier. Richer task distributions and more granular social-cognitive bandits will likely be necessary to fully instantiate human-like ToM in artificial agents.

Conclusion

CoSToM advances the state of the art in ToM alignment for LLMs by integrating causal tracing and targeted activation steering. The resultant alignment is both intrinsic and robust, unlocking human-level social reasoning in practical, downstream dialogue. This work motivates continued research at the intersection of mechanistic interpretability and social cognition in foundation models, especially for safe and adaptive deployment in agentic, multi-party environments.

Reference: "CoSToM: Causal-oriented Steering for Intrinsic Theory-of-Mind Alignment in LLMs" (2604.10031)