AI大语言模型解数学题或仅依赖最后一个token

Alexin大语言模型解数学题或仅依赖最后一个token：近日，一项研究揭示，大语言模型在处理心算类数学题时，几乎所有实际数学计算都集中在序列的最后一个token上完成，而非分散在所有token中。来自加州大学圣克鲁兹分校、乔治·梅森大学和Datadog的研究团队发现，这一发现挑战了Transformer架构模型常见的信息访问模式。

研究人员运用上下文感知平均消融（Context - Aware Mean Ablation, CAMA）和基于注意力的窥视（attention - based peeking, ABP）技术，对Llama - 3 - 8B等基于Transformer架构的大语言模型开展了一系列消融实验。这些实验通过系统性地移除或改变模型部分组件，探究模型完成心算任务所需的“最少计算量”。

在实验过程中，研究人员发现模型内部会形成一种被称为“人人为我（All - for - One, AF1）”的稀疏子图——该子图通过最少的计算层和最有限的信息传递，让模型高效完成运算。在这种机制中，输入Transformer前几层的token主要负责通用准备工作（如识别token、结构编码、预测下一步所需的通用表示），而中间的两层则将信息传递给最后一个token。之后，最后一个token独自完成计算并给出答案。这一过程表明，模型内部将“任务通用型计算”（如token识别、数值与结构编码）与“输入特定型计算”（如实际算术运算）是分开处理的。

研究聚焦与验证：该研究聚焦于心算任务（即涉及两个或三个操作数的算术问题，如42 + 20 − 15等，这些问题可以通过单个token的输出解决，无需模型进行显式的链式思维推理）。实验表明，大语言模型在诸多计算任务上表现出色，而其Transformer架构允许任意token通过自注意力机制立即访问所有先前的token以传递信息，并使每个token能够通过多层感知机（MLP）并行执行各自独立计算。但即便如此，模型内部的信息流和计算过程仍是不透明的。因此，研究人员通过“在模型的初始层抑制token针对特定输入的计算”“在少数几层中限制跨token位置的信息传递路径”“在剩余的层中强制所有计算都在最后一个token上发生”等步骤，来探索模型内部机制。

通过一系列实验，研究人员发现Llama - 3 - 8B等模型在心算任务中只需前14层做任务通用计算，然后通过2层信息传输让最后一个token获取全局信息，剩余层仅进行最后token的自计算。这个几乎保留全部性能的子图被命名为AF1 - llama。进一步验证表明，AF1 - llama在八个任务中总体表现出高忠实度，且关键注意力头集中在少数层，大部分头部冗余。此外，研究还验证了AF1 - llama在Llama - 3 - 8B和Llama - 3.1 - 8B上的泛化情况，以及在Pythia和GPT - J模型上的类似稀疏子图。

研究意义：总体而言，这项工作为大语言模型中的算术推理和跨token计算的机制理解做出了重要贡献，也为优化大语言模型计算效率提供了新思路。此外，CAMA和ABP这类研究方法也可服务于算术任务之外的更广泛应用。

（注：相关研究论文及链接见文中标注。）