Exploring Next Token Prediction in Theory of Mind (ToM) Tasks: Comparative Experiments with GPT-2 and LLaMA-2 AI Models

TL;DR

The paper investigates whether next-token prediction can approximate Theory of Mind reasoning in autoregressive models by comparing GPT-2 and Llama-2-7b-hf on ToM tasks. It uses 10 ExploreToM stories augmented with GPT-4 infills across four temperature settings and five infill levels to examine zero-, first-, and second-order reasoning. The results show that increasing contextual infill generally reduces prediction accuracy, while architecture and temperature modulate confidence and output diversity; Llama-2-7b-hf typically outperforms GPT-2, especially on higher-order reasoning. The findings illuminate how model architecture, context length, and stochasticity interact to shape ToM-like reasoning in LLMs and guide future evaluations and robustness improvements in context-aware AI systems.

Abstract

Language models have made significant progress in generating coherent text and predicting next tokens based on input prompts. This study compares the next-token prediction performance of two well-known models: OpenAI's GPT-2 and Meta's Llama-2-7b-chat-hf on Theory of Mind (ToM) tasks. To evaluate their capabilities, we built a dataset from 10 short stories sourced from the Explore ToM Dataset. We enhanced these stories by programmatically inserting additional sentences (infills) using GPT-4, creating variations that introduce different levels of contextual complexity. This setup enables analysis of how increasing context affects model performance. We tested both models under four temperature settings (0.01, 0.5, 1.0, 2.0) and evaluated their ability to predict the next token across three reasoning levels. Zero-order reasoning involves tracking the state, either current (ground truth) or past (memory). First-order reasoning concerns understanding another's mental state (e.g., "Does Anne know the apple is salted?"). Second-order reasoning adds recursion (e.g., "Does Anne think that Charles knows the apple is salted?"). Our results show that adding more infill sentences slightly reduces prediction accuracy, as added context increases complexity and ambiguity. Llama-2 consistently outperforms GPT-2 in prediction accuracy, especially at lower temperatures, demonstrating greater confidence in selecting the most probable token. As reasoning complexity rises, model responses diverge more. Notably, GPT-2 and Llama-2 display greater variability in predictions during first- and second-order reasoning tasks. These findings illustrate how model architecture, temperature, and contextual complexity influence next-token prediction, contributing to a better understanding of the strengths and limitations of current language models.

Figures (2)

• Figure 1: Prediction results of the GPT-2 model for different temperature settings and infill variations.
• Figure 2: Prediction results of the Llama-2 model for different temperature settings and infill variations.