Diagnosing Memorization in Chain-of-Thought Reasoning, One Token at a Time
Huihan Li, You Chen, Siyuan Wang, Yixin He, Ninareh Mehrabi, Rahul Gupta, Xiang Ren
Abstract
Large Language Models (LLMs) perform well on reasoning benchmarks but often fail when inputs alter slightly, raising concerns about the extent to which their success relies on memorization. This issue is especially acute in Chain-of-Thought (CoT) reasoning, where spurious memorized patterns can trigger intermediate errors that cascade into incorrect final answers. We introduce STIM, a novel framework for Source-aware Token-level Identification of Memorization, which attributes each token in a reasoning chain to one of multiple memorization sources -local, mid-range, or long-rangebased on their statistical co-occurrence with the token in the pretraining corpus. Our token-level analysis across tasks and distributional settings reveals that models rely more on memorization in complex or long-tail cases, and that local memorization is often the dominant driver of errors, leading to up to 67% of wrong tokens. We also show that memorization scores from STIM can be effective in predicting the wrong tokens in the wrong reasoning step. STIM offers a powerful tool for diagnosing and improving model reasoning and can generalize to other structured step-wise generation tasks. 1 * Equal Contribution † Work done in Amazon AGI 1 https://github.com/INK-USC/STIM Model Reasoning Steps: Rick killed 10 + 15 = 25 animals last night. Today Rick killed 3 * 3 = 9 wolves. Today Rick killed 15 -3 = 12 cougars. Today Rick killed 9 + 12 = 21 animals. Total animals killed by Rick is 25 + 21 = 46. So the answer is 46. Input Question: Last night Rick killed ten wolves and 15 cougars while hunting. Today Rick killed three times as many wolves as cougars and three fewer cougars than the previous night. How many animals did Rick kill? Examples of Simultaneous Memorization Sources of Different Strength Question: Here is a list: [pitanga, pitanga, yumberry, yumberry, pitanga, yumberry, pitanga, yumberry, pitanga, pitanga]. How many times does 'pitanga' appear on it? Model Output: 1. The first element is 'pitanga'. ... 10. The tenth element is 'pitanga'. Now, counting the occurrences of 'pitanga': We have 'pitanga' appearing 5 times at positions 1, 2, 5, 7, and 10. So the answer is 5. Analysis: Token 5 is incorrect and has the highest influence from local context, where the n-gram "appearing 5" has a high pretraining frequency. STIM scores are: local = 0.81 , mid = 0.24, long = 0.03, suggesting that the model was influenced by locally frequent patterns rather than actual counting.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 39175b80-cfb7-42fa-8660-c3a072c91ee3Cited by top-tier papers2
- Reasoning or Retrieval? A Study of Answer Attribution on Large Reasoning ModelsYuhui Wang, Changjiang Li, Guangke Chen, Jiacheng Liang et al.ICLR 2026 · 13 citations
- SSVPO: Effective Step-Level Credit Assignment for RL Training of Language ModelsYugu Li, Zehong Cao, Jianglin Qiao, Siyi HuICLR 2026
Builds on7
- Let's Verify Step by StepHunter Lightman, Vineet Kosaraju, Yuri Burda, Harrison Edwards et al.ICLR 2024 · 3,045 citations
- Quantifying Memorization Across Neural Language ModelsNicholas Carlini, Daphne Ippolito, Matthew Jagielski, Katherine Lee et al.ICLR 2023 · 158 citations
- Local Explanation of Dialogue Response GenerationYi-Lin Tuan, Connor Pryor, Wenhu Chen, Lise Getoor et al.NeurIPS 2021 · 13 citations
- Generalization v.s. Memorization: Tracing Language Models' Capabilities Back to Pretraining DataXinyi Wang, Antonis Antoniades, Yanai Elazar, Alfonso Amayuelas et al.ICLR 2025 · 4 citations
- Evaluating n-Gram Novelty of Language Models Using Rusty-DAWGWilliam Merrill, Noah A. Smith, Yanai ElazarEMNLP 2024 · 2 citations
Related papers
- Token Signature: Predicting Chain-of-Thought Gains with Token Decoding Feature in Large Language ModelsPeijie Liu, Fengli Xu, Yong LiICML 2025
- Not All Tokens Matter: Towards Efficient LLM Reasoning via Token Significance in Reinforcement LearningHanbing Liu, Lang Cao, Yuanyi Ren, Mengyu Zhou et al.ACL 2026 · 5 citations
- Scaling Reasoning Hop Exposes Weaknesses: Demystifying and Improving Hop Generalization in Large Language ModelsZhaoyi Li, Jiatong Li, Gangwei Jiang, Linqi Song et al.ICLR 2026 · 1 citation
- Unveiling Factual Recall Behaviors of Large Language Models through Knowledge NeuronsYifei Wang, Yuheng Chen, Wanting Wen, Yu Sheng et al.EMNLP 2024 · 3 citations
- SIM-CoT: Supervised Implicit Chain-of-ThoughtXilin Wei, Xiaoran Liu, Yuhang Zang, Xiaoyi Dong et al.ICLR 2026 · 58 citations
