Collaborative Theory of Mind: Learning to Self-Revise by Anticipating User Responses

Fanzhi Zeng; Anubrata Das; Maytal Saar-Tsechansky; Liu Leqi

cs.AIOct 9, 2026

Collaborative Theory of Mind: Learning to Self-Revise by Anticipating User Responses

Authors: Fanzhi Zeng, Anubrata Das, Maytal Saar-Tsechansky, Liu Leqi

Organizations: The University of Texas at Austin · Autodesk

Abstract

Large language models (LLMs) demonstrate strong individual task performance, yet still struggle to collaborate with users over extended, multi-turn interactions, for example by attempting solutions before understanding the user's needs or by failing to revise mistaken assumptions after receiving feedback. We propose Collaborative Theory of Mind (CToM), a framework that addresses these challenges by explicitly training LLMs to anticipate user responses and to use those predictions to improve their own behavior. CToM drafts a response, predicts how the user will respond, and uses this prediction to guide internal revision before presenting a response. During training, the user's reply to each draft serves as the target for predicting the next turn, and we train the assistant to minimize the negative log-likelihood of that reply. We combine this prediction objective with group-relative policy optimization (GRPO) to optimize both task performance and interaction efficiency. Across mathematical reasoning, coding, and writing tasks, CToM achieves higher collaborative task performance than all baselines; for example, on a collaborative mathematical reasoning task, it improves final accuracy from 74.1% to 79.6% compared with naive GRPO. Ablations show that predicting user responses also makes collaboration more efficient: under the same reward, adding prediction supervision reduces the number of user replies by 8% on coding and 17% on writing while improving task performance. Behavioral analyses of interaction trajectories show that user-response predictions become more accurate after training, and that more accurate predictions are associated with larger improvements in collaborative task performance through self-revision.

Figures & tables

Figure 1: A real example of an LLM without theory of mind and CToM on the same request. (a) Without a model of the user, the LLM answers at once with a long draft that misses the user’s need, and the user must correct it twice. (b) CToM first asks about the tone, then drafts internally: its predicted reply to the first draft flags the weak Ultron section, so it revises before showing anything, and the user accepts the first visible draft. The user of (a) reads three long drafts and writes two corrections, while CToM anticipates the problem and saves effort from the user.

Figure 2: Overview of the CToM training pipeline. Rollout: given the user’s request $u_{0}$ ①, the LLM chooses an action ②; when it drafts, it predicts the user’s reply $\hat{u}_{t+1}^{(k)}$ to each draft $d_{t}^{(k)}$ and revises until its prediction is classified as satisfied, then sends $a_{t}=d_{t}^{(k)}$ to the user ③. The user replies, and this repeats until the conversation $\tau$ ends. Training: each discarded draft is sent to the simulated user in a counterfactual query ④; its reply, or the user’s actual reply for the shown draft, is the target $\tilde{u}_{t+1}^{(k)}$ that supervises the prediction $\hat{u}_{t+1}^{(k)}$ through $\mathcal{L}_{\mathrm{pred}}$ ⑤. The conversation $\tau$ defines $\mathcal{L}_{\mathrm{GRPO}}$ , and the sum of the two losses updates the LLM ⑥.

	Math			Coding			Writing
Method	Score $\uparrow$	Turns $\downarrow$	Tokens $\downarrow$	Score $\uparrow$	Turns $\downarrow$	Tokens $\downarrow$	Score $\uparrow$	Turns $\downarrow$	Tokens $\downarrow$
Prompting Base	69.14 $\pm$ 1.97	1.23 $\pm$ 0.03	1104 $\pm$ 37 0	34.00 $\pm$ 3.19	1.65 $\pm$ 0.04	235 $\pm$ 8 0	13.69 $\pm$ 0.36	1.00 $\pm$ 0.00	259 $\pm$ 12
Naive GRPO	74.08 $\pm$ 1.67	2.16 $\pm$ 0.06	2255 $\pm$ 192	48.91 $\pm$ 3.24	2.01 $\pm$ 0.00	253 $\pm$ 8 0	17.92 $\pm$ 0.29	1.00 $\pm$ 0.00	353 $\pm$ 6 0
CToM	79.62 $\pm$ 1.65	1.78 $\pm$ 0.05	541 $\pm$ 32 0	51.05 $\pm$ 3.08	3.57 $\pm$ 0.08	460 $\pm$ 26	19.08 $\pm$ 0.51	4.89 $\pm$ 0.09	960 $\pm$ 39
CToM w/o effort	82.16 $\pm$ 1.62	1.71 $\pm$ 0.05	871 $\pm$ 45 0	54.42 $\pm$ 3.19	3.98 $\pm$ 0.07	527 $\pm$ 27	19.18 $\pm$ 0.40	6.00 $\pm$ 0.00	1396 $\pm$ 48
CToM w/o prediction	77.81 $\pm$ 1.77	1.66 $\pm$ 0.05	290 $\pm$ 28 0	46.98 $\pm$ 3.13	3.78 $\pm$ 0.08	571 $\pm$ 32	18.03 $\pm$ 0.44	5.69 $\pm$ 0.06	1198 $\pm$ 61

Table 1: Main results for CToM , baselines, and ablations. Scores are multiplied by 100 in all domains; turns and tokens are per-conversation totals of visible assistant output and describe how much each method interacts (user effort is compared in Table 2 ). Subscripts give

\pm 1

standard error of the mean. CToM w/o effort is trained with

w_{\mathrm{effort}}=0

in Eq. ( 4 ). CToM w/o prediction is trained with

\lambda_{\mathrm{pred}}=0

in Eq. ( 6 ). CToM task-only removes both terms.

			Math		Coding		Writing
Method	$\mathcal{L}_{\mathrm{pred}}$	$C_{\mathrm{effort}}$	Replies $\downarrow$	Tokens $\downarrow$	Replies $\downarrow$	Tokens $\downarrow$	Replies $\downarrow$	Tokens $\downarrow$
CToM	✓	✓	2.24 $\pm$ 0.05	89.9 $\pm$ 2.5	2.57 $\pm$ 0.08	68.6 $\pm$ 4.2	3.89 $\pm$ 0.09	86.9 $\pm$ 2.9
CToM w/o effort	✓		2.23 $\pm$ 0.05	92.8 $\pm$ 3.1	2.98 $\pm$ 0.07	77.5 $\pm$ 3.8	5.00 $\pm$ 0.00	113.3 $\pm$ 2.0
CToM w/o prediction		✓	2.28 $\pm$ 0.04	85.7 $\pm$ 2.4	2.78 $\pm$ 0.08	73.4 $\pm$ 4.1	4.69 $\pm$ 0.06	105.8 $\pm$ 2.9
CToM task-only			2.42 $\pm$ 0.04	105.8 $\pm$ 2.8	3.29 $\pm$ 0.07	82.4 $\pm$ 4.0	4.79 $\pm$ 0.05	104.5 $\pm$ 2.7

Table 2: User effort for CToM and its ablations. Replies are the user’s messages after the opening request, and tokens are their total length; both are measured over conversations with at least one reply (Appendix B.2 ). Subscripts give

\pm 1

standard error of the mean. Bold marks the lowest value in each column.

Panel of Figure 3 — Figure 3: Behavioral analysis of CToM . (a) Lexical F1 between predicted and observed user replies at five ordered checkpoints for CToM (solid) and CToM w/o prediction (dashed); the first is the untrained model, and every trained CToM checkpoint outperforms it. (b) Mean math self-revision gain ( $\times 100$ ) for CToM and CToM w/o prediction, at revision caps of one and three. (c) Coding score after a failing first program, with conversations aligned at the turn that first contains a program; a finished conversation keeps its final score.

Figure 4: CToM without and with action-entropy regularization. (a) Observed Ask , Draft , and Commit proportions; the regularizer raises the Ask rate. (b) Task score, multiplied by 100. Off and On denote whether the regularizer is applied.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Hyperparameter	Value
LoRA rank	32
GRPO Rollouts $G$ per prompt	4
Prompt batch size	36
PPO minibatch size (prompts)	12
Actor learning rate	$10^{-4}$
Learning-rate schedule	Cosine

Appendix

Table 3: Shared training hyperparameters.

Setting	Math	Coding	Writing
Training tasks	3,263	480	500
Epochs / final update	2 / 180	16 / 208	16 / 208
Maximum visible assistant rounds $N$	4	6	6
Internal revision cap	3	3	3
Assistant evaluation temperature	1	1	1
Simulator evaluation temperature	0.7	0.7	0.7

Appendix

Table 4: Task-specific training and evaluation settings. All lengths and caps are in tokenizer tokens.

Task	Method	$\Delta$ ASK	$\Delta$ Score
Math	CToM	$+22.14$	$-7.05$
	CToM task-only	$+36.77$	$-4.24$
Coding	CToM	$+17.75$	$-2.12$
	CToM task-only	$+15.71$	$-4.46$
Writing	CToM	$+13.39$	$-0.29$
	CToM task-only	$+8.13$	$-0.33$

Appendix

Table 5: Effect of the action-entropy regularizer. Each difference is the value at

\lambda_{H}=10^{-3}

minus the value at

\lambda_{H}=0

, using the runs in Table 6 (revision cap three in every domain).

\Delta

ASK is in percentage points of all action decisions;

\Delta

Score is the change in task score, multiplied by 100.

Task	Method	Inner	$\lambda_{H}$	ASK	DRAFT	COMMIT	Score
Math	CToM	3	0	0.00	97.59	2.41	79.62 $\pm$ 1.65
	CToM	3	$10^{-3}$	22.14	68.29	9.58	72.57 $\pm$ 1.85
	CToM task-only	3	0	0.00	97.36	2.64	77.19 $\pm$ 1.75
	CToM task-only	3	$10^{-3}$	36.77	55.44	7.79	72.95 $\pm$ 1.85
Coding	CToM	3	0	2.89	70.19	26.92	51.05 $\pm$ 3.08
	CToM	3	$10^{-3}$	20.64	54.39	24.97	48.93 $\pm$ 3.07

Appendix

Table 6: Observed actions and task score of each arm. ASK, DRAFT, and COMMIT are percentages of all recorded action decisions; scores are multiplied by 100, as in Table 1 , with

\pm 1

standard error of the mean.

\lambda_{H}

is the action-entropy coefficient.

Domain	Method	0	50	100	150	200	Final
Math	CToM	$11.91_{\pm 0.17}$	$18.86_{\pm 0.29}$	$19.67_{\pm 0.32}$	$18.45_{\pm 0.30}$	–	$19.70_{\pm 0.31}$
	w/o prediction	$11.8_{\pm 0.2}$	$14.1_{\pm 0.2}$	$13.6_{\pm 0.2}$	$14.0_{\pm 0.2}$	–	$14.0_{\pm 0.2}$
Coding	CToM	$14.00_{\pm 0.58}$	$42.43_{\pm 1.92}$	$44.12_{\pm 2.15}$	$44.71_{\pm 2.04}$	$43.72_{\pm 2.03}$	$43.83_{\pm 2.11}$
	w/o prediction	$13.8_{\pm 0.5}$	$11.4_{\pm 0.5}$	$12.8_{\pm 0.6}$	$12.4_{\pm 0.5}$	$12.8_{\pm 0.5}$	–
Writing	CToM	$10.5_{\pm 0.6}$	$42.45_{\pm 1.72}$	$37.74_{\pm 1.47}$	$37.69_{\pm 1.62}$	$36.60_{\pm 1.52}$	$35.81_{\pm 1.51}$
	w/o prediction	$10.4_{\pm 0.6}$	$13.8_{\pm 0.7}$	$16.3_{\pm 0.7}$	$13.8_{\pm 0.6}$	$15.6_{\pm 0.8}$	$15.8_{\pm 0.7}$

Appendix

Table 7: Lexical prediction F1 (Eq. ( 9 ),

\times 100

) of CToM and CToM w/o prediction at each evaluated checkpoint, with task-clustered SEM. “Final” is the checkpoint in Table 1 (update 180 for math, 208 for coding and writing); a dash marks an update that was not evaluated.

Figure 5: LLM-judge agreement between predicted and observed user replies (

\times 100

) for CToM (solid) and CToM w/o prediction (dashed), with task-clustered SEM, at updates 0, 50, 100, 150, and a final update (180 for math, 200 for coding and writing). The first point is the untrained model.

Domain	Hidden information	Disclosure rule
Math	Reference answer	Judges correctness; never reveals, quotes, or hints at the answer. May name a wrong step.
Coding	Full behavioral specification	Answers the question asked; never pastes the specification or writes code.
Writing	Author’s notes and target length	Gives the minimum fact needed; never pastes material or dictates text.

Appendix

Table 8: Information available to the user simulator. The assistant sees only the opening request.

Explore similar work

Training LLMs for Multi-Agent Collaboration with Verifiable Information Rewards
Oct 13, 2026Keru Chen, Sen Lin, Yingbin Liang +2LLM Agent Training Multi-Agent Collaboration
School of ECEE, Arizona State University · Computer Science Department, University of Houston · Department of ECE, The Ohio State University +1
Humans' ALMANAC: A Human Collaboration Dataset of Action-Level Mental Model Annotations for Agent Collaboration
Jun 4, 2026Jiaju Chen, Yuxuan Lu, Jiayi Su +10Multi-Agent Collaboration AI Agent Benchmarks
Northeastern University · University of Notre Dame · Adobe +3
SEPAL: Separated Expert Pairs with Answer-Level Fusion for Reliable LLM Collaboration
Sep 30, 2026Weijie Ren, Yanwen Zhang, Hao Li +3Multi-Agent LLM Systems Question Answering
Zhejiang University · University of Electronic Science and Technology of China · University of Science and Technology of China

Oct 13, 2026J/K move · Enter open · S save

Training LLMs for Multi-Agent Collaboration with Verifiable Information Rewards

Keru Chen, Sen Lin, Yingbin Liang, Nathaniel D. Bastian+1

School of ECEE, Arizona State University · Computer Science Department, University of Houston · Department of ECE, The Ohio State University+1

LLM Agent Training Multi-Agent Collaboration Multi-LLM Collaboration Reinforcement Learning with Verifiable Rewards Reward Shaping

Open paper AskPDF arXiv

Effective collaboration among large language model (LLM) agents in multi-agent systems requires sharing relevant evidence, requesting missing information, and incorporating teammates' contributions into decisions. Training these behaviors is challenging because final task outcomes provide limited feedback on the value of individual messages. We introduce CoSIR (Collaboration via Surprisal-based Information Rewards), a post-training approach that uses a simple, verifiable environment to teach collaboration behaviors that generalize to more complex, unseen tasks. We propose a Bayesian evaluator that scores shared information by its reduction in team surprisal, defined as the sum of the negative log probabilities assigned to the correct option given each agent's available evidence. Requests are scored by the potential reduction in the same quantity if the requested information were truthfully provided. Separate rewards encourage agents to use the received information in their private choices and calculations. We post-train Qwen3-4B, Qwen3-8B, and Gemma-3-4B-it through self-play and evaluate transfer to seven unseen natural-language collaboration tasks. CoSIR achieves the highest aggregate score, while its performance in general tasks remains competitive. Behavioral analyses show increased disclosure of private evidence and fewer judged failures to use teammates' input.

Abstract

Explore similar work

Training LLMs for Multi-Agent Collaboration with Verifiable Information Rewards

Humans' ALMANAC: A Human Collaboration Dataset of Action-Level Mental Model Annotations for Agent Collaboration

SEPAL: Separated Expert Pairs with Answer-Level Fusion for Reliable LLM Collaboration