How do LLMs count with “thinking” mode on or off?

01 / MOTIVATION & TASK

Counting as a test of retrieval and state tracking.

Long-context reasoning is important for LLM agents working with long documents or on multi-step tasks. We want to understand how models retrieve relevant information and keep track of what they have already processed. Therefore, we use counting as a testbed: the model must find every relevant record without skipping or counting any record twice.

In our Needle-in-a-Haystack task, city-score audit records are scattered throughout a passage of unrelated text, and the model is asked to count them. This separates the difficulty of finding and tracking records from the difficulty of interpreting their meaning. We know both the correct total and which records the model should account for, so we can examine where counting goes wrong.

A LONG PASSAGE ILLUSTRATIVE EXCERPT

… the report described changes to the transport network …

N₁ In the 2024 city score audit, Paris received a score of 92.

… several thousand words of unrelated text …

N₂ In the 2024 city score audit, Vienna received a score of 73.

QUESTION How many audit records are there? 2

We find that a model achieves substantially higher counting accuracy with chain-of-thought reasoning enabled than with it disabled. We want to understand what changes inside the model when CoT reasoning is enabled. Does it change which records are retrieved? Does the model represent how many records it has already processed more clearly? And does that information help it retrieve the next record and produce the final count? We compare Non-thinking, which gives a direct answer, with Thinking, which generates a reasoning trace before answering.

In what follows, the running index k denotes progress through records in the prompt for Non-thinking or through items in the reasoning trace for Thinking. The final count N is the total number of records.

We compare counting accuracy, examine hidden-state representations, and intervene on attention heads and states to test how they affect the answer. We also train small models on a synthetic counting task to study how retrieval develops. We find two retrieval pathways. In Non-thinking, the model attends to multiple records at the answer query, which we call broad retrieval. In Thinking, it retrieves individual records in sequence, which we call targeted retrieval.

02 / BEHAVIOR & EMPIRICAL LAWS

Accuracy depends on both count and context length.

Across twelve model comparison groups, Thinking improves exact counting accuracy by 28–63 percentage points, averaged over the evaluation grid spanning context lengths from 1k to 20k tokens. At a context length of 1k tokens, median Non-thinking accuracy falls from 100% for three records to 30% for ten and 0% for twenty. Thinking remains accurate over a wider count range, although it also deteriorates as count and context length increase.

Models evaluated 12 comparison groups · Table 1

Parsed exact accuracy (%) for N = 5 over 30 paired seeds. Transposed from Table 1.
Model / comparison group	1k tokens		20k tokens
Model / comparison group	Non-thinking	Thinking	Non-thinking	Thinking
Qwen3-4B	43.3	100	3.3	90
Qwen3-8B	60	100	30	96.7
Qwen3-14B	50	100	30	100
Qwen3-32B	96.7	100	46.7	100
Gemma-4-E4B	36.7	100	33.3	43.3
Gemma-4-12B	73.3	100	10	83.3
Gemma-4-26B-A4B	100	100	43.3	93.3
Gemma-4-31B	100	100	63.3	100
NVIDIA-Nemotron-Nano-9B-v2	40	100	3.3	33.3
NVIDIA-Nemotron-3-Nano-4B	46.7	96.7	13.3	73.3
GLM-4-9B-0414 / GLM-Z1-9B-0414 †	20	100	13.3	46.7
Ministral 3 8B Instruct / Reasoning †	73.3	96.7	16.7	70

† The GLM and Ministral groups use separate checkpoints, listed as Non-thinking / Thinking. All other groups switch modes within the same checkpoint. Qwen3, Gemma-4, and Nemotron-3-Nano-4B use the thinking flag; Nemotron-Nano-9B-v2 uses a system instruction. Model configurations are described in Appendix C.1.

Accuracy across models and context lengths

counting accuracy as record count and context length increase — **Counting accuracy across models.** A–B show the median and middle 50% across twelve comparison groups. C–D extend evaluation to 25k–100k tokens for two larger models: rose solid lines denote Non-thinking and blue dashed lines denote Thinking. Each condition uses 30 paired seeds. The left curves are smoothed visual guides; the right curves connect observations. Section 2.

A heuristic view of counting errors

Why does accuracy decline as the number of records grows? We consider two simple approximations: an increasingly imprecise estimate of the total, and an increasing number of opportunities to make a mistake.

Error models, equations, and assumptions

For Non-thinking, imagine a noisy estimate centered on the true count N. Rounding gives the correct integer only when the estimate is within half a count of N. If the standard deviation grows in proportion to N, fewer estimates fall inside this fixed interval as N increases. We model this with Gaussian error whose standard deviation is Nσ(1, L), where L is context length and σ(1, L) is a fitted noise scale in count units.

For Thinking, imagine processing the records one at a time. Even if each step is usually correct, a longer sequence offers more opportunities for failure. If each of N required steps succeeds independently with probability 1 − γ_L, the probability that all steps succeed is (1 − γ_L)^N. A factor of 1 − α accounts for an additional baseline failure probability.

NON-THINKING

p ≈ 2Φ(1 / [2Nσ(1, L)]) − 1

THINKING

p ≈ (1 − α)(1 − γ_L)^N

Here p is exact counting accuracy and Φ is the standard-normal cumulative distribution. These are descriptive approximations, not established accounts of the model’s internal computation. The first does not imply that activation noise itself grows with N. The second assumes independent step errors that do not cancel; γ_L is an effective fitted parameter, not a directly measured retrieval error rate.

descriptive fits of counting accuracy — **Empirical accuracy laws.** Circles and triangles are measured median accuracies; solid and dashed curves fit the Non-thinking and Thinking models. Colors identify context length. The paper also analyzes how count representations become harder to distinguish as count increases, including the roles of noise and separation between adjacent counts (Appendix C.2). Appendix C.

03 / FEATURE REPRESENTATIONS

Representations offer a glimpse of the counter.

Trace parsing, linear probes, and steering

Thinking traces vary considerably in syntax: records may appear in numbered lists, prose descriptions, or recaps, and enumeration may restart. To compare progress states across these formats, we parse the traces into ordered record items and identify their endpoints. Counting is often expressed in natural language or through cumbersome enumeration rather than a concise, consistent procedure, suggesting room to improve how Thinking organizes and tracks progress.

A natural starting point is to ask whether hidden states encode a running index. We test both whether a linear model can predict the index and whether states with different indices can be distinguished from one another. At Non-thinking record endpoints, held-out ridge regression achieves R² = 0.910 for Qwen and 0.827 for Gemma. Yet balanced nearest-centroid accuracy at the same selected layers is only 46% and 40%. Although the class means are ordered by count, individual states from different counts overlap substantially.

This led us to examine the geometry rather than interpret a high regression score as evidence for a single count axis. PCA reveals curved configurations of count-dependent states. The first three components explain 98.1% and 97.5% of the variance of the ten class centroids. An ordered, low-dimensional configuration need not be a straight line.

When linear steering fails

Like many researchers studying model representations, we began by looking for a direction that encodes count. We then tried adding or subtracting this direction to make the model report a larger or smaller count—a linear steering intervention. Despite repeated attempts, we could not find a reliable steering direction: adding or subtracting the fitted direction did not consistently increase or decrease the reported count. One possible reason is that count representations follow a curved geometry rather than a straight line. A linear readout in a high-dimensional space can achieve a high R² while obscuring this structure, so successful decoding need not imply that moving along the fitted direction produces a meaningful change in count.

Thinking produces more distinguishable progress states

In a selected paired Qwen3-8B comparison, completed Thinking-item states form more compact running-index groups than Non-thinking record endpoints. The interactive plot below shows these representations.

Each small dot is a snapshot of the model’s hidden state at one point in a counting task. On the left, we take that snapshot at the end of a record in the prompt; on the right, at the end of an item in the reasoning trace. Color marks the running index: for example, dots colored for index 3 come from the third record or completed trace item. Dots of the same color come from different inputs. PCA projects these high-dimensional states into three dimensions so we can view their arrangement.

COUNT-STATE GEOMETRYREAL PROJECTED STATES

Running-index representations on paired inputs

Qwen3-8B · 10k tokens · N = 10 · secondary confirmation subset

Non-thinking L13 · prompt endpoints46% NCC

Thinking L31 · completed CoT items98% NCC

EXPLORE BOTH VIEWS

RUNNING COUNT

Loading the saved geometry…

How classification accuracy is measured

Contrasting representations. Each point shows a hidden state after a prompt record (Non-thinking) or a reasoning-trace item (Thinking). Colors indicate the running index; numbered points mark the average state for each index. To measure how easily the indices can be distinguished, nearest-centroid classification (NCC) predicts each state’s index from the closest group average. On held-out inputs—inputs not used to fit the classifier—accuracy is 46% for Non-thinking and 98% for Thinking. These scores use 16 PCA components, scaled to equal variance (whitening), rather than only the three dimensions shown here. Both views use Qwen3-8B with 10k-token inputs. See Appendix A for details.

Look at how tightly the dots of each color group together, and how much different colors overlap. In this example, Thinking states form tighter groups by running index, making progress easier to distinguish across inputs.

Additional PCA analyses in the paper

In our paper, we extend this PCA analysis beyond the paired example above. We compare running-index representations across the full set of parsed trace formats, examine final-count representations at the answer query, and test whether the geometry transfers across city, flower, and animal records. Across these analyses, count representations are more compact in Thinking and less sensitive to the topic of the records, which indicates a noise reduction process within the thinking trace. See Appendix E.2 for the full readout comparisons and PCA visualizations.

04 / TWO CAUSAL PATHWAYS

Causal interventions distinguish two different pathways.

The following experiments use Qwen3-8B and Gemma-4-E4B on 10k-token passages. Results are reported in that order. Head selection uses fitting data and is frozen for held-out interventions.

At the Non-thinking answer query, selected heads distribute attention across multiple records. During Thinking, successive trace queries preferentially attend to successive records. We then ablate these heads and patch hidden states to test how they affect counting.

Attention patterns and retrieval-head definitions

broad retrieval over prompt records — **Non-thinking.** At the answer query, selected heads attend broadly across records. Rows: heads; columns: records. Appendix D.2.

targeted retrieval during enumeration — **Thinking.** Successive trace queries target successive records. Marker k precedes item k + 1; marker 0 is initialization. Appendix E.2.

Defining the two types of retrieval heads

We score heads by how much attention they assign to records and how that attention is distributed. One score identifies heads that attend to many records at the answer query; the other identifies heads that attend to the next record during the trace.

Broad-retrieval heads

A broad-retrieval head assigns substantial attention to records and spreads that attention across them. Its score combines the total record-directed attention M_h with effective coverage:

B_h = M_h exp(H_h) / N

H_h is the entropy of attention shares across records. If attention is uniform over k of N records, effective coverage is k/N. A head focused on just one record therefore receives a lower score than one with the same total attention spread across all records.

Broad-retrieval head scores across layers in Qwen and Gemma — Each point is a head. Highlighted: Qwen Top-32 and Gemma Top-6, selected for causal tests. Scores are averaged over prompts with 2–10 records; Gemma includes global-attention layers only. Appendix D.2.

Targeted-retrieval heads

A targeted-retrieval head attends to the prompt record corresponding to the next trace item. At the query q_k before item k, its score is the attention mass assigned to that record’s token span S_k:

T_h(q_k, S_k) = ∑_{s ∈ S_k} A_h(q_k, s)

A_h(q, s) denotes attention from query q to token s. We average over parsed trace items within each seed, then across seeds. A high score indicates reliable attention to the next record as the trace progresses.

Targeted-retrieval head scores across layers in Qwen and Gemma — Each point is a head. Highlighted: Qwen Top-128 and Gemma Top-6, selected for causal tests. Gemma includes global-attention layers only. Appendix E.2.

THINKING OFF

Form → retrieve → consolidate

In Non-thinking, early record representations contribute information that is retrieved at the answer query and consolidated into a final answer state.

Original mechanism schematic · Section 3.2

Non-thinking form–retrieve–consolidate mechanism — **The Non-thinking pathway.** Record states contribute to retrieval at the answer position, followed by consolidation in deeper layers. Section 3.2 and Appendix D.3.

Causal evidence: record states, retrieval heads, and answer states

01 / FORM

Record-state contributions are strongest in early layers.

We first test at which layers record states have the greatest effect on the final count. We corrupt inputs by replacing records with length-matched ordinary text, then patch clean record-span states from a source run into the corrupted target. Reverse patching inserts corrupted states into a clean target. We measure restoration by the decrease in normalized count error and reverse patching by its increase.

Both interventions have effects of approximately 47–50% at the first layer and little effect at late layers. Record states therefore have the largest effect on counting when patched in early layers.

record-span patching by layer — Clean-to-corrupted restoration and corrupted-to-clean patching of record spans, measured by changes in normalized count error. Appendix D.3.

02 / RETRIEVE

Broad-retrieval heads contribute to the reported count.

To test whether these heads affect the count, we zero the outputs of 32 selected Qwen heads or six selected Gemma heads at the answer query and compare with layer-matched random head banks.

The normalized absolute count change is 54.5% and 24.7%, compared with 4.0% and 6.0% for random controls. The selected heads contribute more strongly to the answer than these controls.

broad-retrieval head ablation — Retrieval-head ablation versus layer-matched random and original controls. Effect: absolute count change divided by the true count. Appendix D.3.

03 / CONSOLIDATE

Deep answer-state transfer changes the predicted count.

We test whether replacing the answer state with one from a different-count prompt makes the answer follow the source count. We patch answer-query states from a source prompt into a target prompt with a different count, one layer at a time, using pairs with correct baseline answers.

At the final layer, the patched answer matches the source count in 95.8% and 100% of cases. This supports an answer-readout role for the transferred state.

source-count match rate after answer-state transfer — Proportion of predictions matching the source count after answer-state patching. Appendix D.3.

THINKING ON

Retrieve → update → read

In Thinking, the model retrieves records one by one. The state after each item helps guide the next retrieval, and the model uses the completed trace to produce the final count.

Original mechanism schematic · Section 4.2

Thinking retrieve–update–read mechanism — **The Thinking pathway.** N_k is a prompt record, q_k the query before retrieving it, and H_k the state after completing its trace item. The previous item state guides the next query. Section 4.2 and Appendix E.3.

Causal evidence: next-record retrieval, progress, and final count

01 / RETRIEVE

Targeted-retrieval heads support trace continuation.

To test whether retrieving the next record depends on these targeted heads, we ablate them starting at the selected trace marker and continuing at every subsequent token. This increases the next-item failure rate by 60% and 80%, compared with 10% and 6.7% for random heads, across ten confirmation transitions per model.

Final-count failure increases by 100 percentage points in both models. The intervention establishes a contribution of the selected heads to continuing the trace and completing the count.

targeted-retrieval head ablation effects — Retrieval-head ablation versus layer-matched random and original controls: next-item and final-count failure rates. Qualitative inspection also suggests compensatory retrieval in Qwen’s unablated later heads. Appendix E.3.

02 / UPDATE

Completed-item states influence subsequent retrieval.

We test whether replacing the state after one item with a state from one item further ahead makes the model skip the next record. To avoid explicit numbering as a cue, we use unnumbered traces for both Qwen and Gemma. We remove non-record filler to align source and target token positions.

We patch the full state span of source item k into target item k−1, aligning the ends of the two item spans. We then test whether the next generated record is k+1 and whether enumeration continues from there without another patch.

The model skips the next record in 24/30 and 21/30 trials, compared with 0/30 and 2/30 when patched with its own states. After a successful skip, it usually continues from the shifted position without another patch.

completed-item state transfer and subsequent retrieval — Counter-state patching from a later trace item to an earlier one, without explicit indices. Later columns condition on preceding success and a remaining item. Appendix E.3.

We also tried patching in the reverse direction, transferring an earlier item state to a later position. However, repeating a record prompted the model to reconsider its enumeration, making it difficult to determine whether the patched state directly guided the next retrieval.

03 / READ

The model uses the completed trace to produce the final count.

We test whether the model still needs the original records once the trace is complete. We hold the generated trace fixed and separately zero out prompt-record and trace representations while preserving token positions. Removing prompt-record states leaves exact-count accuracy unchanged at 97% and 70%; removing trace states reduces it to 1% and 15%. Once the trace is complete, the model relies mainly on it to produce the final count.

We then patch answer-query states between different-count prompts with correct baseline answers, one layer at a time. At the last layer, the patched answer matches the source count in 97.5% and 90% of the 40 shared directed prompt pairs.

source-count match rate across intervention layers — Proportion of predictions matching the source count after answer-state patching. Appendix E.3.

Together, these experiments support two counting pathways. Non-thinking retrieves broadly and consolidates record information into a final count. Thinking retrieves records sequentially, tracks progress without explicit numbering, and reads the final count from the completed trace.

05 / TRAINING DYNAMICS

Retrieval heads emerge
during training.

Training setup and results

Can the two retrieval strategies emerge through autoregressive training? We train two Transformers from scratch, each with four layers, eight heads per layer, and approximately 12.7 million parameters. Inputs contain 256 shuffled Shakespeare characters and a query specifying three target characters, with one to ten matching occurrences.

The Non-thinking model reports the count directly. The Thinking model first enumerates matching characters in source order, without explicit indices. Architecture, input distribution, training budget, and final-count loss weight are matched; the two modes are separately trained models.

UNSHUFFLED ILLUSTRATION · TARGET CHARACTERS: T, b, n

<BOS> <CountChar> T b n <Sep>

… To be, … or not to be …

Non-thinking:<Ans> <4> <EOS>

Thinking:<Think> <Sep>T <Sep>b <Sep>n <Sep>b </Think><Ans> <4> <EOS>

Because RoPE and the regular synthetic format may let the models exploit token position as a shortcut for counting, counter readouts are less informative here; we focus instead on the emergence of retrieval heads.

Learning to retrieve

Both models initially predict the full sequence. At step 1,500, we switch the loss to task outputs: the final count for Non-thinking, and the enumeration trace plus final count for Thinking.

The accuracy gap widens later in training. By step 10,000, Thinking reaches 95% count accuracy, while Non-thinking plateaus at 21%. We track individual heads throughout training to examine how their retrieval patterns develop.

free-generation count accuracy during training — **Counting accuracy during training.** Evaluated on the same 100 held-out inputs. The dashed line marks the objective switch. Section 5 and Appendix G.4.

Broad and targeted retrieval emerge in different layers

Using the same retrieval scores as in the pretrained models, we track how attention changes during training. Broad-retrieval heads emerge in the Non-thinking model’s first layer, while targeted-retrieval heads emerge in the Thinking model’s last layer.

Layer-wise broad retrieval during Non-thinking training — **Non-thinking.** Broad retrieval emerges in the first layer. The dashed line marks the switch to task-output loss. Section 5.

Layer-wise targeted retrieval during Thinking training — **Thinking.** Targeted retrieval emerges in the last layer. The dashed line marks the switch to task-output loss. Section 5.

Implications for data and training

These results suggest that enumeration traces can shape the retrieval strategy a model learns, even without explicit numbering. Training on such traces with a standard autoregressive objective produces targeted retrieval and higher counting accuracy in this setting, without specialized RL post-training. This points to the potential of high-quality CoT data to improve reasoning and support the development of broader capabilities.

06 / FUTURE WORK

Future directions

We still do not know exactly how the model updates its progress or whether the same mechanisms operate in more complex tasks.

Linking geometry to causal control
Count representations may follow a curved geometry. Geometry-aware interventions could clarify how these representations support count updates.
Identifying the update circuit
State transfers support a role for completed-item states in guiding retrieval, but transfer record content and count information together. Identifying the circuit that updates progress when a record is retrieved would provide more direct evidence for the proposed counter mechanism.
Understanding rethinking
What makes a model reconsider an earlier step, and when does that improve its answer? We would like to understand how it detects a problem and changes its subsequent reasoning.
Extending to complex reasoning
Do similar retrieval and state-tracking mechanisms operate when reasoning involves search, branching, or tracking several changing states?

07 / REFERENCES

References and related work

Wei et al. (2022). Chain-of-Thought Prompting Elicits Reasoning in Large Language Models. Reasoning traces improve performance on arithmetic, symbolic, and commonsense tasks.
Hsieh et al. (2024). RULER: What’s the Real Context Size of Your Long-Context Language Models? Extends needle retrieval to controlled long-context tasks involving multiple records, tracing, and aggregation.
Wu et al. (2025). Retrieval Head Mechanistically Explains Long-Context Factuality. Identifies attention heads that retrieve information from context and tests their contribution through ablation.
Behrens, Biggio, and Zdeborová (2025). Counting in Small Transformers: The Delicate Interplay between Attention and Feed-Forward Layers. Characterizes distinct counting strategies and how attention and feed-forward layers implement them in small transformers.
Engels et al. (2025). Not All Language Model Features Are One-Dimensionally Linear. Demonstrates circular, multidimensional features and tests their causal role, motivating analysis beyond a single linear direction.
Gurnee et al. (2025). When Models Manipulate Manifolds: The Geometry of a Counting Task. Studies curved count representations and the attention computations that construct and use them during text linebreaking.
Heimersheim and Nanda (2024). How to Use and Interpret Activation Patching. Clarifies how prompt choice, intervention direction, and evaluation metrics affect the interpretation of patching results.
McGrath et al. (2023). The Hydra Effect: Emergent Self-repair in Language Model Computations. Documents compensation after ablation, showing why the measured effect of removal can understate a component’s contribution.
Cabannes et al. (2024). Iteration Head: A Mechanistic Study of Chain-of-Thought. Tracks how attention mechanisms for iterative computation emerge during controlled training.
Li, Guo, and Andreas (2025). (How) Do Language Models Track State? Identifies alternative state-tracking mechanisms in permutation tasks and relates their emergence to training and generalization.

How to cite

The paper can be cited as follows.

Bibliography

Liang Twist Shan, Tianyu Hu, Hao Yan, and Yiqiao Zhong. “Targeted Retrieval, Compact Representations: How CoT Reasoning Improves Long-Context Counting.” arXiv preprint arXiv:2609.38958, 2026.

BibTeX

@misc{shan2026targetedretrieval,
  title = {Targeted Retrieval, Compact Representations:
           How {CoT} Reasoning Improves Long-Context Counting},
  author = {Shan, Liang Twist and Hu, Tianyu and
            Yan, Hao and Zhong, Yiqiao},
  year = {2026},
  eprint = {2609.38958},
  archivePrefix = {arXiv},
  primaryClass = {cs.AI},
  url = {https://arxiv.org/abs/2609.38958}
}