Elamur Research⌕
☰
home / papers / 2025 / All for One: LLMs Solve Mental Math at the Last Token With Information Transferred From Other Tokens

All for One: LLMs Solve Mental Math at the Last Token With Information Transferred From Other Tokens

Paper · 2025 · Siddarth Mamidanna, Daking Rai, Ziyu Yao et al.

Summary

The paper asks how much of a transformer's theoretical capacity for cross-token communication LLMs actually use when they do "mental math": one-token answers to two- or three-operand arithmetic, with no chain-of-thought. The authors progressively strip computation out of Llama-3-8B using two new techniques. Context-Aware Mean Ablation (CAMA) replaces early-layer activations with context-conditioned averages, removing input-specific information while keeping representations in-distribution. Attention-Based Peeking (ABP) restricts which positions a token may attend to. The result is a sparse three-stage subgraph called All-for-One (AF1):

  • In early layers, tokens perform only task-general processing.
  • In a few middle layers, information flows to the last token.
  • In the remaining layers, all real computation happens at that last position.

AF1 keeps high accuracy across symbolic and verbal arithmetic prompts and partly transfers to Pythia and GPT-J. Because even A + B + C is solved without intermediate results being stored at earlier tokens, the findings suggest LLMs compose little across positions. CAMA and ABP may also be useful interpretability tools beyond arithmetic.

Key points

  • For Llama-3-8B and Llama-3.1-8B, the waiting stage (L_wait) can cover close to half of all layers, and as few as 2 transfer layers (L_transfer) suffice to recover strong performance.
  • For Pythia and GPT-J, the best AF1 configurations are (L_wait, L_transfer) = (9, 7) and (9, 8). Their faithfulness on two-operand tasks ranges from roughly 0.44 to 0.79, with ×B the highest and ÷B the lowest.
  • Alternative waiting mechanisms all reach zero faithfulness. These include direct embedding copy (DEC), random token mean ablation (RTMA), self-peek as waiting (SPAW) and isolated forward pass (IFP).
  • AF1 performs poorly on tasks needing deeper semantic or contextual understanding.
  • The method assumes each number maps to a single token, so models such as Qwen and Gemma, which split numbers into digit tokens, are not yet covered. Code is available at github.com/siddarth-pm/all-for-one.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🎓University of California, Santa CruzSanta Cruz · United States36 works · 3,515 citations🎓George Mason UniversityFairfax · United States24 works · 1,288 citations

 Fields, topics and concepts

📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 16,667 worksArtificial Intelligence · 16,194 worksArts and Humanities · 16,039 worksLanguage and Linguistics · 15,891 works

🧭 Topics 1Research subjects and languages the work is about.Language models & pretraining · 3,320 works

Cite

Cite as Siddarth Mamidanna, Daking Rai, Ziyu Yao and Yilun Zhou. 2025. All for One: LLMs Solve Mental Math at the Last Token With Information Transferred From Other Tokens. Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.18653/v1/2025.emnlp-main.1565

@inproceedings{mamidanna2025all,
  title = {{All for One: LLMs Solve Mental Math at the Last Token With Information Transferred From Other Tokens}},
  author = {Siddarth Mamidanna, Daking Rai, Ziyu Yao and Yilun Zhou},
  booktitle = {Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing},
  year = {2025},
  doi = {10.18653/v1/2025.emnlp-main.1565},
  url = {https://doi.org/10.18653/v1/2025.emnlp-main.1565}
}

References

40 references, 8 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Leonard Bereska and Stratis Gavves. 2024. Mechanistic interpretability for AI safety - a review. Transactions on Machine Learning Research. Survey Certification, Expert Certification. find
  2. Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O’Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, and 1 others. 2023. Pythia: A suite for analyzing large language models across training and scaling. In International Co… find
  3. Eden Biran, Daniela Gottesman, Sohee Yang, Mor Geva, and Amir Globerson. 2024. Hopping too late: Exploring the limitations of large language models on multi-hop queries. In Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pages 14113–14130.
  4. Nicola Cancedda. 2024. Spectral filters, dark signals, and attention sinks. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 4792– 4808. find
  5. Lawrence Chan, Adrià Garriga-Alonso, Nicholas Goldwosky-Dill, Ryan Greenblatt, Jenny Nitishinskaya, Ansh Radhakrishnan, Buck Shlegeris, and Nate Thomas. 2022. Causal scrubbing, a method for rigorously testing interpretability hypotheses. AI Alignment Forum. https://www.alignmentforum.org/posts/ ↗
  6. Bilal Chughtai, Alan Cooney, and Neel Nanda. 2024. Summing up the facts: Additive mechanisms behind factual recall in llms. arXiv preprint arXiv:2402.07321. — [arXiv:2402.07321](https://arxiv.org/abs/2402.07321) find
  7. Róbert Csordás, Christopher D. Manning, and Christopher Potts. 2025. Do language models use their depth efficiently? Preprint, arXiv:2505.13898v2. find
  8. Subhabrata Dutta, Joykirat Singh, Soumen Chakrabarti, and Tanmoy Chakraborty. 2024. How to think step-by-step: A mechanistic understanding of chain-of-thought reasoning. Trans. Mach. Learn. Res. find
  9. Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, and 6 others. 2021. A mathe… find
  10. Mor Geva, Jasmijn Bastings, Katja Filippova, and Amir Globerson. 2023. Dissecting recall of factual associations in auto-regressive language models. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pages 12216–12235. find
  11. Nicholas Goldowsky-Dill, Chris MacLeod, Lucas Sato, and Aryaman Arora. 2023. Localizing model behavior with path patching. arXiv preprint arXiv:2304.05969. find
  12. Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, and 1 others. 2024. The llama 3 herd of models. arXiv preprint arXiv:2407.21783.
  13. Tal Haklay, Hadas Orgad, David Bau, Aaron Mueller, and Yonatan Belinkov. 2025. Position-aware automatic circuit discovery. Preprint, arXiv:2502.04577. find
  14. Roee Hendel, Mor Geva, and Amir Globerson. 2023. In-context learning creates task vectors. In Findings of the Association for Computational Linguistics: EMNLP 2023, pages 9318–9333.
  15. Evan Hernandez, Arnab Sen Sharma, Tal Haklay, Kevin Meng, Martin Wattenberg, Jacob Andreas, Yonatan Belinkov, and David Bau. 2023. Linearity of relation decoding in transformer language models. CoRR. find
  16. Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. 2020. Transformers are rnns: Fast autoregressive transformers with linear attention. In International conference on machine learning, pages 5156–5165. PMLR. find
  17. Maximilian Li and Lucas Janson. 2024. Optimal ablation for interpretability. Advances in Neural Information Processing Systems, 37:109233–109282. find
  18. Zhaoyi Li, Gangwei Jiang, Hong Xie, Linqi Song, Defu Lian, and Ying Wei. 2024. Understanding and patching compositional reasoning in LLMs. In Findings of the Association for Computational Linguistics: ACL 2024, pages 9668–9688, Bangkok, Thailand. Association for Computational Linguistics. find
  19. Davide Maltoni and Matteo Ferrara. 2024. Arithmetic with language models: From memorization to computation. Neural Networks, 179:106550. find
  20. Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. 2022. Locating and editing factual associations in gpt. Advances in neural information processing systems, 35:17359–17372. find
  21. Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. 2023. Progress measures for grokking via mechanistic interpretability. arXiv preprint arXiv:2301.05217.
  22. Yaniv Nikankin, Anja Reusch, Aaron Mueller, and Yonatan Belinkov. 2024. Arithmetic without algorithms: Language models solve math with a bag of heuristics. arXiv preprint arXiv:2410.21272. nostalgebraist. 2020. Interpreting gpt: the logit lens. AI Alignment Forum. https://www.lesswrong.com/ posts/Ac… ↗
  23. Daking Rai and Ziyu Yao. 2024. An investigation of neuron activation as a unified lens to explain chain-of-thought eliciting arithmetic reasoning of llms. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 7174–7193. find
  24. Daking Rai, Yilun Zhou, Shi Feng, Abulhair Saparov, and Ziyu Yao. 2024. A practical review of mechanistic interpretability for transformer-based language models. arXiv preprint arXiv:2407.02646. find
  25. Jie Ren, Qipeng Guo, Hang Yan, Dongrui Liu, Quanshi Zhang, Xipeng Qiu, and Dahua Lin. 2024. Identifying semantic induction heads to understand in-context learning. In Findings of the Association for Computational Linguistics ACL 2024, pages 6916– 6932. find
  26. Mansi Sakarvadia, Aswathy Ajith, Arham Khan, Daniel Grzenda, Nathaniel Hudson, André Bauer, Kyle Chard, and Ian Foster. 2023. Memory injections: Correcting multi-hop reasoning failures during inference in transformer-based language models. In Proceedings of the 6th BlackboxNLP Workshop: Analyzing an… find
  27. Alessandro Stolfo, Yonatan Belinkov, and Mrinmaya Sachan. 2023. A mechanistic interpretation of arithmetic reasoning in language models using causal mediation analysis. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pages 7035–7052. find
  28. Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, and 1 others. 2024. Gemma: Open models based on gemini research and technology. arXiv preprint arXiv:2403.08295. find
  29. Ian Tenney, Dipanjan Das, and Ellie Pavlick. 2019. Bert rediscovers the classical nlp pipeline. arXiv preprint arXiv:1905.05950.
  30. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. Advances in neural information processing systems, 30.
  31. Ben Wang and Aran Komatsuzaki. 2021. GPT-J-6B: A 6 Billion Parameter Autoregressive Language Model. https://github.com/kingoflolz/mesh-transformer-jax. ↗
  32. Boshi Wang, Xiang Yue, Yu Su, and Huan Sun. 2024. Grokked transformers are implicit reasoners: A mechanistic journey to the edge of generalization. In ICML 2024 Workshop on Mechanistic Interpretability. find
  33. Kevin Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. 2022. Interpretability in the wild: a circuit for indirect object identification in gpt-2 small. arXiv preprint arXiv:2211.00593. find
  34. Zhengxuan Wu, Atticus Geiger, Thomas Icard, Christopher Potts, and Noah Goodman. 2023. Interpretability at scale: Identifying causal mechanisms in alpaca. Advances in neural information processing systems, 36:78205–78226.
  35. Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. 2023. Efficient streaming language models with attention sinks. arXiv preprint arXiv:2309.17453. find
  36. An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, and 1 others. 2025. Qwen3 technical report. arXiv preprint arXiv:2505.09388.
  37. Sohee Yang, Elena Gribovskaya, Nora Kassner, Mor Geva, and Sebastian Riedel. 2024. Do large language models latently perform multi-hop reasoning? In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 10210–10229, Bangkok, Thailand.… find
  38. Wei Zhang, Chaoqun Wan, Yonggang Zhang, Yiu-Ming Cheung, Xinmei Tian, Xu Shen, and Jieping Ye. 2024. Interpreting and improving large language models in arithmetic calculation. In International Conference on Machine Learning, pages 59932–59950. PMLR. find
  39. Tianyi Zhou, Deqing Fu, Vatsal Sharan, and Robin Jia. 2024. Pre-trained large language models use fourier features to compute addition. arXiv preprint arXiv:2406.03445. find
  40. A + B − C A + B + C Heads Removed Accuracy Heads Removed Accuracy 59 Least Important L15H31 L16H1 L15H13 L16H21 L15H3 - 98.8% 90.6% 54.8% 6.6% 2.0% 0.8% - 58 Least Important L16H1 L15H3 L16H2 L16H3 L16H21 L15H13 98.6% 88.0% 57.8% 27.0% 5.6% 3.2% 2.4% find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/mamidanna-2025-all-one-llms-solve-mental/