ToM2Vec: Continuous Theory-of-Mind Tokens for Multi-Agent Coordination
Nithin Sivakumaran, Ali Payani, Hyunji Lee, Mohit Bansal, Elias Stengel-Eskin
cs.AI, cs.CL, cs.MA
2026-10-22
Self-play trains 10 continuous ToM tokens on Qwen3-8B. Overcooked average rises 94.7 to 106.7, negotiation deals 62% to 71%, and an unseen GPT-5.4 partner still adds 17.3 points.
When several LLM agents share a task, one's own next move is usually the easy part. The partner's move is not. Theory of Mind (ToM) is the step from visible behavior to beliefs, desires, and next actions the other agent never states. A shared pot, a six-party negotiation with private preferences, and Hanabi, where each player sees every card but their own, all break when that inference is off.
Prose is a poor place to keep several partner hypotheses alive. Once a sentence is sampled, the model tends to lock onto one partner state. Figure 1 shows the pattern with Qwen3-8B in Overcooked. With no ToM, both agents head for the dish dispenser and block each other. A model that verbalizes the partner's reasoning does look across the kitchen, then commits: the partner is about to pass a plate, so walk over and take it. The continuous-token agent keeps another branch open. The partner may grab the plate, so go for an onion instead.
ToM2Vec compresses that judgment into a fixed budget of continuous vectors, K = 10. The tokens never enter the vocabulary. Each step takes the last-layer hidden state and feeds it back as the next input embedding, for 10 autoregressive steps. The mechanism matches Coconut-style latent reasoning: a hidden state can carry several unfinished lines of thought without saying any of them aloud.
Supervision is self-play only. The loss has no task reward. The acting agent and the partner share one set of weights, initialized from Qwen3-8B. Cosine alignment needs a shared representation space, so the two roles start as the same model. The partner samples a full chain of thought from its private context. The trace is split on sentence boundaries into 8 chunks. The hidden state at the center token of each chunk is kept, and the hidden states at the start and end of the trace are added, producing 10 target vectors. The acting agent never sees the private context. It predicts the 10 vectors from its own observation and history. The loss is the gap in position-wise cosine similarity, L = 1 − (1/10) Σ cos(prediction, target).
At test time the partner's chain of thought is unavailable. The agent emits the 10 ToM tokens first, and later thoughts and actions are conditioned on them. Figure 2 draws this on multi-party negotiation: training aligns a compression of the partner's real calculus, and deployment keeps only that compression. Language ToM uses the same self-play loop, with the target swapped for a text summary of the partner's reasoning.
Every agent starts from Qwen3-8B. Overcooked uses three layouts and 30 seeds each. The metric is mean team score per episode.
| Method | Cramped Room | Asymmetric Adv. | Coordination Ring | Avg. |
| Untrained | 90.0 | 91.3 | 102.7 | 94.7 |
| Prompted ToM | 86.7 | 96.0 | 106.7 | 96.4 |
| Language ToM | 92.0 | 98.7 | 102.0 | 97.6 |
| Self-Consistency | 91.3 | 97.3 | 107.3 | 98.7 |
| SimToM | 90.7 | 94.0 | 105.3 | 96.7 |
| Adaptive ToM | 92.7 | 102.7 | 101.3 | 98.9 |
| Oracle | 110 | 118 | 125 | 117.7 |
| ToM2Vec | 98.0 | 107.7 | 114.3 | 106.7 |
ToM2Vec averages 106.7, 12.0 above the untrained model, and leads every layout among methods without privileged access. Prompted ToM gains 1.7 on average and drops on Cramped Room, from 90.0 to 86.7. Language ToM averages 97.6 and trails by 6.0, 9.0, and 12.3. On Coordination Ring the text version sits 0.7 below the base model. The continuous version sits 11.6 above it. Oracle, given the partner's actual next action and intent, averages 117.7 and still leads by 11 points. Self-Consistency samples three reasoning chains and takes the majority. SimToM and Adaptive ToM are earlier prompt-style ToM methods. Adaptive ToM is compared only on Overcooked.
| Task | Untrained | ToM2Vec | Nearest baseline |
| Negotiation deal rate | 62.0% | 71.0% | Prompted ToM 68.0% |
| Negotiation own score | 59.2 | 65.6 | SimToM 62.1 |
| Negotiation collective utility | 54.7 | 59.2 | Language ToM, also 59.2 |
| Decrypto decoder win | 10.4% | 14.8% | Language ToM 13.6% |
| Decrypto interceptor win | 87.4% | 93.5% | Language ToM 93.0% |
| Hanabi (max 25) | 8.9 | 11.3 | Language ToM 10.6 |
| Collab Capture rate | 0.67 | 1.00 | Prompted ToM also 1.00 |
| Collab Escape rate | 0.75 | 0.92 | SimToM also 0.92 |
Negotiation uses 100 games. A deal counts only if at least five of six parties accept. Decrypto runs 1,000 games in each role, with the other agents left as untrained Qwen3-8B. The decoder still wins only 14.8% of games. Collab Capture has 6 seeds. ToM2Vec finishes in 4.32 turns, fewer than Prompted ToM at 4.76, which also captures every episode. Collab Escape has 12 seeds. SimToM matches the 0.92 escape rate and is faster, 7.80 turns against 8.00.
Partner transfer is reported only for Overcooked, on 20 seeds, after training exclusively against Qwen3-8B. With GPT-5.4 the three layouts move from 135, 193, and 174 to 152, 214, and 188, a 17.3 point average gain. With Gemma-4-E4B they move from 65, 90, and 87 to 68, 105, and 94. Every layout is positive, and the margin is smaller.
A probe checks whether the vectors carry a mental model or only action-correlated noise. Three six-way questions, 300 items each, put chance at 16.7%. Belief asks what the partner thinks this agent is trying to do. Desire asks which subtask the partner is pursuing. Action asks for the partner's next primitive move. ToM2Vec scores 52.3%, 62.3%, and 64.0%, ahead of the next-best baseline by 5.6, 2.0, and 2.7 points. Replacing the context used to build the vectors with another layout raises perplexity on the partner's true reasoning from 1.412 to 1.702. A sabotage-oriented partner context lands at 1.613. The true context gives the lowest perplexity on 98.6% of items.
The training signal is hidden-state alignment, with no environment reward, so one set of weights transfers to cooperative, competitive, and mixed-motive tasks. Writing the partner's thoughts out in text buys little, and on one layout it hurts. The self-play partner can be an open 8B model, and the partner at deployment can be a closed model the training run never saw. On Overcooked the gain survives that swap.
Prompting barely moves the score. Oracle stays about 11 points ahead. The gain is real, and it is incremental.
Training needs the partner's chain of thought. The partner at deployment can change. The partner that supplies supervision cannot, because both sides must share a representation space, so training uses one base model for both roles. The method also needs a training stage. Prompt-only ToM does not. Every task is in English.
Collab Capture and Collab Escape use 6 and 12 seeds. A capture rate moving from 0.67 to 1.00 does not separate noise from gain. Transfer across model families is not shown for negotiation, Decrypto, or Hanabi. Belief accuracy of 52.3% beats chance and still sits far from a reliable second-order ToM. The paper accounts for the gap over text by arguing that continuous vectors can hold several partner hypotheses at once. The measurements are the score gap and sensitivity to a corrupted context. There is no control that spends the same 10 hidden steps on a target other than the partner's thoughts, so the share of the gain that is just extra thinking budget is still open. Significance tests are in the appendix, and the appendix's false-belief results are not quantified in the main text.