Source-linked AI summary

Targeted Retrieval, Compact Representations: How CoT Reasoning Improves Long-Context Counting

Liang Twist Shan, Tianyu Hu, Hao Yan, Yiqiao Zhong

arXiv:2609.38958v1cs.AIcs.CLcs.LGstat.AP

TL;DR

The paper asks how internal state tracking supports long-context retrieval and reasoning. It studies this question with a needle-in-a-haystack counting task and mechanistic analyses, finding that Thinking improves counting accuracy and uses targeted retrieval with compact progress representations. The results provide evidence for counter states that evolve during CoT retrieval, while direct identification of the update circuits remains future work.

  • Problem

    How internal state-tracking mechanisms support long-context retrieval and reasoning remains an open question.

    Source: §1 · p. 1
  • Method

    The study uses long-context needle-in-a-haystack counting and analyzes retrieval, count representations, causal interventions, and small autoregressively trained models.

    Source: §1 · pp. 2–3
  • Results

    Thinking substantially improves counting accuracy, and the analysis contrasts broad retrieval with noisy representations against targeted retrieval with compact progress representations.

    Source: §1 · p. 3
  • Takeaways & Limitations

    Causal interventions provide evidence that counter states guide continued enumeration, jointly carrying count information and record content.

    Source: §6 · p. 12

    The update circuits that implement counter-state changes as needles are retrieved remain unidentified, and the study uses counting as its testbed.

    Source: §6 · p. 12

Abstract

from arXiv · show

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.

1 Introduction

The paper uses long-context counting to study how retrieval and progress tracking interact, contrasting broad retrieval in Non-thinking models with targeted retrieval in Thinking models. It situates this analysis amid open questions about retrieval, counting representations, and CoT state tracking, and tests whether related mechanisms also emerge in small autoregressive models.

1 Introduction

  • Long-context counting asks models to count records scattered through a passage, combining retrieval and aggregation with explicit record and count labels for analyzing state tracking.The task adapts needle-in-a-haystack benchmarks and reduces semantic ambiguity compared with open-ended reasoning.Source: §1 · pp. 1–2
  • Thinking couples successive needle retrievals to compact progress representations, while Non-thinking broadly retrieves needles and forms noisy representations.The paper argues that targeted retrieval may help avoid omissions and double counting; causal interventions support counter states evolving in CoT traces.Source: §1 · p. 3
  • The paper compares counting accuracy across model families and reports that Thinking substantially improves accuracy over Non-thinking.The evaluation varies needle count and context length.Source: §1 · p. 3
  • Small transformers trained from scratch with or without CoT traces reproduce key aspects of the Thinking and Non-thinking mechanisms.The experiment tests whether these mechanisms emerge under autoregressive training.Source: §1 · p. 3

1.1 Related work

  • Prior work examines retrieval heads, counting geometry, and state tracking, but how retrieval mechanisms support task execution and connect to model behavior remains insufficiently explained.The paper combines geometric analysis and causal interventions to study this internal pathway.Source: §1.1 · pp. 3–4

2 Thinking vs. Non-thinking: performance across model families

Across 12 comparison groups, Thinking improves counting accuracy, especially at larger counts, while performance still declines when large counts coincide with longer contexts. The proposed explanation is that broad retrieval accumulates representational noise, whereas repeated retrieval steps can amplify errors as context length grows.

2 Thinking vs. Non-thinking: performance across model families

  • At L = 1k, median accuracy across 12 model groups falls from 100% at N = 3 to 30% at N = 10 and 0% at N = 20 for Non-thinking.These values are reported for Figure 2A.Source: §2 · p. 4
  • 28–63 percentage points: Thinking improves accuracy by this amount across 12 groups averaged over the 1k–20k grid, though accuracy drops at large counts combined with longer contexts.Figure 2 compares accuracy across counts and context lengths; its panels summarize group medians and interquartile ranges and paired-seed curves.Source: §2 · p. 5, Figure 2 · p. 5
  • A plausible explanation is that broad retrieval accumulates representational noise as needle counts rise, while longer contexts can amplify per-step retrieval errors across multiple steps.The latter implication applies to the proposed model in which each retrieval step can fail independently.Source: §2 · p. 5

3 Non-thinking mode: form, retrieve, and consolidate

Non-thinking counting follows a form–retrieve–consolidate pathway: broad answer-query retrieval gathers multiple needles, while count information forms in shallow layers and is consolidated at the answer state. The study finds that these representations remain noisy and that selected retrieval heads and answer states causally affect count predictions.

3 Non-thinking mode: form, retrieve, and consolidate

  • Robustness tests include instructed enumeration with Thinking disabled and other retrieval tasks; the analysis primarily uses Qwen3-8B and Gemma-4-E4B with approximately 10k-token prompts.Instructed enumeration is also presented as a way to study closed-source models through verifiable intermediate outputs; retrieval heads are evaluated on held-out seeds.Source: §3 · pp. 5–6

3.1 Broad attention, noisy representations

  • At the answer query, selected heads attend across multiple needle spans, although individual heads show position preferences.Figure 3A reports needle-attention shares for each model’s six highest-ranked heads.Source: §3.1 · p. 6, Figure 3 · p. 6
  • 81.2% for Qwen Top-32 heads and 80.3% for Gemma Top-6 heads: mean effective coverage across prompts with 2–10 needles.The score rewards both needle-directed attention and broad coverage; the highlighted head sets are shown in Figure 3B.Source: Figure 3 · p. 6, §3.1 · p. 6
  • Needle-end states predict running index well by ridge regression, but linear steering does not reliably shift the answer in the intended direction.Held-out R2 is 0.910 for Qwen and 0.827 for Gemma; steering was tested on ten initially correct N = 3 prompts per model.Source: §3.1 · pp. 6–7
  • The first three principal components explain 98.1% of Qwen and 97.5% of Gemma class-mean variance, yet held-out individual-state NCC accuracy is 46% and 40%.Individual states overlap despite the low-dimensional structure of their class means.Source: §3.1 · p. 7

3.2 Causal analysis of the mechanism

  • At shallow layers, restoring clean needle-span states into corrupted inputs, or reversing the patch, changes normalized count error by 47–50% at L1; effects fall below 5% past Qwen L23 and Gemma L18.These interventions support count-encoding representations forming in shallow layers.Source: §3.2 · pp. 7–8, Figure 4 · p. 7
  • Selected retrieval-head ablation changes the generated count by 54.5% for Qwen and 24.7% for Gemma, versus 4.0% and 6.0% for layer-matched random heads.The ablation is applied at the answer query; Figure 4C compares selected heads with random and original controls.Source: Figure 4 · p. 7, §3.2 · p. 8
  • At the final layer, answer-state patches produce predictions matching the source count in 95.8% of Qwen and 100% of Gemma cases.Subspace ablation also supports a causal role for count-related information in the answer state.Source: Figure 4 · p. 7, §3.2 · p. 8

4 Thinking mode: retrieve, update, and read

In Thinking mode, successive trace queries target the next needle, and trace-item states form compact representations of counting progress. Interventions support a retrieve–update–read pathway in which counter states guide continued enumeration and the completed trace supports the final count.

4.1 Targeted attention, compact representations

  • At the query before each trace item, the shown heads mainly attend to the next record as enumeration proceeds; analysis is restricted to successfully parsed trace items.Figure 5A marks needle k + 1 at Marker k; Figure 5B highlights the Qwen Top-128 and Gemma Top-6 heads.Source: §4.1 · p. 8, §4 · p. 8, Figure 5 · p. 9
  • Held-out NCC accuracy is 98% for Thinking trace-item states versus 46% for Non-thinking prompt-needle states; answer-state readout also drops less under needle-domain changes in Thinking.Replacing city needles with flower needles reduces accuracy by 2 percentage points for Qwen and none for Gemma in Thinking, versus 17 and 15 points in Non-thinking.Source: §4.1 · p. 9

4.2 Causal analysis of the mechanism

  • Ablating targeted retrieval heads raises next-item failure by 60 percentage points for Qwen and 80 for Gemma, versus 10 and 6.7 for random heads; final-count failure rises by 100 points in both models.Figure 6B compares selected-head ablation with layer-matched random and original controls.Source: §4.2 · p. 9, Figure 6 · p. 10
  • Without explicit indices, counter-state patching yields the next record in 24/30 Qwen and 21/30 Gemma trials, versus 0/30 and 2/30 with own-state patches.Among first-step successes, the following item matches the next record in 24/24 and 20/21 trials without another patch.Source: §4.2 · p. 10, Figure 6 · p. 10
  • Removing trace representations reduces exact-count accuracy to 1% for Qwen and 15% for Gemma, while removing prompt-record representations leaves it at 97% and 70%.Final-layer answer-state patches match the source count in 97.5% of Qwen and 90% of Gemma predictions.Source: §4.2 · p. 10, Figure 6 · p. 10

5 Synthetic Experiment: Emergence of Two Retrieval Mechanisms

The controlled training experiment compares small autoregressive counting models with and without enumeration traces. The trace-trained model achieves higher free-running accuracy and develops targeted retrieval, while the model without traces develops broad retrieval.

5 Synthetic Experiment: Emergence of Two Retrieval Mechanisms

  • Both four-layer, eight-head models share the modified Shakespeare data, architecture, and training budget; only Thinking generates an enumeration trace before the answer.Training uses target counts 1–10 and a two-stage schedule, switching to task-output loss at step 1,500.Source: §5 · p. 11
  • Free-running accuracy reaches 95% for Thinking but plateaus at 21% for Non-thinking; broad retrieval emerges in the latter’s first layer and targeted retrieval in the former’s last layer.The comparison uses one training seed per mode and 100 held-out inputs.Source: §5 · p. 11, Figure 7 · p. 11
  • Enumeration traces without explicit numbering appear to influence whether models learn broad or targeted retrieval, and standard autoregressive training can yield strong counting accuracy without specialized RL post-training.The authors present these findings as evidence about training dynamics and data.Source: §5 · pp. 11–12

6 Limitations and Future Work

The authors identify limits to the mechanistic evidence and the counting testbed, and point to circuit identification and more complex reasoning as future work. They also provide code and appendix details for reproduction.

6 Limitations and Future Work

  • Interventions provide evidence that counter states guide continued enumeration, but identifying circuits that implement needle-by-needle updates remains future work.The transferred states jointly carry count information and record content.Source: §6 · p. 12
  • The study uses counting as a clean testbed; extending the analysis to reasoning with search, branching, and multiple evolving states is left for future work.Reproduction code is available, with appendix details on prompts, synthetic datasets, training hyperparameters, and representation measurements.Source: §6 · p. 12, Reproducibility · p. 12

Reproducibility

  • The paper provides code for reproducing results and appendix details on prompts, datasets, training hyperparameters, and representation measurements.The appendices also include additional supporting analyses.Source: Reproducibility · p. 12

AI Use Statement

The authors disclose generative-AI assistance for literature search, synthetic data, code, repository organization, writing, and figure editing. They state that they reviewed the technical work and take responsibility for the final content.

AI Use Statement

  • Generative-AI tools were used for literature searches, synthetic data generation, code implementation, repository organization, writing, and figure editing.The authors state these tools were not used to design the conceptual framework, hypotheses, or research methodology.Source: AI Use Statement · p. 12
  • The authors say they reviewed all technical claims, experimental designs, implementations, and reported results.They also state that the tools were not used to interpret results or suggest the paper structure.Source: AI Use Statement · p. 12
  • The authors take responsibility for the final content, including text, claims, and artifacts produced with generative-AI assistance.Source: AI Use Statement · p. 12

A Details of the main-text figures

The appendix documents how to interpret the main-text figures, the benchmark and prompting protocols, and an exploratory analysis of Non-thinking representation noise. It also names work that inspired the figures.

A Details of the main-text figures

  • Figure 1 notes describe prompt-attention bins and successive needle-normalized attention, with marker positions indicating initialization and completed trace items.The illustrated examples use ten-record passages of approximately 10k tokens.Source: A Details of the main-text figures · p. 18
  • Figure 2 plots accuracy by target count and passage length; long-context panels compare Qwen3-32B and Gemma-4-31B, while Figure 3 reports needle-attention shares and head-level retrieval scores.Figure 2 uses group medians and interquartile ranges for its short-context benchmark; unparseable outputs count as errors.Source: Figure 2 · p. 18, A Details of the main-text figures · p. 19
  • Figure 4 defines patching and ablation outcomes, Figure 5 specifies targeted-attention scoring, and Figure 6 describes counter-state transfer and ablation comparisons.Figure 4 also reports source-count match rates; Figure 6 notes that its schematic does not establish a fully mediated chain on one common cohort.Source: A Details of the main-text figures · pp. 19–20
  • Figure 7 tracks free-running accuracy and layer-wise retrieval scores, with the dashed line marking the switch to task-output loss at step 1,500.Its accuracy evaluation uses eleven checkpoints on the same 100 held-out inputs, with one training seed per mode and no error bars.Source: Figure 7 · p. 20, A Details of the main-text figures · p. 20

B Additional related work

  • The figure designs draw on work about induction heads, circuit tracing, attention tracing, representation manifolds, and scaling-law plots.Source: B Additional related work · p. 20

C.1 Prompts and experimental design

  • The behavioral benchmark inserts city-score records into shuffled Paul Graham essays and evaluates 12 comparison groups across counts, passage lengths, and paired seeds.The short-context grid has 14 counts, eight passage lengths from 1k to 20k, and 30 paired seeds per condition.Source: §C.1 · pp. 20, 21, 23
  • The benchmark uses a shared counting prompt with mode-specific response instructions; Non-thinking decoding is greedy, while Thinking uses group-specific sampling settings.The response instructions prohibit listing records for Non-thinking and specify concise reasoning for Thinking.Source: §C.1 · pp. 21, 23

C.2 Non-thinking: relative noise

  • Relative noise rises overall in two Non-thinking models, mainly as adjacent-count mean separation shrinks; the exploratory analysis does not establish growing absolute noise or model use of the measured direction.The reported high/low ratios are 7.93 for Qwen and 8.64 for Gemma; intervals are not simultaneous bands.Source: §C.2 · p. 24, Figure 8 · p. 24

C.3 Count dependence: models, fits, and complete group results

Across the model groups, Thinking has higher overall mean counting accuracy than Non-thinking, while both modes show local failures at long contexts. Length-model comparisons favor a coarse linear trend for Thinking, but do not identify a unique length law.

C.3 Count dependence: models, fits, and complete group results

  • The length-fit γ is an effective parameter from final-answer outcomes, not a direct measure of retrieval failure, and fit quality does not establish independent step failures.The study fits each group separately and validates by holding out each count; shared seeds prevent treating conditions as independent for confidence intervals.Source: §C.3 · p. 25
  • Cross-model median fits achieve held-count R2 of 0.828 for Non-thinking and 0.924 for Thinking, with held-count RMSE of 15.2 and 3.6 points, respectively.The Non-thinking transition shows systematic departures from its fitted curve.Source: §C.3 · p. 26
  • Thinking’s overall mean accuracy is higher in every group.The complete comparison covers 12 groups, 14 counts, and eight lengths.Source: Table 2 · p. 26, §C.3 · p. 26, Figure 10 · p. 27

C.4 Complete long-context observations

  • At N = 1 and L = 100k, Qwen accuracy is 22/30 for Non-thinking and 14/30 for Thinking; Gemma’s N = 20 accuracy varies from 0/30 at 5k to 21/30 at 8k and 0/30 at 15k.Twelve Qwen Thinking traces hit the output limit; the cause of Gemma’s local increase is not identified.Source: §C.4 · p. 27, Figure 11 · p. 28, Figure 12 · p. 29

C.5 Which length functions describe Thinking?

  • Linear γ(L) has lower held-length RMSE than logarithmic γ(L) in 11 of 12 short-context groups, but linear-λ and reciprocal-success fits have similar errors.Median held-length RMSEs are 4.9 versus 9.0 points; frozen short-to-long errors remain 44.0 points for Qwen and 18.4 for Gemma even with linear γ.Source: §C.5 · p. 31, Figure 14 · p. 31, Table 3 · p. 32

D.1 Experimental design

  • The Non-thinking experiments use matched prompts and define expected count over candidate answers 1–10, excluding other outputs through renormalization.The default split uses seeds 1234–1253 for selection and fitting and 1254–1263 for held-out evaluation.Source: §D.1 · pp. 32–33

D.2 Broad retrieval and count representations

The analyses test how broad and targeted retrieval relate to count representations and generated answers. Readouts and interventions show separable count information and support causal roles for retrieval and answer-state components, with geometry and transfer results bounded by their evaluation settings.

D.2 Broad retrieval and count representations

  • Broad-retrieval running-index NCC reaches 46%/40% held-out accuracy for Qwen/Gemma, while final-count NCC peaks at 67.5%/57% on cross-validation folds.Three components explain 98.1%/97.5% of running-index centroid variance, although individual states overlap.Source: §D.2 · p. 35, Figure 16 · p. 35

D.3 Causal tests of the form–retrieve–consolidate pathway

  • On initially correct prompts, ranked-head ablation causes correct-to-wrong rates of 67.5% versus 8.8% for random banks in Qwen and 33.3% versus 12.5% in Gemma.The comparison uses Top-32 Qwen and Top-6 Gemma banks.Source: §D.3 · p. 37, Figure 18 · p. 37
  • Final-layer answer-state patches match the source count on 95.8% of Qwen pairs and 100% of Gemma pairs; same-count and self-patch controls yield zero matches at every layer.Further single-layer and cumulative patches produce source-count matches totaling 1628/1686 for Qwen and 1809/1884 for Gemma.Source: §D.3 · p. 38, Figure 19 · p. 38
  • Removing aligned retrieval and answer components increases normalized count error by 12.6/9.8 and 53.9/52.0 percentage points, respectively, for Qwen/Gemma.The results support causal contributions from both components; joint removal leaves slightly more error than ordinary-region restoration.Source: §D.3 · p. 39, Figure 21 · p. 40

E.1 Experimental design

  • The experiments split seeds into selection/fitting and held-out sets; parsed trace items determine available running states, and repeated records remain indexed by occurrence.The primary analyses cover N = 1–10 and do not filter states by answer correctness.Source: §E.1 · pp. 40–41, §E.2 · p. 42

E.2 Targeted retrieval and compact count representations

  • Held-out running-index NCC/logistic accuracy is 58.9/68.1% for Qwen and 71.8/78.0% for Gemma; final-count accuracy is 99/100% and 70/71%, respectively.Visible indices, token position, trace format, and total count can covary with the running index and are not separated by these readouts.Source: §E.2 · p. 42, Figure 23 · p. 43
  • PCA3 explains 35.2/41.4% of running-state variance and 53.4/58.6% of final-state variance for Qwen/Gemma, while class means summarize heterogeneous, overlapping states.These analyses establish label separability; the paper treats causal tests as addressing different specified populations.Source: §E.2 · p. 43, Figure 24 · p. 44
  • City-trained count readouts are evaluated on held-out city, flower, and animal trajectories, with topic changes also altering tokenization and generated traces.Running-state populations need not align across topics.Source: §E.2 · p. 44, Table 7 · p. 45

E.3 Causal tests of the retrieve–update–read pathway

Causal interventions link retrieval and count readout to the trace, while held-out representations make running indices and final counts decodable. These tests use selected retrieval heads and controlled counting cohorts.

E.3 Causal tests of the retrieve–update–read pathway

  • Qwen next-item accuracy stayed at least 9/10 through K = 64 but fell to 4/10 at K = 128; final-count accuracy fell to 0/10 at K = 64.For Gemma, Top-1 left both accuracies at 1/10 and K = 2, 4, 6 left both at zero; persistent masks also affected later generation.Source: Figure 26 · p. 46, §E.3 · pp. 46–47
  • Forward state transfer induced the source-implied next record in 24/30 Qwen and 21/30 Gemma trials, while backward transfer repeated the last completed record in 29/30 and 22/30.After backward transfer, the following item was correct in only 5/29 Qwen cases versus 18/22 Gemma cases; the passage cautions that trace-format differences limit this comparison.Source: §E.3 · pp. 48–49, Figure 27 · p. 48
  • Whole-trace blanking lowered exact-count accuracy from 97% to 1% for Qwen and 70% to 15% for Gemma, while prompt-record blanking preserved more accuracy.Answer-state patches also transferred source-count predictions in 39/40 Qwen and 36/40 Gemma pairs at the final layer; full-state transfer does not isolate count from numeral encoding.Source: §E.3 · p. 49, Figure 28 · p. 49

F.1 Experimental design

  • The unnumbered Bullet enumeration experiment tests retrieval, count representations, and transfer using roughly 10k-token passages and held-out seeds.Strict success required correct totals, ordered record lists, and no truncation; Qwen succeeded on 297/300 inputs and Gemma on 229/300.Source: §F.1 · pp. 49–50

F.2 Targeted retrieval and compact count representations

  • Held-out running-index readout reached 69.4/75.3% for Qwen and 56.7/68.0% for Gemma, while final-count readout reached 100/100% and 88/85%, respectively.Values are NCC/logistic balanced accuracy at Qwen L19/L30 and Gemma L24/L38 for running index/final count; readouts do not separate content, syntax, position, and progress.Source: Figure 30 · p. 51, §F.2 · p. 52
  • The PCA analyses establish label separability, but independent projection bases prevent direct cross-panel distance comparisons.The displayed PCA3 space also differs from the whitened PCA16 readout space.Source: Figure 31 · p. 52, §F.2 · p. 53

F.3 Causal tests of the retrieve–update–read pathway

Selected retrieval heads affect subsequent retrieval and counting, and state patches transfer progress-related behavior under unnumbered enumeration. The paper also reports a controlled synthetic-training setup and distinct broad- and targeted-head ablation outcomes.

F.3 Causal tests of the retrieve–update–read pathway

  • At the largest head doses, selected-minus-clean next-item/final-count failure increases were 50/50 percentage points for Qwen and 100/70 for Gemma.Final-count contrasts against random banks were 30.0 points for Qwen and 76.7 for Gemma; persistent masks can also affect later generation.Source: §F.3 · pp. 53–54, Figure 32 · p. 53
  • Item-span patches transfer record content and progress jointly, so they do not isolate a pure count update.Patch norms were unmatched, and self adoption was 0/30 in each model and direction.Source: Figure 33 · p. 54, §F.3 · p. 55
  • Forward item-span patches produced the source-implied successor in 25/30 Qwen and 21/30 Gemma trials; backward patches repeated the last completed record in 29/30 and 21/30.After these patches, the following item was correct in 24/25 and 19/21 forward trials, and 22/29 and 17/21 backward trials.Source: Figure 33 · p. 54, §F.3 · p. 55
  • Trace blanking reduced correct counts from 100 to 12 for Qwen and from 79 to 30 for Gemma, while prompt-record blanking yielded 91 and 62.Whole-state answer patches matched source counts in 31/40 Qwen and 40/40 Gemma pairs at the final layer, but do not isolate count from numeral encoding.Source: §F.3 · pp. 55–56, Figure 34 · p. 55

G Further details for the synthetic experiment

  • The synthetic experiment trains separate decoder-only models from scratch for direct-answer and enumeration formats on a small character corpus.Thinking enumerates matching characters in source order without explicit indices; there is one training run per mode.Source: G Further details for the synthetic… · p. 56, §G.1 · pp. 56–57

G.1 Experimental design

  • After step 1,500, training switches from next-token loss to weighted task-output losses, including retrieved-character tokens for Thinking.The objective switch is preset, and the evaluation metric remains full-sequence cross-entropy.Source: §G.1 · pp. 58–59, Figure 35 · p. 59
  • Validation and training-region loss curves largely overlapped and did not identify a clear onset of overfitting on the approximately 1 MB corpus.The paper notes that the loss metric remains fixed across the objective switch.Source: §G.1 · p. 59, Figure 35 · p. 59

G.2 Broad and targeted retrieval heads

  • Non-thinking Top-4 broad-head ablation reduced final-count accuracy to 7% versus a 15% control mean; Thinking Top-4/Top-8 targeted-head ablations yielded 85.2%/68.2% next-needle accuracy.The corresponding Thinking control means were 93.2% at Top-4 and 77.3% at Top-8.Source: §G.2 · p. 61, Figure 37 · p. 61

G.3 Count representations and causal interventions

Count representations are decodable at selected layers, but probing alone does not show that models use the decoded information. Interventions test sensitivity to prompt and trace states, continuation after transplanted states, and transfer of answer predictions.

G.3 Count representations and causal interventions

  • Running-index states also carry structural cues because each Thinking occurrence adds two trace tokens, so progress covaries with trace length and position.The PCA projections are fitted separately by mode and endpoint, so their coordinates are not directly comparable.Source: §G.3 · p. 61, Figure 38 · p. 62
  • At NCC-selected layers, final-count decoding reached 99% for Thinking at L2 and 24% for Non-thinking at L4.Thinking running-index accuracy ranged from 15% to 34% across blocks; decoding alone does not establish model use of this information.Source: §G.3 · p. 63, Figure 39 · p. 63
  • Next-needle accuracy fell from 96% to 21% after prompt-target ablation and to 79% after trace ablation; on inputs with prior trace items, either trace ablation reduced accuracy from 94.3% to 70.0%.The next-needle prediction is sensitive to both prompt targets and available trace history.Source: §G.3 · pp. 63–64, Figure 40 · p. 63
  • A one-character source transplant increased source-continuation matches by 38.3 points over self-patching and 16.0 points over a norm-matched random control.Full-item transplantation transfers content and progress jointly, and differing eligible pairs and wide intervals limit conclusions about sustained continuation.Source: §G.3 · p. 64, Figure 41 · p. 65
  • Trace ablation reduced Thinking accuracy from 95% to 59%, while prompt-target ablation left it at 95%.Answer-state transfer raised agreement with the source prediction from 0% to 100% for Thinking by L3; both modes reached 100% at L4.Source: §G.3 · p. 64, Figure 41 · p. 65

G.4 Learning dynamics of retrieval and counting

Training analyses track retrieval attention and count decoding over checkpoints, revealing distinct timing for attention patterns, running-index states, and final-count readout. A targeted-head intervention further tests the role of a successor-like attention head in trace processing.

  • Non-thinking develops broad retrieval in shallow layers, while Thinking develops targeted retrieval in deeper layers; at step 400, maximum Uh is 0.267 versus maximum Th at 0.006.Scores are averaged equally across evaluation inputs.Source: §G.4 · p. 65, Figure 42 · p. 66
  • Zeroing L2H4’s output at generated trace-character positions after <Sep> lowers final-count accuracy from 95% to 45%, versus 79–94% for same-layer controls.Exact-trace accuracy also falls from 87% to 39%, supporting L2H4’s contribution to trace processing.Source: §G.4 · p. 65
  • At step 400, Thinking final-count decoding is 100% with a correct teacher-forced trace, versus 12% free-generation count accuracy and 2% exact-trace accuracy.The final count becomes readable given a correct trace before the model can reliably generate that trace.Source: §G.4 · pp. 66–67, Figure 42 · p. 66
  • Running-index decoding falls from 73.7% at step 200 to 31.4% at the final checkpoint; at step 200, maximum Uh is 0.025, so its role in the early peak remains unverified.Source: Figure 42 · p. 66, §G.4 · p. 67

G.5 Relation to the main experiments

The paper relates controlled retrieval and counter-state experiments to its main findings, then extends behavioral and head-ablation tests to record retrieval and category counting. Thinking performs better on both additional tasks, while targeted-head interventions impair next-entity generation.

G.5 Relation to the main experiments

  • Scratch-trained synthetic experiments reproduce broad retrieval in Non-thinking and targeted retrieval in Thinking; selected-head ablations affect final-count prediction and next-needle retrieval, respectively.The synthetic Thinking assay supports local retrieval, while the main experiments additionally test final-count effects.Source: §G.5 · p. 67
  • Full-item state transplants increase source-continuation adoption at short horizons, supporting a role for trace-item states in subsequent enumeration while jointly transferring content and progress.Source: §G.5 · p. 67
  • Thinking improves final-count decoding despite weakly separated intermediate running-index states; with a generated trace, final answering depends on trace states, while prompt-target ablation leaves accuracy unchanged.Late answer-state transplants steer both modes toward the source run’s prediction.Source: §G.5 · p. 67

H Further details for additional tasks

H.1 Experimental design

  • Both tasks evaluate Qwen3-8B and Gemma-4-E4B without additional training, using 300 inputs from 30 seeds and passages of approximately 10,000 tokens.The inputs are shared across models and modes; only selection inputs determine head rankings.Source: §H.1 · pp. 67–68

H.2 Behavioral comparison between the two modes

  • Thinking achieves higher final-answer accuracy for both models on both tasks; Qwen remains high across positions and category counts, while Gemma generally declines at later positions and larger counts.Gemma Thinking outputs include 12 unparsed retrieval answers and 16 unparsed category answers; comparisons include each mode’s instructions and generation budget.Source: §H.2 · pp. 69–70, Table 11 · p. 70, Figure 43 · p. 70

H.3 Retrieval-head distributions and causal tests

Task-specific targeted-head ablations impair next-entity accuracy more than matched random ablations in the additional tasks. Broad-head ablation effects are weaker and less consistent, and the two assays do not support direct comparisons of head importance.

H.3 Retrieval-head distributions and causal tests

  • Heads are ranked by task-, model-, and condition-specific mean scores on selection inputs; targeted retrieval requires an aligned transition, and selection and intervention inputs are disjoint.Targeted ablation acts at the registered transition and subsequent decoding positions, whereas broad ablation acts at the final-answer query.Source: §H.3 · pp. 70–71
  • For category counting at largest K, Selected versus Random next-entity accuracy is 16.5% versus 91.3% for Qwen and 17.6% versus 53.6% for Gemma.Qwen’s Selected accuracy decreases monotonically over scanned K values; Gemma’s separation is pronounced at K = 6 and 8, with mixed smaller-K effects.Source: §H.3 · p. 72, Figure 45 · p. 73
  • Broad-head ablation lowers final-answer accuracy at some K, but its separation from Random is weaker and less consistent; low Clean accuracy limits observable declines.Broad-assay Clean accuracies are 13.00/26.00% for Qwen/Gemma retrieval and 25.00/29.00% for category counting.Source: §H.3 · pp. 72–73
  • For kth-record retrieval at largest K, Selected versus Random next-entity accuracy is 10.9% versus 98.1% for Qwen and 44.7% versus 87.9% for Gemma.The layer-matched controls show impairment beyond removing the same number of heads from the same layers.Source: §H.3 · p. 72, Figure 45 · p. 73

H.4 Relation to the main experiments

  • Different outcomes, eligible inputs, and intervention durations preclude direct comparisons of head importance between the broad and targeted assays.The additional-task results are consistent with the main experiments: targeted ablation impairs next-entity accuracy more than layer-matched random ablation, while broad effects are weaker and less consistent.Source: §H.4 · p. 73
Loading 2609.38958v1…