TLDR: A study adapted a neuroscientific “contrast localizer” to find crucial units for Theory of Mind (social) and mathematical tasks in LLMs and VLMs. Surprisingly, low-activation units sometimes caused bigger performance drops than high-activation ones, and math-related units often impacted social tasks more than social-specific units. This suggests current contrast localizers may not accurately pinpoint task-specific causal units in AI, highlighting a need for broader testing and deeper analysis of model internals.
A recent research paper explores how large language models (LLMs) and vision-language models (VLMs) process complex tasks like understanding social cues (Theory of Mind) and solving mathematical problems. The study, titled “Evaluating Contrast Localizer for Identifying Causal Units in Social & Mathematical Tasks in Language Models,” investigates whether a neuroscientific technique, known as a contrast localizer, can effectively pinpoint the specific internal components within these AI models that are crucial for these cognitive abilities.
The concept of a “contrast localizer” comes from neuroscience, where it’s used to identify specific brain regions that become active during certain cognitive tasks. For instance, to find language-selective areas in the human brain, researchers might compare brain activity when someone is listening to language versus non-language sounds. This paper adapts this idea to AI models, aiming to identify “causal units” – parts of the model whose activity is directly responsible for a particular function.
The researchers, Yassine Jamaa, Badr AlKhamissi, Satrajit Ghosh, and Martin Schrimpf, conducted extensive experiments across 11 LLMs and 5 VLMs, ranging significantly in size from 3 billion to 90 billion parameters. For Theory of Mind (ToM) tasks, they used stories designed to test understanding of false beliefs (e.g., a character believes something that isn’t true) as the “positive” condition, contrasted with stories that don’t require such social understanding (false photographs) as the “negative” condition. For mathematical reasoning, they contrasted “hard” arithmetic problems with “easy” ones.
After identifying units that showed high activation during these “positive” tasks, the team performed “ablations” – essentially turning off or setting to zero the activity of these selected units. They then observed the impact on the model’s performance on established ToM and mathematical benchmarks. This was compared against ablating units with low activation or randomly selected units, to see if the highly activated units were indeed causally important.
Also Read:
- Unpacking LLM Decisions: Why Human Risk Theories Fall Short with Linguistic Uncertainty
- Playing Detective with AI: Auditing Models for Covert Actions
Unexpected Findings Challenge Assumptions
The results of this study presented some surprising insights that challenge initial assumptions about how these contrast localizers work in AI models. Contrary to what might be expected, the researchers found that sometimes, units with low activation during a task, when ablated, led to larger performance drops than ablating the highly activated units. This suggests that simply being highly active doesn’t necessarily mean a unit is causally essential for a task.
Even more intriguingly, the study revealed a cross-task effect: units identified by the mathematical localizer often impaired the model’s performance on Theory of Mind tasks more significantly than units identified by the ToM localizer itself. This raises questions about the specificity of these identified units and whether they are truly specialized for one cognitive function over another, or if there’s a more general underlying mechanism at play.
These findings suggest that while the contrast localizer method is a promising neuroscientific approach, its direct application to LLMs and VLMs for identifying task-specific causal units might be more complex than anticipated. The authors emphasize the need for broader and more diverse stimulus sets for these localizers. For instance, the mathematical localizer could include more complex operations like division and multiplication, or even logic puzzles, beyond simple arithmetic. Similarly, the ToM localizer could benefit from a wider range of social scenarios.
Furthermore, the paper discusses the granularity of analysis. The current study focused on output activation units from each transformer block. However, the authors suggest that examining intermediate representations and hidden activations within the model’s architecture could provide a more comprehensive understanding of how complex reasoning processes are distributed and executed within these large AI models. This deeper dive could reveal more nuanced causal relationships.
This research contributes to the growing field of mechanistic interpretability, aiming to open the “black box” of large language models. While it highlights challenges in directly mapping high-level cognitive functions using current contrast localizer designs, it also points towards crucial directions for future research to better understand the internal workings of these powerful AI systems. You can read the full paper here: Evaluating Contrast Localizer for Identifying Causal Units in Social & Mathematical Tasks in Language Models.
