Brittle Minds, Fixable Activations: Understanding Belief Representations in Language Models (2406.17513v3)

Abstract: Despite growing interest in Theory of Mind (ToM) tasks for evaluating LMs, little is known about how LMs internally represent mental states of self and others. Understanding these internal mechanisms is critical - not only to move beyond surface-level performance, but also for model alignment and safety, where subtle misattributions of mental states may go undetected in generated outputs. In this work, we present the first systematic investigation of belief representations in LMs by probing models across different scales, training regimens, and prompts - using control tasks to rule out confounds. Our experiments provide evidence that both model size and fine-tuning substantially improve LMs' internal representations of others' beliefs, which are structured - not mere by-products of spurious correlations - yet brittle to prompt variations. Crucially, we show that these representations can be strengthened: targeted edits to model activations can correct wrong ToM inferences.

• The paper demonstrates that larger models and fine-tuning methods, like instruction-tuning and RLHF, significantly boost belief representation accuracy.
• It reveals that mental state representations are sensitive to prompt variations, highlighting the need for careful prompt engineering.
• Activation steering with Contrastive Activation Addition markedly enhances ToM performance, achieving accuracy levels near 90%.

Benchmarking Mental State Representations in LLMs: A Professional Overview

The paper "Benchmarking Mental State Representations in LLMs" by Matteo Bortoletto, Constantin Ruhdorfer, Lei Shi, and Andreas Bulling presents an extensive investigation into the internal representations of mental states, specifically beliefs, within various LLMs (LMs). This research addresses the emergent area of evaluating Theory of Mind (ToM) capabilities in LMs, which is pivotal for enhancing human-computer interaction through improved understanding of human mental states.

Summary of Contributions

The research benchmarks mental state representations across a range of LM families, model sizes, fine-tuning approaches, and prompt variations. The primary contributions of the paper include:

1. Probing Experiments Across Diverse LMs:
   • The paper encompasses both base and fine-tuned versions of the Pythia and Llama-2 models, ranging from 70 million to 70 billion parameters.
   • Probing experiments reveal that accuracy in representing others' beliefs increases with model size, and fine-tuning, particularly with instruction-tuning and RLHF, significantly enhances the quality of these representations.
2. Sensitivity to Prompt Variations:
   • The authors explore the robustness of LM representations to different types of prompt variations, including random token insertion, misleading prompts, time specification, and initial belief revelation.
   • It is observed that while oracle representations remain robust, representations of beliefs from another’s perspective are sensitive to prompt variations, with significant improvements when initial beliefs are included in the prompt.
3. Investigation of Memorization:
   • Probing accuracy is analyzed with dimensionality reduction to identify potential memorization effects. Probes trained on principal components of the activation data demonstrated that high accuracy can be retained with significantly fewer parameters, indicating low memorization issues.
4. Enhancing Reasoning via Activation Steering:
   • Building on previous work on activation editing, the paper demonstrates improvements in LMs' ToM performance without probe training using Contrastive Activation Addition (CAA).
   • Results showcase substantial performance gains across forward belief, forward action, and backward belief tasks, often reaching or exceeding 90% accuracy.

Analytical Insights and Implications

Probing Experiments and Fine-tuning

The probing experiments indicate a crucial aspect of model design and training: larger models inherently develop richer and more accurate internal representations of mental states. This is consistent with broader trends in NLP showing that scaling up model parameters correlates with improved performance across various tasks. The effectiveness of fine-tuning, particularly through methods involving human feedback, suggests that exposure to diverse and informative data during training is pivotal for nurturing ToM reasoning capabilities.

Sensitivity to Prompts

The sensitivity of belief representations to prompt variations underscores the necessity for careful prompt engineering when deploying LMs in real-world applications. This finding implies that achieving robust mental state understanding requires not only advanced model architectures and training regimes but also carefully designed prompts that can mitigate ambiguity and bias in generated responses.

Memorization Concerns

Addressing memorization in probing provides reassurance regarding the interpretability of probe outcomes. The absence of strong memorization effects, as evidenced by the successful use of principal components, suggests that the probes offer genuine insights into the models’ internal representations without overfitting to the training data.

The Role of Activation Steering

The success of CAA in improving reasoning performance without dedicated probe training marks a notable advancement. This technique's efficacy in steering model activations points to a practical and computationally efficient method for enhancing model performance on specific tasks, particularly those involving nuanced understanding like ToM. This could be crucial for future work aimed at achieving highly adaptable and generalizable AI systems.

Future Speculations and Theoretical Implications

The findings present pathways for future research and development. One critical direction is the extension of these benchmarks to new data sets and model architectures, given the rapid evolution of LMs. Exploring the intersection of ToM capabilities with other cognitive faculties like emotional intelligence and moral reasoning could provide a holistic framework for AI understanding of human mental states.

Furthermore, the implications of fine-tuning methodologies in shaping ToM capabilities suggest that future models could benefit from more sophisticated and specialized instruction-tuning regimes that encompass broader and deeper aspects of human cognition. This could lead to AI systems that not only understand beliefs and emotions but can also navigate complex social interactions with a higher degree of empathy and contextual awareness.

In conclusion, this paper provides extensive benchmarks and insights into how LMs represent mental states. The robustness analysis, fine-tuning impacts, and activation steering advancements contribute to a deeper understanding of the internal mechanisms that underpin ToM capabilities in AI, highlighting both current capabilities and future potential in this vital area of research.