💬 NLP

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

The paper introduces CoSToM, a framework that enhances Large Language Models' intrinsic Theory-of-Mind capabilities by using causal tracing to identify critical internal layers and applying targeted activation steering to align internal knowledge with stable, high-quality social reasoning behaviors.

Original authors: Mengfan Li, Xuanhua Shi, Yang Deng

Published 2026-04-14

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Mengfan Li, Xuanhua Shi, Yang Deng

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Robot Actor" vs. The "Real Thinker"

Imagine you are watching a very talented actor on stage. The script says, "The character is sad," and the actor immediately starts crying. They look perfect. But if you ask them, "Why are you sad?" while they are off-stage, they might say, "I don't know, the script just told me to cry."

This is the current state of many Large Language Models (LLMs) when it comes to Theory of Mind (ToM). ToM is the human ability to understand that other people have their own thoughts, beliefs, and desires that are different from your own.

Current AI models are great at mimicking this behavior when you give them a specific prompt like, "Pretend you are a negotiator who understands the other person's feelings." But if you remove that prompt and just say, "Negotiate with me," the AI often forgets everything. It stops thinking about the other person's feelings and starts spitting out random or rude answers.

The researchers ask: Does the AI actually "know" what the other person is thinking deep inside its brain, or is it just acting on command?

The Solution: COSTOM (The "Mind-Steering" System)

The authors created a framework called COSTOM (Causal-oriented Steering for ToM). Think of it as a way to tune the AI's brain so it actually thinks about other people's feelings, rather than just pretending to.

They did this in three steps:

1. The Detective Work: Finding the "Empathy Switch"

First, they needed to find where in the AI's massive brain the concept of "other people's feelings" lives.

The Analogy: Imagine a giant library (the AI) with millions of books (layers of data). You want to find the specific shelf where the "Empathy" books are stored.
The Discovery: Using a technique called Causal Tracing, they acted like detectives. They "plugged in" a probe to different parts of the AI's brain while it read a conversation. They found that the AI does understand feelings, but it happens very early in its processing—like the first few pages of a book. If you wait until the end of the book (the deeper layers), the AI forgets the feelings and just focuses on finishing the sentence.

2. The Tuning: "Steering" the Brain

Once they found the "Empathy Switch" (the early layers), they didn't want to rebuild the whole library. They wanted to just tweak that specific shelf.

The Analogy: Imagine a car engine. Usually, to make a car go faster, you might try to rebuild the whole engine (Full Fine-Tuning). But COSTOM is like adding a small, precise turbocharger just to the air intake valve.
How it works: They used a "Gradient Bridge." Imagine a teacher standing next to the AI. When the AI thinks about a feeling, the teacher checks if it's right. If the AI is wrong, the teacher sends a signal backwards through the brain, but only to the specific "Empathy Switch" layers found in step 1. This "steers" the AI to keep those feelings active and strong, even as it moves to the later stages of generating a response.

3. The Result: A Natural Conversationalist

Finally, they tested the tuned AI in real conversations (like negotiating for firewood or persuading someone to buy a violin).

The Outcome: Before COSTOM, the AI was like a robot that needed a reminder to be nice. After COSTOM, the AI became like a natural human. It didn't need a prompt saying "be empathetic." It naturally understood what the other person wanted and negotiated or persuaded them effectively.
The Magic: It wasn't just "better at math." It was better at socializing. It could maintain a coherent, friendly, and strategic conversation without getting confused or rude.

Why This Matters

Most previous methods tried to fix AI by giving it better instructions (Prompting) or retraining the whole model (Fine-tuning).

Prompting is like telling a student, "Remember to be polite!" It works for a second, but they might forget.
Full Fine-tuning is like making the student re-learn the entire school curriculum just to be polite. It's expensive and slow.

COSTOM is like giving the student a permanent, internal habit of empathy. It changes the way the AI thinks at a fundamental level, ensuring that understanding other people's minds becomes a stable part of its personality, not just a trick it performs when asked.

Summary in One Sentence

The researchers found the specific "empathy circuits" inside AI brains, gave them a gentle nudge to stay active, and turned robotic actors into genuine social thinkers who can negotiate and persuade naturally.

1. Problem Statement

While Large Language Models (LLMs) often perform well on static Theory of Mind (ToM) benchmarks (e.g., inferring beliefs or desires in controlled scenarios), they frequently fail to generalize these capabilities to complex, dynamic social interactions like negotiation or persuasion.

The Gap: There is a critical misalignment between an LLM's internal knowledge (which may correctly infer a user's mental state) and its external behavior (which often fails to utilize that inference in dialogue).
The Limitation of Current Methods: Existing approaches rely heavily on "prompt scaffolding" (explicitly instructing the model to "think about beliefs") or external neuro-symbolic frameworks. These methods are brittle; once the explicit instruction is removed, the model reverts to incoherent generation, suggesting it lacks intrinsic, stable ToM cognition.
Core Question: Do LLMs possess intrinsic ToM reasoning, and can this internal knowledge be manipulated to produce stable, high-quality social behaviors without relying on external prompts?

2. Methodology: The COSTOM Framework

The authors propose COSTOM (Causal-oriented Steering for ToM alignment), a framework that moves from mechanistic interpretation to active intervention. It operates in three stages:

A. ToM Interpretation (Locating Representations)

Causal Tracing: The authors employ causal tracing to "read" implicit mental states from the model. They use a dual-model setup: a Context Encoder (processing dialogue history) and a Probe Decoder (frozen).
Mechanism: The encoder's hidden activations at specific layers ( $h_\ell$ ) are injected into the decoder, which is tasked with answering ToM-focused questions (e.g., "What is Agent 1's desire?").
Goal: Identify which layers contain the necessary information to reconstruct Beliefs, Desires, and Intentions (BDI).

B. ToM Steering (Active Alignment)

Activation Intervention: Once ToM-sensitive layers are identified, COSTOM steers the model's internal activations to align with ground-truth mental states.
Gradient Bridge Mechanism:
- A supervised loss ( $L_{ToM}$ ) is calculated based on the difference between the decoder's output and the true BDI labels.
- Crucially, gradients are backpropagated through the frozen decoder, across the patched activation interface, and into the Context Encoder.
- Parameter Efficiency: Only the layers preceding the intervention point (identified as shallow/early layers) are updated using LoRA (Low-Rank Adaptation). The deeper layers remain frozen. This "steers" the encoder to spontaneously generate ToM-enriched representations with minimal parameter updates.

C. ToM Leveraging (Downstream Generation)

Inference: During actual dialogue generation, the tuned encoder (now containing aligned ToM representations) processes the dialogue history.
Generation: The refined activations are passed to the decoder, which is prompted with task-specific instructions (e.g., "Negotiate for firewood") rather than ToM questions. The decoder translates the internal ToM reasoning into coherent, socially appropriate responses.

3. Key Contributions

ToM Interpretation Discovery: The paper systematically traces ToM features and reveals a counter-intuitive finding: ToM-related information is predominantly encoded in the early (shallow) layers of LLMs (e.g., Layers 0–3), not the deep layers where high-level reasoning is typically assumed to reside.
Efficient Intervention Framework: COSTOM introduces a lightweight, plug-and-play alignment method. By targeting only specific shallow layers via activation steering, it induces stable, human-like social reasoning without full fine-tuning.
Bridging Inference and Behavior: The framework successfully translates internal ToM alignment into improved downstream dialogue performance, proving that intrinsic cognition can be externalized into stable behaviors.

4. Experimental Results

The authors evaluated COSTOM on NEGOTIATIONTOM and PERSUASIVETOM datasets using Llama-3-8B and Qwen2.5-7B.

RQ1 (Interpretation): Causal tracing confirmed that ToM features (Belief, Desire, Intention) are most decodable in shallow layers. "Desire" and "Belief" peaked in early layers, while "Intention" showed a slightly deeper but still relatively shallow peak.
RQ2 (Steering Efficacy):
- Stability: COSTOM prevented "representation collapse." In baseline models, ToM information decayed rapidly in deep layers; COSTOM maintained high-fidelity representations across all layers.
- Magnitude: It rescued signals from near-total collapse (e.g., improving Agent 1's desire accuracy from ~2% to ~51% in deep layers) and amplified early-layer signals (e.g., boosting persuader desire tracking to ~99%).
RQ3 (Dialogue Generation):
- Performance: COSTOM significantly outperformed baselines (Zero-shot, Prompt-based MindDial, Full-Layer LoRA, and LatentQA).
- Metrics: On the Negotiation task with Llama-3, COSTOM achieved a ToM reasoning score of 0.499, nearly double the Full-Layer LoRA baseline (0.245) and a 6x improvement over the Zero-shot baseline (0.081).
- Qualitative: Human and LLM judges noted that COSTOM reduced identity confusion and semantic drift, producing responses that were empathetic and strategically aligned with inferred mental states.

5. Significance and Implications

Mechanistic Insight: The paper challenges the assumption that complex reasoning resides solely in deep layers, showing that fundamental social reasoning (ToM) is encoded early in the processing pipeline.
Beyond Prompt Engineering: It demonstrates that LLMs possess latent ToM capabilities that are often suppressed or unaligned. By intervening at the activation level, these capabilities can be unlocked without massive retraining.
Practical Application: COSTOM offers a scalable, parameter-efficient solution for building dialogue agents that are genuinely socially aware, applicable to negotiation, persuasion, and recommendation systems.
Limitations: The method currently requires access to open-weights models to manipulate internal activations and gradients, limiting its application to closed-source APIs. Additionally, its generalizability to non-strategic social contexts remains to be fully explored.

In summary, COSTOM provides a robust mechanism to "steer" LLMs toward intrinsic Theory of Mind, transforming them from models that merely mimic social reasoning when prompted into agents that inherently possess and act upon social understanding.