Language models improve teamwork by predicting user feedback
Collaborative Theory of Mind: Learning to Self-Revise by Anticipating User Responses
Artificial Intelligence
Summary
Large language models often have trouble working well with people over many back-and-forth exchanges, like jumping to answers too soon or not fixing mistakes after being told. The authors propose a way for these models to guess how users might respond to their answers before showing them. They train the models to adjust their replies based on these guesses to better match what users need. This method helps the models perform tasks like math thinking, coding, and writing more accurately and with fewer back-and-forth replies.
What this means in practice
- •For software developers: Improve coding assistants that adapt their suggestions by anticipating developer feedback during multi-turn interactions.
- •For technical writers: Enhance writing support tools that revise drafts internally by predicting how users might critique or request changes.
Authors
Fanzhi Zeng, Anubrata Das, Maytal Saar-Tsechansky, Liu Leqi
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.