PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2610.12872v1cs.AI

TL;DR

LLMs can perform strongly on individual tasks yet struggle to collaborate over multiple turns as user needs become clearer. CToM trains assistants to predict user replies and revise drafts internally, and it outperforms baselines across math, coding, and writing tasks. The experiments use a simulated user, leaving performance with real users to be tested.

  • Problem

    LLMs may act before understanding users’ needs and struggle to revise mistaken assumptions during multi-turn collaboration.

    Source: §1 · p. 1

  • Method

    CToM trains assistants to predict user responses to drafts and use those predictions to revise internally before responding.

    Source: §1 · p. 2

  • Results

    CToM improves task score over the strongest baseline by 7.5%, 4.4%, and 6.5% on math, coding, and writing, respectively.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Across math, coding, and writing, findings support that anticipating user feedback helps assistants revise more effectively and reduces user input effort.

    Source: §6 · p. 11

    Experiments use a simulated user, so performance and user preference in real-world interaction remain to be tested.

    Source: §6 · p. 11

Abstract

from arXiv · show

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.

1 Introduction

The introduction frames a gap between strong individual task performance and effective multi-turn collaboration, then proposes CToM to predict user responses and revise drafts before they are shown. It reports gains across mathematical reasoning, coding, and writing, alongside findings on efficiency and interaction strategies.

1 Introduction

  • Strong individual task performance does not reliably ensure effective collaboration: assistants may act before understanding user needs and require repeated corrections.The figure contrasts a user correcting multiple drafts with CToM revising before presenting its first visible draft.Source: §1 · p. 1, Figure 1 · p. 2
  • CToM predicts user feedback on a draft and revises internally before showing a response, with real user replies supervising those predictions during training.The framework combines supervised prediction loss with GRPO for task behavior.Source: §1 · p. 2
  • The framework operationalizes theory of mind as anticipating user replies and using those predictions to improve the assistant’s own response.The motivation draws on the role of anticipating others’ reactions in human collaboration.Source: §1 · p. 2
  • CToM improves task score over the strongest baseline by 7.5%, 4.4%, and 6.5% on math, coding, and writing, respectively.The evaluation trains a Qwen3-8B assistant with a Qwen3-32B user simulator.Source: §1 · p. 2
  • Prediction supervision raises task score in every ablation comparison, while the combined CToM variant requires the least user input effort among those variants.The reported effort measures include the number and length of user replies.Source: §1 · p. 3
  • After training, user-reply predictions become more accurate and are associated with more effective revision within and across conversation turns.The introduction also reports that encouraging more clarification questions lowers task score in every domain.Source: §1 · p. 3

2 Related Work

Related work examines interactive evaluation, collaborative language models, and theory-of-mind capabilities in LLMs. The paper positions its approach as training assistants to anticipate user responses and revise proactively.

Interactive Evaluation

  • Multi-turn evaluation work highlights premature assumptions and argues for assessing both interaction quality and final task performance [7] [9].User intent and task completion may become clearer over the course of interaction.Source: §2 · p. 3

Collaborative LLMs

  • Collaborative LLM research spans trace analysis, multi-agent coordination, proactive collaboration, and simulated-interaction training [10] [20].The paper contrasts its general-purpose approach with methods that use specific collaboration rewards or proactive personalization.Source: §2 · p. 3

Theory of Mind

  • Prior theory-of-mind work studies mental-state inference, specialized agents, and settings with communication or unequal information access.FANToM evaluates conversational belief tracking, while ToM-SWE models user goals, constraints, and preferences [24].Source: §2 · pp. 3–4
  • Prior work evaluates or trains agents to detect belief discrepancies, while static-benchmark improvements may not transfer to dynamic environments [16].The related-work discussion identifies collaborative-task RL as formative ground for implicit theory-of-mind capabilities.Source: §2 · p. 4

3 Problem Setting

The problem setting models collaboration as an assistant completing a task through a multi-turn conversation with a user. The assistant lacks direct access to some task requirements, and the objective values task performance alongside interaction effort.

3 Problem Setting

  • Each task instance has an initial request and task information, including user requirements or goals that are available to the user but not directly to the assistant.The assistant completes the task through a multi-turn conversation.Source: §3 · p. 4
  • At each turn, the assistant responds to the visible dialogue, and a continuing user reply is drawn from a response distribution conditioned on the dialogue, assistant message, and task information.The conversation ends under a task stopping rule or when its interaction budget is exhausted.Source: §3 · p. 4
  • The policy objective maximizes expected conversation return, instantiated as task performance with penalties for interaction effort measured by turns and user-visible message length.The return may generally combine task success with user-relevant costs such as effort, time, or experience.Source: §3 · p. 4

4 Method: Collaborative Theory of Mind (CTOM)

CToM trains an assistant to anticipate user feedback, revise drafts before showing them, and balance task success with the effort users spend interacting. It asks when progress depends on information only the user can provide.

4 Method: Collaborative Theory of Mind (CTOM)

  • CToM combines asking when only the user can supply missing information with hidden draft–predict–revise behavior, supervised by user replies and optimized for task success and user effort.Shown drafts use actual replies as targets; discarded drafts use simulated counterfactual replies.Source: §4 · p. 4

4.1 Prediction-Guided Interaction

  • Each round offers Ask, Draft, or Commit; Draft predicts user feedback and revises internally until satisfaction is predicted or the revision cap is reached.Internal drafts and predictions remain invisible to the user.Source: §4.1 · p. 5

4.2 Learning to Predict User Feedback

  • The prediction loss trains replies to each draft using actual feedback for the shown draft and simulated counterfactual feedback for discarded drafts.Targets are trained by minimizing negative log-likelihood under the same context used for prediction.Source: §4.2 · p. 6

4.3 Interaction Cost

  • Interaction cost combines normalized user-turn count with clipped, normalized token totals for visible user and assistant messages.Internal drafts and predictions are excluded from the visible-message totals.Source: §4.3 · p. 6, §4.5 · p. 7

4.4 Combined Reward

  • The conversation reward combines task-evaluator score and interaction cost, with task weight 1 and effort weight 0.01.Prediction quality and improvements between internal drafts receive no separate reward.Source: §4.4 · p. 6

4.5 Joint Training and Inference

  • Joint training combines GRPO with reply-prediction loss, capping prediction-gradient influence at half the GRPO gradient; inference uses the revision loop without training signals.Setting λpred to zero removes prediction supervision while retaining the same prompt and revision loop.Source: §4.5 · pp. 6–7

5 Experiments

Across collaborative math, coding, and writing, CToM leads the baselines; the experiments also examine user effort, prediction accuracy, revision, and action choices.

5.1.1 Datasets

  • Evaluation uses Level-5 MATH problems, BigCodeBench Python tasks, and MediumDocEdit writing; coding and writing prompts leave some requirements unspecified.Source: §5.1.1 · pp. 7–8

5.2 Main Results

  • CToM beats every baseline in all three domains, with relative gains over the strongest baseline of 7.5% on math, 4.4% on coding, and 6.5% on writing.Source: Table 1 · p. 8, §5.2 · p. 8

5.1.2 Baselines and Ablations

  • Baselines include an untrained minimal-prompt model and Naive GRPO; ablations remove CToM’s effort penalty, prediction supervision, or both while retaining its interactive prompt.Source: §5.1.2 · p. 8

5.2 Main Results

  • Prompting Base and Naive GRPO use at most 2.2 visible turns on average and end coding and writing conversations after the first program or draft.Their low turn and token counts reflect collapsed interaction rather than efficient collaboration.Source: §5.2 · p. 9

5.3 User Effort

  • Under the same reward, adding prediction supervision reduces user replies by 8% on coding and 17% on writing while also raising task scores.Across ablations, prediction supervision raises task scores in all six comparisons; removing CToM’s effort penalty raises scores further but increases visible tokens by 15–61%.Source: §5.2 · pp. 8–9, §5.3 · p. 9, Table 2 · p. 9

5.4 Behavioral Analysis

  • Trained CToM checkpoints beat the untrained model on reply-prediction F1, and prediction-supervised models gain more from math self-revision at both tested caps.Final reply F1 is 19.7, 43.8, and 35.8 for CToM versus 14.0, 12.8, and 15.8 without prediction on math, coding, and writing; cross-domain levels are not directly comparable.Source: §5.4 · pp. 9–10, Figure 3 · p. 10
  • After a failing first program, CToM reaches 38.4 after one more turn and ends at 46.2, versus 33.8 and 40.5 without prediction.Source: §5.4 · p. 10, Figure 3 · p. 10
  • Entropy regularization keeps Ask frequent across domains, but regularized runs score lower at their selected checkpoints.Without regularization, Ask nearly disappears on math and coding but remains common on writing, where requirements are less specified.Source: §5.4 · p. 10

6 Conclusion and Discussion

CToM combines supervised prediction of user replies with GRPO to revise drafts before sharing them. The authors argue that anticipating feedback supports more effective revision and useful interaction strategies.

6 Conclusion and Discussion

  • CToM combines supervised user-reply prediction with GRPO on a reward for task success and user effort, revising drafts before sharing them.The assistant predicts the user's reply to a draft and uses that prediction in an internal revision loop.Source: §6 · p. 10
  • Figure 4 reports Ask, Draft, and Commit proportions and task scores multiplied by 100, comparing runs with action-entropy regularization off and on.The regularizer raises the Ask rate.Source: Figure 4 · p. 11
  • Learning to anticipate user feedback helps assistants revise more effectively, reduce user input, and use useful interaction strategies across math, coding, and writing.Source: §6 · p. 11
  • Because experiments use a simulated user, real-user reply prediction, task performance, and user preferences remain to be tested.Source: §6 · p. 11
  • Predicting the next reply is only one form of theory of mind, and CToM primarily serves task completion rather than modeling users' personalities or preferences.Source: §6 · p. 11
  • CToM uses its user model only for Draft, leaving open how theory of mind should inform asking, question choice, and committing.Source: §6 · p. 11

A Implementation Details

The implementation uses group-relative advantages, a clipped policy objective with a reference-policy penalty, and a separate prediction-gradient adjustment.

A.1 GRPO Objective

  • GRPO advantages use sample standard deviation with denominator G−1 and ϵ = 10^-6; equal-reward groups receive zero advantage, and no critic is used.Source: §A.1 · p. 14
  • The loss uses a dual-clipped surrogate and a KL penalty toward the frozen reference policy, with λKL = 0.01 and a k3 estimator.Log-ratios and the KL estimate are clamped for numerical stability.Source: §A.1 · p. 14
  • Prediction and GRPO gradients are computed separately; the prediction gradient is rescaled by η before the gradients are added and clipped.Source: §A.1 · p. 14

A.2 Task Scores and Numerical Settings

Task scores are defined differently for math, coding, and writing, while α and β are shared across domains and maximum lengths are task-specific.

A.2 Task Scores and Numerical Settings

  • Math scores are 1 for a visible boxed answer equivalent to the reference and 0 otherwise; coding scores are unit tests passed, and writing uses ROUGE-L F1.Discarded drafts are not scored.Source: §A.2 · p. 14
  • All domains use α = 0.5 and β = 0.4, with task-specific maxima N, Bu, and Ba.Lengths are measured in tokenizer tokens.Source: §A.2 · p. 14

A.3 Satisfaction Classifier

A fixed lexical classifier uses predicted replies to decide whether to revise a draft, and its labels distinguish replies to incorrect versus correct math answers.

A.3 Satisfaction Classifier

  • After each draft, a fixed classifier labels the predicted reply Satisfied or Unsatisfied; an Unsatisfied prediction triggers revision if k < Kmax = 3.Otherwise, the draft is shown; the same classifier is used in training, evaluation, and the three ablations.Source: §A.3 · p. 14
  • The classifier uses mined lexical cues: pushback such as a question mark or phrases like “wrong” and “could you” is Unsatisfied, while empty, end-of-conversation, or clear-acceptance replies are Satisfied.Source: §A.3 · pp. 14–15
  • On mined math replies, the rules label 95.6% of replies to incorrect answers and 19.6% of replies to correct answers as Unsatisfied, a difference of 0.76.The classifier uses reply wording alone and never checks the answer.Source: §A.3 · p. 15

B Experiment Details and Additional Results

The appendix specifies task construction and interaction-cost accounting for the experiments. It defines how user effort is counted and reports task-clustered standard errors.

B.1 Task construction

  • Coding requests retain the public function signature while behavioral requirements remain available to the simulator; writing tasks group articles and keep full specifications separate from the public request.Training conversations and counterfactual replies are generated online, without using later source-conversation turns as assistant demonstrations.Source: §B.1 · p. 15

B.2 Interaction-Cost Counting

  • Interaction cost counts user messages, user-message tokens, and assistant-message tokens, measured in tokenizer tokens and normalized by task-specific maxima.The listed user-message count includes the opening request.Source: §B.2 · p. 15
  • User-effort reply and token averages include only conversations with at least one reply after the opening request, and count only those subsequent replies.For example, CToM’s 2.24 replies and 89.9 tokens on math correspond to about 3.24 assistant turns and two short replies per conversation.Source: §B.2 · pp. 15–16
  • Table 5 differences compare each value at λH = 10^-3 with λH = 0; ∆ASK is in percentage points and ∆Score is multiplied by 100.The reported standard errors are clustered by task.Source: Table 5 · p. 16, §B.2 · p. 16

B.3 Decision-Entropy Ablations

The appendix details action-entropy ablations and the measurement of within-round self-revision gains on math. Entropy regularization increased ASK decisions but reduced task scores in every matched comparison.

B.3 Decision-Entropy Ablations

  • The entropy regularizer enters the joint-loss gradient before the prediction-gradient cap, so it can also change CToM’s effective prediction-loss weight.CToM task-only has no prediction loss, isolating the regularizer’s effect on actions.Source: §B.3 · p. 16
  • All six matched action-entropy comparisons increased ASK and had lower task scores.Action shares pool recorded decisions rather than averaging conversation-level proportions; Table 6 reports ASK, DRAFT, and COMMIT shares and task scores.Source: §B.3 · p. 16, Table 6 · p. 17

B.4 Within-Round Self-Revision Gains on Math

  • Self-revision gain is the mean strict-answer score change from the initial to final draft over revised rounds, with no intervening user feedback.The measure averages rounds, not tasks or conversations, and excludes ASK-only rounds and rounds without a usable draft.Source: §B.4 · p. 17
  • The gain estimate uses 365 math tasks with eight conversations each and one evaluation seed per run, treating revised rounds as independent without task or conversation clustering.Error bars show ±1 standard error of the mean.Source: §B.4 · p. 17

B.5 Prediction Learning Across Checkpoints

This section describes how predicted user replies are evaluated, including the user-simulator instructions, judge, and scoring criteria, and reports that CToM’s predictions scored higher than those of untrained models across checkpoints.

B.5 Prediction Learning Across Checkpoints

  • At math update 0, 1,540 of 5,495 pairs (28%) scored zero because a prediction or reply was empty or missing, so part of the math gain reflects more reliable prediction emission.Because the pairing conventions differ, score levels are comparable within a run but not across domains.Source: §B.5 · p. 18
  • CToM’s final judge scores were 46.6 on math, 47.1 on coding, and 41.7 on writing, versus 29.9, 30.4, and 33.0 for CToM w/o prediction.Every trained CToM checkpoint also scored above the untrained model under the judge, which assessed reply substance rather than just shared wording.Source: §B.5 · p. 18, Figure 5 · p. 19

C Prompts

  • The appendix states that prompts were copied from saved evaluation records or producing code; all assistant prompts disable Qwen3’s thinking mode with /no_think.Long lines are wrapped for display, braces mark placeholders, and the Qwen3 chat template uses <|im_start|> and <|im_end|>.Source: C Prompts · p. 19

C.2 Assistant Prompts for Prompting Base and Naive GRPO

  • Prompting Base and Naive GRPO use one visible response per round without planning, prediction, or internal revision, and are instructed to ask clarifying questions before final answers on coding and writing tasks.Math conversations end when a boxed answer passes evaluation; coding and writing end with the marker “I WANT TO ANSWER:”.Source: §C.2 · pp. 22–23

C.3 User-Simulator Prompts

  • The user simulator receives task-specific prompts and hidden reference or specification information, while the assistant sees only the opening request and visible chat history.Simulator instructions vary by domain but generally require concise, natural replies without revealing hidden answers or requirements wholesale.Source: §C.3 · pp. 23–25
  • The user-simulator prompt specifies instructions for cases when the material does not answer the question or the AI submits a final document, and requires JSON with status, thought, and response fields.The response field is for the human's concise reply.Source: §C.3 · p. 26

C.4 LLM Judge for Prediction Agreement

  • The prediction-agreement judge compares predicted and observed replies using semantic similarity, intent alignment, specificity, and tone/style consistency.GPT-4o-mini scores each pair at temperature zero, with the parsed score clipped to [0, 1].Source: §C.4 · p. 26
  • Predictions must use the user's first-person voice; third-person descriptions are invalid and receive 0.0, while valid predictions are rated on a 0.0–1.0 match scale.The scale ranges from no match to near-perfect agreement in meaning, intent, and key details.Source: §C.4 · p. 26

C.1 Assistant Prompts for CTOM and Its Ablations

The prompts use a plan-and-action protocol with DRAFT, ASK, and COMMIT decisions. Only DRAFT is followed by an anticipated-user reply, which can guide internal revision before the draft is shown.

C.1 Assistant Prompts for CTOM and Its Ablations

  • The simulator sees only the shown draft, an ASK question, or the last shown draft at COMMIT, not plans, internal drafts, or predictions.Revisions rely on the earlier draft and the assistant’s own predicted reply remaining in context.Source: §C.1 · p. 19
  • The math prompt treats the human as able to judge correctness but not solve the problem for the assistant.A draft accepted by the anticipated reply is sent to the human as the answer.Source: §C.1 · p. 20
  • After a draft, the assistant predicts the user’s likely response; if it identifies a problem, the assistant gets another turn to improve the draft before the user sees it.The categorical addendum reserves the anticipated-user slot for DRAFT; ASK receives no anticipated reply.Source: §C.1 · pp. 19, 20, 22
  • The coding prompt supplies the exact function signature but withholds behavioral requirements, while the writing prompt makes facts, tone, and length user-specific.The simulator receives access to the private task information.Source: §C.1 · pp. 20–21
  • ASK is for missing information that could materially change the solution; an anticipated reply is a prediction, not confirmed information.The prompt directs the assistant to correct known implementation mistakes and seek clarification for unknown requirements.Source: §C.1 · p. 21
Loading 2610.12872v1…