大语言模型在解心算题时，只依赖最后一个 token？

最近，来自加州大学圣克鲁兹分校、乔治 · 梅森大学和 Datadog 的研究人员发现：在心算任务中，几乎所有实际的数学计算都集中在序列的最后一个 token上完成，而不是分散在所有 token 中。

这意味着，相较于在 Transformer 和多层感知机（MLP）中常见的全局信息访问——即每个 token 在预测时都能查询并利用整个上文信息——在诸如心算这样的特定任务中，全局访问其实并不是必需的。

这是怎么一回事？

心算只要最后一个 token？！

总的来说，研究人员采用了上下文感知平均消融（Context-Aware Mean Ablation, CAMA）和基于注意力的窥视（attention-based peeking）技术对Llama-3-8B等 Transformer 架构的模型进行了一系列的消融实验。

这些实验通过系统性地移除或改变模型的一部分，探究能让模型依然表现良好的" 最少计算量 "。

在这一过程中，研究人员发现模型内部会形成一个稀疏子图（sparse subgraph）——他们把它称为" 人人为我 "（All-for-One, AF1）。

这个子图通过最少的计算层和最有限的信息传递，让模型高效完成运算。

在 " 人人为我 " 中，输入 Transformer 前几层（L_wait）的 token 并没有做跟 " 自己数值 " 相关的计算，而是 " 等待 "，并主要承担一些通用的准备工作（比如识别 token、结构编码、预测下一步所需的通用表示）。

然后，在中间的两层（L_transfer）里，它们就将信息传递给最后一个 token。

之后，最后一个 token 独自完成计算并给出答案。

这一过程表明，模型内部将任务通用型计算（如 token 识别、数值与结构编码）与输入特定型计算（如实际算术运算）是分开的。

（注：这篇研究聚焦于心算任务，即涉及两个或三个操作数的算术问题（例如 42+20 − 15），这些问题可以通过单个 token 的输出解决，而无需模型进行显式的链式思维推理。）

接下来，我们具体来看。

众所周知，大语言模型在许多计算任务上表现出色，而其中一个重要原因是其采用了 Transformer 架构。

与 RNN 不同，Transformer 允许任意 token 通过自注意力机制立即访问所有先前的 token 以传递信息，并使每个 token 能够通过多层感知机（MLP）并行执行各自的独立计算。

但即便如此，模型内部的信息流和计算过程仍然是不透明的。

因此，为了揭开大语言模型的 " 黑箱 "，研究人员采用了以下三个步骤来进行探索。

首先，在模型的初始层抑制 token 针对特定输入的计算。

研究人员发现，在传统 Transformer 的每一层中，token 都能访问所有之前的 token，但对于简单的心算任务，每个 token 可能未必从一开始就要获得全局信息。

由此，研究人员引入了等待期（L_wait）：让在前面的 L_wait 层中的 token 独立计算，只执行任务通用操作（如理解数字、识别算术结构），而不访问其他 token。

为了实现这一点，他们使用了上下文感知平均消融（CAMA）。

CAMA 的作用是屏蔽掉 token 之间的输入特定信息，同时保留每个 token 的普遍计算能力，使模型能够在不依赖具体输入的情况下完成基础准备工作。

接下来，在少数几层中限制跨 token 位置的信息传递路径。只让最后 token 在 L_transfer 层中访问所有 token，其余层只关注自己。

最后，在剩余的层中强制所有计算都在最后一个 token 上发生。

由于 CAMA 只能从第一层开始，因此，研究人员引入了一种可以在任意层控制信息访问的方法——基于注意力的窥视（ABP）。

它通过修改注意力掩码（attention mask），精确指定每个 " 查询 "（query）token 可以关注哪些 " 键 "（key）。

在这篇论文中，研究人员主要使用了以下两种模式：

完全窥探 ( Full-peeking ) :   token 可以关注所有在它之前的 token，这是标准的因果注意力。在 AF1 的传递阶段，最后一个 token 使用此模式来收集信息。

自我窥探 ( Self-peeking ) :   token 只能关注它自己，在传递和计算阶段，所有非末尾的 token 都使用此模式；在计算阶段，最后一个 token 也切换到此模式。

实验验证

在完成方法和操作流程的构建后，研究者进行了一系列实验来发现、验证和分析 AF1 子图。这里主要涉及到 Llama-3-8B 和 Llama-3.1-8B，以及在 Pythia 和 GPT-J 模型上的验证。

首先，通过三阶段消融与窥视实验，研究人员发现 Llama-3-8B 在 A+B+C 任务中只需前 14 层做任务通用计算（CAMA 层），然后通过 2 层信息传输让最后的 token 获取全局信息，剩余层仅进行最后 token 的自计算。

这个几乎保留全部性能的子图被命名为AF1_llama。

接下来，研究人员又进一步验证了 AF1_llama 在 Llama-3-8B 和 Llama-3.1-8B 上的表现。

实验表明，AF1_llama 在八个任务中总体表现出高忠实度。

更进一步，实验进一步验证了第 15 和 16 层的信息传输在 Llama-3-8B 中的重要性。

研究表明，仅少数注意力头对算术计算关键，即使移除近 60 个头部，模型仍能保持约 95% 的准确率，表明大部分注意力头冗余，而关键头集中在少数层。

此外，为了探究 AF1_llama 是否可以在 Llama-3-8B 上泛化到表示 A+B 和 A − B 运算的其他算术形式，研究进一步将口头描述运算以及将运算嵌入到应用题或 Python 代码中。

实验表明，AF1_llama 在不包含额外语义上下文的直接算术任务中仍保持了相当高的准确率。

然而，它在需要语义理解的任务上，如应用题和 Python 代码，完全失败了，这表明它需要额外的组件来处理其他能力，比如理解自然语言或 Python 程序输入。

最后，研究人员在 Pythia 和 GPT-J 中也发现了类似 AF1 的子图，但与 Llama 不同，这些模型的等待期更短（L_wait ≈ 9 – 11）、信息传输层更长，且性能边界不如 Llama 清晰。

尽管忠实度普遍低于 Llama，但对二元运算任务的子图仍能恢复超过一半的原始模型准确率。

总体而言，这项工作为大语言模型中的算术推理和跨 token 计算的机制理解做出了贡献。此外，它通过 CAMA 和 ABP 提供了方法论上的创新，可服务于算术任务之外的更广泛应用。

参考链接

[ 1 ] https://x.com/rohanpaul_ai/status/1966788202238448101

[ 2 ] https://www.arxiv.org/pdf/2509.09650

一键三连「点赞」「转发」「小心心」

欢迎在评论区留下你的想法！

—  完  —

  年度科技风向标「2025 人工智能年度榜单」评选报名开启啦！我们正在寻找 AI+ 时代领航者  点击了解详情

❤️‍   企业、产品、人物 3 大维度，共设立了 5 类奖项，欢迎企业报名参与   

一键关注 点亮星标

科技前沿进展每日见