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

📄 arXiv: 2509.09650v1 📥 PDF

作者: Siddarth Mamidanna, Daking Rai, Ziyu Yao, Yilun Zhou

分类: cs.CL

发布日期: 2025-09-11

备注: EMNLP 2025 Main Conference

💡 一句话要点

LLM通过末位Token计算解决心算问题，信息由其他Token传递

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 心算 可解释性 信息传递 注意力机制

📋 核心要点

  1. 现有方法难以解释LLM在心算任务中的内部计算机制和信息传递方式。
  2. 提出Context-Aware Mean Ablation (CAMA)和Attention-Based Peeking (ABP)两种技术，用于识别关键计算子图。
  3. 实验表明，All-for-One (AF1)子图对于LLM执行心算任务至关重要，且具有跨模型迁移能力。

📝 摘要（中文）

大型语言模型（LLM）在众多计算任务中表现出卓越的能力，但其内部运作机制仍不明确。理论上，因果自注意力机制和多层感知机层的结合使得每个token都能访问并计算基于所有先前token的信息。实际上，这种操作在多大程度上存在？本文针对心算任务（即通过下一个token预测进行直接数学计算，无需显式推理）分三步研究了这个问题：抑制初始层中特定于输入的token计算，限制接下来几层中跨token位置的信息传递路径，并强制所有计算在剩余层中的最后一个token处发生。通过提出的两种技术，上下文感知均值消融（CAMA）和基于注意力的窥视（ABP），我们识别出一个All-for-One子图（AF1），该子图在各种心算任务上具有很高的准确性，其中有意义的计算发生在非常晚的阶段（就层深度而言），并且仅在最后一个token处发生，该token接收来自特定中间层中其他token的信息。对各种模型和算术表达式的实验表明，该子图对于高模型性能是充分且必要的，可以在不同模型之间转移，并且适用于各种输入样式。对不同CAMA和ABP替代方案的消融实验揭示了它们相对于其他方法的独特优势，这些优势可能具有独立的意义。

🔬 方法详解

问题定义：论文旨在探究大型语言模型（LLM）在执行心算任务时，信息是如何在不同token之间传递和计算的。现有方法难以解释LLM内部的运作机制，特别是如何利用上下文信息进行计算，以及哪些token在计算中起关键作用。现有方法缺乏对LLM内部计算过程的细粒度控制和分析能力。

核心思路：论文的核心思路是通过逐步限制和控制LLM中不同层的信息传递和计算方式，从而识别出对心算任务至关重要的计算子图。具体来说，通过抑制初始层的输入特定计算，限制中间层的信息传递路径，并强制所有计算在最后一层进行，来观察模型性能的变化，从而推断出关键的计算路径和token。

技术框架：论文的技术框架主要包括三个阶段：1) 抑制初始层中特定于输入的token计算；2) 限制中间层中跨token位置的信息传递路径；3) 强制所有计算在剩余层中的最后一个token处发生。在这个框架下，论文提出了两种关键技术：Context-Aware Mean Ablation (CAMA) 和 Attention-Based Peeking (ABP)。CAMA用于抑制特定token的计算，而ABP用于控制token之间的信息传递。

关键创新：论文最重要的技术创新点在于提出了CAMA和ABP两种技术，并利用它们识别出了All-for-One (AF1) 子图。AF1子图表明，LLM在执行心算任务时，大部分计算发生在最后一层，并且最后一个token接收了来自其他token的信息。与现有方法相比，该方法能够更精确地定位LLM中负责特定任务的关键计算模块。

关键设计：CAMA的关键设计在于使用上下文感知的均值来替代被消融的token的表示，从而避免引入额外的噪声。ABP的关键设计在于使用注意力机制来控制token之间的信息传递，允许特定token“窥视”其他token的信息。论文还详细描述了如何选择合适的层进行CAMA和ABP操作，以及如何评估AF1子图的性能。

📊 实验亮点

实验结果表明，All-for-One (AF1) 子图对于LLM执行心算任务至关重要，移除该子图会导致模型性能显著下降。AF1子图具有跨模型迁移能力，即在一个模型上识别出的AF1子图可以成功应用于其他模型。CAMA和ABP技术在消融实验中表现出优于其他方法的性能。

🎯 应用场景

该研究成果可应用于提升LLM的可解释性和可控性，帮助理解LLM如何进行数值计算和推理。此外，该方法可以用于优化LLM的结构，例如通过减少不必要的计算和信息传递，从而提高效率和降低计算成本。该研究对于开发更高效、更可靠的LLM具有重要意义。

📄 摘要（原文）

Large language models (LLMs) demonstrate proficiency across numerous computational tasks, yet their inner workings remain unclear. In theory, the combination of causal self-attention and multilayer perceptron layers allows every token to access and compute information based on all preceding tokens. In practice, to what extent are such operations present? In this paper, on mental math tasks (i.e., direct math calculation via next-token prediction without explicit reasoning), we investigate this question in three steps: inhibiting input-specific token computations in the initial layers, restricting the routes of information transfer across token positions in the next few layers, and forcing all computation to happen at the last token in the remaining layers. With two proposed techniques, Context-Aware Mean Ablation (CAMA) and Attention-Based Peeking (ABP), we identify an All-for-One subgraph (AF1) with high accuracy on a wide variety of mental math tasks, where meaningful computation occurs very late (in terms of layer depth) and only at the last token, which receives information of other tokens in few specific middle layers. Experiments on a variety of models and arithmetic expressions show that this subgraph is sufficient and necessary for high model performance, transfers across different models, and works on a variety of input styles. Ablations on different CAMA and ABP alternatives reveal their unique advantages over other methods, which may be of independent interest.