揭示大语言模型中的数学推理：内部机制的方法学研究

# 大语言模型中数学推理的解开：内部机制的方法论调查
来源：https://arxiv.org/html/2604.15842
Tanja Baeumel1, 2, 3Josef van Genabith1, 2Simon Ostermann1, 2, 31德国人工智能研究中心 (DFKI) 2萨尔兰德大学 3欧洲可信人工智能研究中心 (CERTAIN) tanja\.baeumel@dfki\.de

###### 摘要

大语言模型（LLM）已展示了令人印象深刻的能力，但它们在处理推理密集型任务时的内部机制仍未得到充分探索。为了推进对模型内部处理机制的理解，我们通过检查任务执行期间的内部机制，研究了LLM如何执行算术运算。使用早期解码，我们追踪了如何在各层中构造下一个token的预测。我们的实验表明，虽然模型很早就识别了算术任务，但正确结果的生成仅在最后几层出现。值得注意的是，算术能力强的模型在注意力和MLP模块之间表现出清晰的劳动分工，其中注意力模块传播输入信息，而MLP模块聚合信息。这种分工在能力较弱的模型中并不存在。此外，成功的模型似乎以函数方式处理更具挑战性的算术任务，这表明推理能力超越了事实回忆。

## 1 引言

大语言模型（LLM）日益令人印象深刻的能力正在引发人们对促进其卓越性能的底层机制的浓厚兴趣。最近在可解释性方面的工作使研究界能够逐步加深对transformer模型如何检索和使用隐式存储在参数中的事实信息的理解。然而，我们对LLM用来解决非事实和推理密集型任务的机制的理解还不够深入。

在本研究中，我们研究了大语言模型（LLM）解决数学推理的机制。我们专注于一项既易于评估又易于解释的任务：具体来说，我们研究了LLM如何执行基本算术运算（例如，请计算 143 + 81 =），并分析了表现出不同算术能力水平的模型的内部机制差异。

我们通过早期解码的可解释性方法研究模型内部机制：我们观察模型的下一个token预测在各层中的构造方式，通过在每个注意力和MLP模块更新后对残差流进行反嵌入（即将其映射到单词），揭示了各个模块对结果生成过程的贡献。

请参考图1：数学推理期间模型内部机制的可视化。注意力块以红色显示，MLP块以蓝色显示。我们进行了三组实验来理解（1）任务如何以及何时被识别以及结果何时生成，（2）输入信息如何以及何时通过网络各层传播，（3）操作数如何结合，通过查看具有2个和3个操作数的设置。我们的主要发现如图1所示：

- 我们展示了LLM在早期层识别需要执行算术任务，但正确输出的生成发生在后期。
- 我们找到了强有力的证据表明LLM以函数方式解决算术任务，其中一个操作数充当参数，另一个操作数根据运算符和"参数操作数"进行改变，表明处理能力超越了事实回忆。
- 我们展示了算术能力强的模型在注意力和MLP模块之间表现出清晰的任务分工：与以前的工作类似，我们发现MLP模块聚合信息，而注意力模块传播信息。这种任务分工在算术能力较弱的模型中不存在。

## 2 方法

为了研究仅解码器的变压器语言模型用于接近和解决算术任务的机制，我们通过早期解码将残差流后的每个注意力和MLP模块投影到词汇空间，观察模型的下一个token预测在各层中是如何构造的，从而产生中间预测。

这使我们能够理解该模块当前关注的token，从而揭示各个模块对解决任务的贡献。我们在两个解码器仅有的语言模型上分析简单算术任务的中间预测，这两个模型的算术能力差异很大，以便理解什么机制使模型能够为算术任务生成正确结果。

请参考图2：早期解码的可视化。残差流由LM头进行解码，一次是在注意力输出添加到残差流之后（后注意力预测），一次是在MLP输出添加之后（后MLP预测）。

### 2.1 早期解码

我们用来研究LLM用于解决算术任务的机制的方法最初被称为"logit lens"。它允许深入了解transformer模型如何在生成过程中更新下一个token预测，方法是将预测token的残差流投影到中间层的词汇空间。

变压器LLM中的每个注意力和MLP组件以当前最终token表示（即残差流）作为输入。在流中，每个MLP和注意力组件的输出通过将模块的输出添加到输入表示来更新。第k层的下一个token预测的表示因此包括前一层模块对预测token表示所做的所有先前加法更新。

语言建模头（LM头），即线性预测层，用于LLM根据最后一个token表示生成下一个token的概率分布。该LM头也可以应用于中间层的中间表示。这有效地实现了一个反嵌入机制，允许研究生成过程中每个步骤的最可能的中间预测。因此，这样的中间预测允许追踪各个模块对表示所做的更改，并确定各个模块对结果生成过程的贡献。该方法在图2中可视化。

### 2.2 任务和数据

为了为观察LLM在算术任务上的行为创建一个受控环境，我们生成一个人工数据集，其中包含算术任务，我们根据运算符（求和，以下称为add，和减法，以下称为sub）、操作数大小、操作数数量和结果大小来控制。

我们使用"Please calculate operand 1 ∘ operand 2 ="类型的查询提示模型，其中operand1, operand2 ∈ ℕ 且 ∘ ∈ {+, -}，例如"Please calculate 306 + 136 ="，正确的回应是"442"。

对于每个运算符∘ ∈ {+, -}，我们创建一个包含较小数字的数据集，即操作数和结果，add_small 和 sub_small，以及一个包含较大数字的数据集 add_large 和 sub_large。对于小数据集，操作数和结果都≤99。在数据集创建中，我们确保所有操作数和结果都是整数≤520。我们选择这个上限是因为这确保使用了单个token的数字表示：GPT-2 XL和GPT-NeoX-20B分词器的词汇表包含0到520之间所有整数的单个token。较高的数字可能被编码为多个token。每个数据集包含500个唯一查询。所有实验都以零射击方式完成，没有模型的训练或适配。

在本文的其余部分，我们专注于对add数据集的评估，因为我们发现两个操作数的数据集的结构相似结果。我们在必要时报告加法和减法之间的差异。对sub数据集的完整评估在附录中提供。

### 2.3 模型

我们使用两个仅解码器的变压器语言模型进行实验：GPT-NeoX-20B 和 GPT-2 XL。我们在Huggingface上使用自由提供的EleutherAI/gpt-neox-20b和openai-community/gpt2-xl变体。我们确认了之前关于GPT-2 XL无法解决简单算术任务的发现，而GPT-NeoX-20b表现良好（附录中的表2）。因此，在本文的其余部分，我们专注于GPT-NeoX-20b。不过，我们在附录中展示了GPT-2 XL的结果，因为与GPT-NeoX-20b相比的内部机制差异很有趣。

## 3 实验集1：任务识别和结果生成

通过调查MLP和ATTN模块之后的中间预测（IP），我们首先检查两个基本问题：模型何时识别它需要执行算术任务？以及它何时开始生成结果？我们主要在后MLP预测中找到这些问题的答案：

- 模型很早就识别出它需要解决数值任务，并在中间层之前考虑不特定的数值token，在那里加载操作数（第3.1节）
- 正确的结果仅在最后一层预测（第3.2节）。

### 3.1 任务何时被识别？

#### 模型早期预测数值token。

我们观察不同层中后MLP和后ATTN IP中数值token的概率质量。图3显示了add_large数据集上分配给数值token的平均概率质量。在后MLP IP中，大约第9层出现了数值预测比例的急剧增加，这可能表明模型开始识别任务的数值性质。我们在add_small数据集中发现了类似的一般趋势。

#### 数值token仅在最后一层被高置信度预测。

图4显示了add_large数据集上后MLP中数值token在前1和前10预测中的平均比例。对于早-中层（第9至17层）和中-后层（第30至43层），后MLP中前10个预测的数值预测数量在10-40%之间。但是，模型仅在最后一层（第44层）将数值token预测为首选预测。

#### 注意力层的行为差异显著。

在注意力模块之后观察到的数值预测的模式非常不同。我们在特定层观察到数值预测的显著峰值，特别是在第9、12、14和21层，平均而言，概率质量的87%-99%在数值token上。以前的工作表明注意力模块负责在位置之间传播信息，因此我们推测这些峰值可能表明重要的数值信息可能在最终token位置被注意到或通过突出的注意力模块传播到最终token位置。我们在第4.1节更详细地研究这些发现。

请参考图3(a)和(b)：平均在add_large数据集上的（a）后ATTN和（b）后MLP中间预测中数值token的组合概率质量。

请参考图4(a)和(b)：平均在add_large数据集上的（a）前1和（b）前10个后MLP中间预测中数值token的比例。

### 3.2 正确结果何时生成？

#### 数值预测到后期层与正确结果无关。

为了理解何时生成结果，我们检查数值IP与整个层的正确结果的相似性。我们分析了绝对误差，定义为预测数字与正确数字之间的差异，在前10和前1个预测的token中（图5）。我们的分析揭示了几个关键发现：首先，在后MLP IP中，预测的数值token与正确结果之间的相似性在早期和中期层中通常较低，但在第35层之后开始逐步增加，这对应于正确结果被分配更高概率的层（图6）。在第28层中，后ATTN预测的前1和前10个的绝对误差明显减少，表明模型在该层对正确结果的量级有了合理的近似。

请参考图5(a)-(d)：（a）前1和（b）前10个后ATTN中间预测中数值token的绝对误差（即与正确结果的差异）。（c）前1和（d）前10个后MLP中间预测中数值token的绝对误差。所有结果平均在add_large数据集上。

#### 正确结果仅在非常后期的层中生成。

我们还分析了在哪一层生成正确结果，并评估了结果token的概率和其在每一层最可能token中的位置。

正确结果通常在模型的后期层出现（图6）。这由结果token的概率显著增加和其位置相应减少来证明，对于大加法任务，正确结果在第35-40层出现，在最后第44层中概率显著增加和排名减少。对于较小的加法任务，正确结果token早至第26层出现，并在第32-34层左右被锚定为最高预测token。因此，正确结果在更容易的任务生成过程中比更复杂的任务出现更早。这可能表明对更容易的任务与更具挑战性的数学推理任务采用不同的内部机制。

请参考图6(a)和(b)：在所有数据点平均的中间层后MLP预测中正确结果的位置，分别为（a）add_small和（b）add_large。

## 4 实验集2：输入传播

我们在前一部分的发现提供了对模型在何处识别任务和结果出现的见解。但是，执行这些