ToM2Vec: Continuous Theory-of-Mind Tokens for Multi-Agent Coordination

AI summary

Overview

Research area: Multi-agent AI / large language model agents, specifically theory-of-mind (ToM) modeling and latent (continuous) reasoning.

Technical level: Advanced. The paper assumes familiarity with LLM agent architectures, decoder-only transformer hidden states, self-play training, and partially observable multi-agent decision processes (Dec-POMDPs).

Scope: The paper proposes and evaluates ToM2Vec, a training method that teaches an LLM agent to predict a fixed budget of continuous latent tokens summarizing a partner agent's private reasoning, and tests whether this improves coordination across five cooperative, mixed-motive, and competitive environments.

What This Paper Is About

LLM agents that must cooperate or compete with another agent need a working theory of mind — an estimate of what the partner believes, wants, and will do next. The authors argue that natural language is a poor medium for this, because writing out a partner's intent commits the agent to a single guess, whereas the agent really needs an abstract representation that can hold several hypotheses at once. ToM2Vec trains agents to build that representation in latent space: each agent emits a small number of continuous "ToM tokens" that abstract over its partner's thoughts, and then conditions its own reasoning and actions on those tokens.

Key Contributions

1. A latent ToM training objective. The main agent autoregressively predicts K = 10 continuous latent tokens from its own local context, supervised by a cosine-similarity loss against latent tokens extracted from the partner's actual reasoning trace during self-play. This ToM loss is the only training signal; the method does not directly optimize task reward.

2. A concrete target-extraction scheme. The partner's reasoning trace is segmented at sentence delimiters, K − 2 interior chunks are formed and the hidden state at the center token of each chunk is taken, plus the hidden states of the start and end tokens, yielding K = 10 targets aligned position-wise with the main agent's predictions.

3. A multi-domain evaluation. ToM2Vec is tested on Overcooked-AI (three layouts, 30 seeds), multi-party negotiation (100 samples), Decrypto (1000 samples), Hanabi (100 seeds), and Collab Capture (6 seeds) / Collab Escape (12 seeds), against untrained, prompted-ToM, Language ToM, Self-Consistency, SimToM, Adaptive ToM, and oracle-partner baselines.

4. Evidence that the tokens encode a ToM rather than task heuristics. The authors run Bayesian ToM probes on beliefs, desires, and next actions; a perplexity-based corruption test (altered layout, adversarial intent); a Sally-Anne false-belief test on a custom cupboard layout; and a belief-revision experiment tracking east/west partner movement.

Main Findings

• Overcooked gains: ToM2Vec scores 98.0, 107.7, and 114.3 on Cramped Room, Asymmetric Advantages, and Coordination Ring (average 106.7), versus 90.0, 91.3, and 102.7 (average 94.7) for the untrained Qwen3-8B agent — a 12.0-point average improvement. Best baselines were Adaptive ToM at 98.9 average and Self-Consistency at 98.7.

• Oracle headroom: Giving the agent the partner's actual private thoughts raises the average to 117.7, i.e. 23.0 points over the base model — showing that partner information is highly valuable but prompting alone captures little of it (Prompted ToM averages 96.4, a 1.7-point gain, and drops below the base model on Cramped Room).

• Latent beats verbalized: Language ToM training, which instead decodes a text summary of the partner's reasoning, averages 97.6 in Overcooked and trails ToM2Vec by 6.0, 9.0, and 12.3 points on the three layouts. On Coordination Ring, Language ToM falls 0.7 points below the base model while ToM2Vec improves it by 11.6.

• Negotiation: Deal-success rate rises from 62.0% to 71.0%, own score from 59.2 to 65.6, and collective score from 54.7 to 59.2. Deal rate and own score are reported as significantly higher than all baselines.

• Decrypto: Decoder win rate improves from 10.4% to 14.8% and interceptor win rate from 87.4% to 93.5% (stated as 4.4 and 6.1 percentage points), beating both prompted ToM and Language ToM in both roles.

• Hanabi: Score of 11.3 out of 25, versus 8.9 for the untrained agent (a 2.4-point gain) and 10.6 for the strongest baseline, Language ToM (a 0.7-point margin).

• Collab Capture and Escape: Capture rate rises from 0.67 to 1.00, with average turns falling from 5.00 to 4.32 (Prompted ToM also reaches 1.00 but at 4.76 turns). Escape rate rises from 0.75 in Table 5 to 0.92 (the body text states 0.76 as the starting value), tying the best success rate at an unchanged turn count of 8.00.

• Transfer to unseen partners: Training only against a Qwen3-8B partner, ToM2Vec still improves Overcooked scores over 20 seeds with a GPT-5.4 partner (135/193/174 to 152/214/188, a 17.3-point average gain) and with a Gemma-4 partner (65/90/87 to 68/105/94). Neither model was seen during training.

• Probes of partner mental state: On 300 samples per question with six options each (random = 16.7%), ToM2Vec reaches 52.3% on beliefs, 62.3% on desires, and 64.0% on next actions, beating the next-best baseline by 5.6, 2.0, and 2.7 points. Beliefs are the hardest category for every method, which the authors attribute to second-order ToM.

• Context sensitivity: Latents derived from the true partner context assign perplexity 1.412 to the partner's real reasoning, versus 1.702 for an altered layout and 1.613 for an adversarial-intent partner; the correct chain-of-thought receives the lowest perplexity 98.6% of the time (0.6% and 0.8% for the two corruptions).

• False belief and revision: On a custom cupboard-based Sally-Anne test over 50 seeds, ToM2Vec scores 0.88 versus 0.74 for the untrained agent, tying Self-Consistency for best, while state understanding stays within 0.08 across methods. In a symmetric layout, moving the partner east or west shifts the perplexity of dish- versus onion-directed chains-of-thought in the expected direction.

Methodology in Plain English

The setup is a two-agent game where each agent sees only part of the world. Both agents start from the same base model, Qwen3-8B, which matters because the training signal compares vectors living in a shared representation space.

Training runs in self-play. The partner is asked to reason about the current situation, producing a text reasoning trace. The researchers chop that trace at sentence boundaries, divide it into K − 2 even chunks, and record the model's internal hidden state at the middle token of each chunk, plus the hidden states of the first and last tokens — ten vectors in total. These ten vectors are the "answer key."

The main agent, meanwhile, is given only what it can legitimately observe: its own history, the visible board state, and role instructions. It produces ten of its own hidden-state vectors one after another, without ever decoding them into words. A cosine-similarity loss pulls each of its predicted vectors toward the corresponding target vector from the partner's trace. Nothing else is optimized — no reward, no action imitation.

At test time the partner's private reasoning is gone. The agent simply emits its ten latent tokens from its own context, feeds them back into itself as the prefix for subsequent generation, and then produces its reasoning and action conditioned on them. The intent is that these tokens carry a compressed, multi-hypothesis sketch of what the partner is thinking, so the agent can adapt before the partner acts. Inference uses vLLM; training uses AdamW with a learning rate of 10⁻⁶ and a batch size of 4 with a single gradient-accumulation step.

Why This Matters

The paper argues against a common default — telling an agent to "think about your partner in words" — and shows empirically that a learned continuous representation does better on the same base model. It also shows the benefit survives contact with stronger, unseen partners, including a proprietary model, which suggests a practical recipe: train on open-weight models, deploy alongside frontier ones.

Impact on research: It connects two previously separate lines — latent/continuous chain-of-thought work (e.g. Coconut-style hidden-state feedback) and theory-of-mind work for LLM agents — and offers a diagnostic toolkit (Bayesian probes, corrupted-context perplexity, false-belief layouts, belief-revision curves) for testing whether a representation is genuinely partner-modeling rather than a task shortcut.

Real-world applications:

• Human-AI or AI-AI teaming in shared physical spaces (kitchens, warehouses, assembly), where anticipating a teammate's next move on shared resources matters more than modeling their long-term goals.
• Automated negotiation and deal-making assistants that must balance self-interest against the risk of no agreement.
• Multi-agent game AI and playtesting, including adversarial settings where one agent tries to intercept another's signals.
• Mixed human-AI workflow orchestration where agents must adapt to a counterpart's strategy on the fly without explicit communication.

Industry relevance: The method is a lightweight fine-tuning objective rather than a new architecture, and the demonstrated transfer from an 8B open-weight training partner to a frontier proprietary test partner speaks directly to deployment scenarios where teams mix models from different vendors.

Future Directions

• Removing the chain-of-thought requirement. The authors flag that ToM2Vec needs access to the partner's reasoning during training as its main limitation; a natural next step is distilling or imputing those targets so the method works without privileged traces.
• Breaking the shared-representation assumption. Training currently requires agent and partner to be the same base model so their hidden states are comparable. Learning a cross-model alignment would allow training against heterogeneous partners directly rather than relying on emergent transfer.
• Tuning the latent budget. All experiments use K = 10 tokens with no reported ablation over K, so how much abstraction capacity is optimal — and whether it should adapt to task complexity — remains open.
• Extending beyond English, two-agent, text-only settings. The authors note their tasks are English-only and constrained by available models and datasets; scaling to more than two agents, longer horizons, and other modalities is untested.

Target Audience

Researchers and engineers working on LLM agents, multi-agent systems, and machine theory of mind who are comfortable with latent-reasoning methods and self-play training. It is also relevant to practitioners building multi-agent products who want evidence on whether to prompt for partner modeling or train for it, and to cognitive-science-adjacent readers interested in operationalizing beliefs, desires, and actions as measurable probes in LLM agents.

Authors’ abstract

Forming an accurate theory-of-mind (ToM) is key to successful multi-agent collaboration and requires reasoning under uncertainty to capture multiple possible partner states. This is hard to do in the text representation space native to large language model (LLM) agents, which more easily describe a single partner state. To address this challenge, we introduce ToM2Vec, which learns to form abstract theories-of-mind in a rich latent space via a fixed budget of continuous reasoning tokens. During training, supervision for these tokens comes from the partner's actual thoughts, which the ToM tokens are trained to abstract and match. At test-time, ToM tokens are generated by each agent and act as conditioners for that agent's thoughts and actions. Across several environments covering collaborative, competitive, and mixed-motive settings (Overcooked-AI, multiparty negotiation, Decrypto, Hanabi, and Collab Escape/Capture), we evaluate ToM2Vec's ability to support multi-agent interactions and collaboration, finding consistent improvements over the base model and a variety of competitive trained and prompted baselines. We show that latent reasoning is critical, with performance dropping when reasoning is expressed in text. Moreover, we show that ToM2Vec training generalizes to partner model families not seen at train-time, including a proprietary model (GPT-5.4). In our analysis, we provide evidence that the latent tokens encode a ToM, rather than task-specific information: ToM2Vec's tokens improve the agent's ability to directly answer questions about partner beliefs, desires, and next actions.