Targeted Retrieval, Compact Representations: How CoT Reasoning Improves Long-Context Counting
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.
Related
- Counting Representations of Quadratic Forms
- INFUSER: Influence-Guided Self-Evolution Improves Reasoning
- Counting and Covering in Nearest-Neighbour Representations of Boolean Functions
- Collaborative representations for targeted causal inference under outcome selection
- Recursive Reasoning Exploration: Mathematical Research through Evolving Reasoning Graphs : case study: An Exact Certificate-Degree Theorem for Two-Term Tensor Decompositions : Phase III of the Math Instruments Program
BibTeX
@misc{shan2026targeted,
title = {{Targeted Retrieval, Compact Representations: How CoT Reasoning Improves Long-Context Counting}},
author = {Liang Twist Shan and Tianyu Hu and Hao Yan and Yiqiao Zhong},
howpublished = {arXiv},
year = {2026},
month = sep,
doi = {10.48550/arxiv.2609.38958},
eprint = {2609.38958},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2609.38958},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Targeted Retrieval, Compact Representations: How CoT Reasoning Improves Long-Context Counting
AU - Liang Twist Shan
AU - Tianyu Hu
AU - Hao Yan
AU - Yiqiao Zhong
PY - 2026
DA - 2026/09/30
PB - arXiv
DO - 10.48550/arxiv.2609.38958
UR - https://arxiv.org/abs/2609.38958
AB - 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.
ER -