Mechanistic Interpretability of Large-Scale Counting in LLMs through a System-2 Strategy

TL;DR

This paper introduces a System-2 inspired strategy that decomposes large counting tasks into smaller parts, enabling large language models to overcome architectural limitations and improve counting accuracy.

Contribution

It proposes a simple test-time decomposition method inspired by System-2 processes, with mechanistic analysis showing how it enhances counting in LLMs.

Findings

The strategy improves counting accuracy on large-scale tasks.

Latent counts are stored and transferred via dedicated attention heads.

Experimental analysis confirms the mechanism's effectiveness.

Abstract

Large language models (LLMs), despite strong performance on complex mathematical problems, exhibit systematic limitations in counting tasks. This issue arises from the architectural limits of transformers, where counting is performed across layers, leading to degraded precision for larger counting problems due to depth constraints. To address this limitation, we propose a simple test-time strategy inspired by System-2 cognitive processes that decomposes large counting tasks into smaller, independent sub-problems that the model can reliably solve. We evaluate this approach using observational and causal mediation analyses to understand the underlying mechanism of this System-2-like strategy. Our mechanistic analysis identifies key components: latent counts are computed and stored in the final item representations of each part, transferred to intermediate steps via dedicated attention heads, and aggregated in the final stage to produce the total count. Experimental results demonstrate that this strategy enables LLMs to surpass architectural limitations and achieve higher accuracy on large-scale counting tasks. This work provides mechanistic insight into System-2 counting in LLMs and presents a generalizable approach for improving and understanding their reasoning behavior.