Chain of thought reasoning improves counting in long texts
Targeted Retrieval, Compact Representations: How CoT Reasoning Improves Long-Context Counting
Artificial IntelligenceComputation and LanguageMachine Learning
Summary
Counting items in long texts is hard for large language models, but reasoning step-by-step helps. The researchers showed that when models think through their process (chain-of-thought), they can focus better on each item and keep a mental tally, even if they don’t explicitly number things. They discovered two ways models find items: either by broadly looking at many at once or by focusing on one at a time while updating their count. This understanding might help improve how AI reads and processes large amounts of information.
What this means in practice
- •For natural language processing engineers: Improve long-document analysis by implementing chain-of-thought methods to increase counting accuracy and focused information retrieval.
- •For ai system designers: Design language models with internal state tracking and targeted retrieval to better handle dispersed data in extended contexts.
Authors
Liang Twist Shan, Tianyu Hu, Hao Yan, Yiqiao Zhong
Abstract
Large language models (LLMs) have been rapidly improving in long-context tasks, powered by Chain-of-Thought (CoT) reasoning. However, the internal mechanisms underlying this improvement remain unclear. We investigate these mechanisms through a needle-in-a-haystack (NIAH) counting task, where an LLM is asked to count the number of records dispersed in a long text. Across twelve model comparison groups, Thinking (or reasoning) improves counting accuracy over Non-thinking, with pronounced gains at larger counts. This motivates our mechanistic analysis, which identifies two contrasting mechanisms: (i) broad retrieval, where Non-thinking models broadly attend to multiple needles; (ii) targeted retrieval, where Thinking models use enumeration in CoT traces to successively retrieve needles. Targeted retrieval concentrates attention on individual needles and is accompanied by more compact internal representations. Moreover, causal intervention analysis suggests that Thinking models use the CoT trace to maintain and update an internal counter as needles are successively retrieved, even without explicit numbering. In small controlled experiments, both retrieval mechanisms and counter states emerge under standard autoregressive training. Together, our results connect long-context retrieval with representation geometry of counting, supporting a state-tracking account of CoT reasoning.