AI解数学题只靠最后一个token

tayun

henry 发自 凹非寺
量子位 | 公众号 QbitAI

大语言模型在解心算题时，只依赖最后一个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+2015），这些问题可以通过单个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和AB运算的其他算术形式，研究进一步将口头描述运算以及将运算嵌入到应用题或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


来源：网易