大模型算术题靠最后一字完成计算,信息从中间层传递而来。
All for One: LLMs Solve Mental Math at the Last Token With Information Transferred From Other Tokens
- 让模型只在最后一步计算,其他步骤屏蔽部分信息
- 实验发现只需最后一层单个词参与运算即可达到高准确率
- 适合研究模型内部机制或提升推理效率的读者
大型语言模型在诸多计算任务中表现优异,但其内在机制仍不清晰。理论上,因果自注意力与多层感知机允许每个词元访问并基于所有前置词元进行计算。实践中,这种能力究竟有多强?本文针对直接算术计算(即通过预测下一个词元完成数学计算,无显式推理过程)任务,分三步探究:在初始层抑制特定输入的词元计算,在后续几层限制信息传播路径,并强制剩余层的所有计算集中在最后一个词元。通过提出的两种技术——上下文感知均值消融(CAMA)和基于注意力的窥探(ABP),我们识别出一个‘全为一’子图(AF1),该子图在多种算术任务上表现优异,其中有效计算发生在极深层且仅集中于最后一个词元,它从若干特定中间层接收其他词元的信息。实验表明,该子图对高性能是充分且必要的,可跨模型迁移,适用于多种输入风格。对CAMA与ABP不同变体的消融分析揭示了其相对于其他方法的独特优势,可能具有独立研究价值。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。