1. 引言

当今人工智能领域最激烈的争论之一围绕着大型语言模型 (LLM) 的本质展开。当像 GPT-4 这样的模型解决复杂的逻辑谜题时,它是在真正地进行推理——应用逻辑规则推导出答案吗?还是仅仅扮演了一只“随机鹦鹉”,从其庞大的训练数据中检索记忆下来的模式?

尽管许多研究集中在通过数据质量或模型规模来提高推理能力,但一项名为 “An Analysis for Reasoning Bias of Language Models with Small Initialization” (小初始化语言模型推理偏置分析) 的迷人新研究关注了一个基本却常被忽视的设计选择: 参数初始化

研究人员提出了一个惊人的发现: 我们初始化模型权重的规模可以从根本上改变其学习个性。具体来说, 较小的初始化规模会使模型倾向于推理 , 而较大 (标准) 的规模则会推动模型倾向于记忆。

为了直观地展示这一点,请看一个 GPT-2 模型的训练动态,该模型在混合了两个数据集的数据上进行训练: PrOntoQA (一个逻辑推理数据集) 和 TinyStories (一个简单的叙事数据集) 。

Figure 1. Comparison of training loss between PrOntoQA and TinyStories in one next-token prediction training for this mix dataset. The red line represents the training loss on the PrOntoQA dataset, while the blue line depicts the training loss on the TinyStories dataset.

如图 1 所示, (使用小权重初始化的) 模型学习推理任务 (PrOntoQA,红线) 的速度明显快于叙事记忆任务 (TinyStories,蓝线) 。这表明当参数从较小的值开始时,训练动态中存在一种内在的“推理偏置”。

在这篇文章中,我们将解构这篇论文以理解发生这种情况的原因。我们将探索嵌入的几何结构、梯度流的数学原理,以及标签的分布如何在模型完成训练之前就塑造了它的“思维”。

2. 背景: 设定

为了严格测试推理与记忆,我们不能仅依赖混乱的自然语言。我们需要一个受控环境。作者利用基于“锚函数 (Anchor Functions) ”的合成任务框架。

2.1 合成组合任务

核心思想是创建充当算术问题的 Token 序列。模型看到一个序列并必须预测一个标签。序列包含:

  1. 键 (Keys, \(z\)) : 变量或起点。
  2. 锚点 (Anchors, \(a\)) : 修饰符或操作。
  3. 噪声 (Noise) : 不相关的 Token。

输入序列 \(X\) 的一般结构如下:

Equation defining the sequence structure with keys and anchors.

这里,\(z\) 代表键,\(a\) 代表锚点。研究人员根据这些序列定义了两种截然不同的映射 (任务) :

1. 推理映射 (The Reasoning Mapping, \(\mathcal{F}_{rsn}\)) 这个任务要求模型学习一条规则。具体来说,标签是键和锚点的总和。如果模型学会了加法规则,它就可以泛化到未见过的数字。

Equation for Reasoning Mapping: z_p plus the sum of anchors.

2. 记忆映射 (The Memory Mapping, \(\mathcal{F}_{mem}\)) 这个任务强制模型进行记忆。对于特定的键和锚点对,标签是一个随机分配的数字 \(y\)。输入和输出之间没有逻辑规则连接;这纯粹是一个查找表。

Equation for Memory Mapping: random label assignment.

2.2 可视化数据

数据集的构建方式使得模型看到的序列在结构上看起来完全相同,但需要不同的认知过程 (加法 vs. 查找) 来解决。

Figure 2. Schematic diagram of the synthetic composition task.

在图 2 中,左侧显示了记忆映射 。 橙色 (键) 和绿色 (锚点) Token 的特定组合映射到一个任意目标 \(y\)。右侧的推理映射显示目标是通过数学推导出来的 (\(78+15=93\)) 。

关键的实验设置是在包含这两种类型任务的数据集上训练模型,并观察它更倾向于学习哪一种,以及它的泛化能力如何。

3. 核心现象: 初始化规模

深度学习模型通常使用从正态分布 \(\mathcal{N}(0, \sigma^2)\) 中抽取的随机权重进行初始化。标准差 \(\sigma\) 通常定义为 \(d^{-\gamma}\),其中 \(d\) 是层宽,\(\gamma\) (gamma) 决定了规模。

  • 小初始化 (Small Initialization) : 大 \(\gamma\) (例如 \(\gamma=0.8\)) 。权重非常接近于零。
  • 大/标准初始化 (Large/Standard Initialization) : 小 \(\gamma\) (例如 \(\gamma=0.3\) 或 \(0.5\)) 。权重分布更分散。

3.1 权衡

研究人员使用不同的初始化率在这些合成任务上训练了 Transformer。结果如图 3 所示,令人震惊。

Figure 3. Loss and prediction accuracy of the models on different datasets under varying initialization scales.

让我们拆解图 3 (A 面板) :

  • \(\gamma = 0.3\) (大初始化,左栏) : 蓝线 (记忆) 和紫线 (推理训练集) 迅速下降。然而,橙线 (推理测试集 )保持高位。这是典型的过拟合。模型记住了两个任务的训练数据,但未能学会推理规则。
  • \(\gamma = 0.8\) (小初始化,右栏) : 紫线 (推理训练集) 和橙线 (推理测试集) 一起下降,而且它们的下降速度比蓝线 (记忆) 更具快

结论: 小初始化抑制了记忆,促进了可泛化规则的学习。当权重从很小开始时,模型“更喜欢”寻找底层逻辑 (\(z+a\)) ,而不是记忆个别的结果。

4. 机制: 为什么小初始化有利于推理

要理解为什么会发生这种情况,我们必须深入底层。作者将 Transformer 简化为一个名为 Emb-MLP (嵌入层 + 多层感知机) 的模型,以便从数学上分析梯度流。

4.1 嵌入空间

语言模型的“大脑”是它的嵌入空间——在这里 Token 被转换为向量。这个空间的几何形状决定了模型觉得什么容易学。

当初始化很小时,嵌入向量 \(w^{emb}\) 很微小。在反向传播过程中,这些向量的增长方式在很大程度上取决于与每个 Token 相关的标签分布

特定 Token \(s\) 的梯度流 (权重如何变化) 可以用以下方程近似:

Gradient flow equation for embeddings.

该方程表明,Token 嵌入的变化与该 Token 相关的平均标签 (\(P^s\)) 成正比。

4.2 标签分布: 关键区分点

这是论文的关键见解。

  1. 记忆任务: 标签是随机的。对于任何给定的记忆锚点,相关标签在可能的输出中均匀分布。
  2. 推理任务: 标签是结构化的。对于一个推理锚点 (比如加法任务中的数字“5”) ,相关标签相比于键正好偏移了 5。

因为记忆标签是随机且均匀的,梯度方程中的“平均标签”项对于所有记忆 Token 往往趋向于一个均匀常数。它们在梯度看来都是一样的。

然而,推理标签具有偏移的均值。Token “5”的平均标签与 Token “10”的平均标签是不同的。

结果是什么? 推理 Token 在训练早期就形成了独特的、结构化的嵌入,而记忆 Token 仍然聚集在一起,无法区分。

4.3 可视化嵌入

我们可以在训练期间嵌入的余弦相似度矩阵中清楚地看到这种区别。

Figure 4. Cosine similarity matrices for memory and reasoning anchors.

在图 4 (A 面板) 中:

  • 顶行 (记忆锚点) : 热图大部分是红/黄色的,表示不同的记忆锚点之间具有高相似度 (接近 1.0) 。模型无法轻易区分它们。
  • 底行 (推理锚点) : 我们看到了漂亮的对角线模式。当你远离对角线时,相似度会下降。这意味着数字“11”与“12”相似,但与“20”不同。

这种结构化的几何形状纯粹源于结构化标签上的梯度流。它在高维空间中创建了一个“数轴”表示。因为推理 Token 是截然不同的,模型可以轻松使用它们来计算输出。由于记忆 Token 难以区分,从而阻碍了记忆任务的学习。

5. 扩展到 Transformer

这个逻辑适用于 Transformer 的复杂架构吗?是的。

研究人员分析了使用小初始化训练的完整 Transformer 模型的嵌入空间。

Figure 5. Embedding structure of a Transformer model with small initialization scale.

图 5 证实了这一理论。面板 A (底部) 显示,Transformer 中的推理锚点自然地以分层方式组织自己。面板 B 显示了 PCA 投影,其中推理锚点 (绿色) 和键 (橙色) 形成了独特的、有序的结构,而记忆 Token 则聚集成一团。

5.1 注意力机制作为聚合器

在小初始化下,第一层的注意力矩阵表现出一种特定的方式: 它充当一个平均算子

Equation showing the attention output as an average.

由于权重很小,softmax 函数不会出现尖锐的峰值;它将注意力大致均匀地分散开来。这允许模型从整个上下文中聚合信息。

关键是,Value 投影矩阵 (\(W^V\)) 将自身与推理锚点对齐。

Figure 6. Characteristics of the first attention module under small initialization.

图 6 (D 面板) 显示,\(W^V\) 的奇异向量与推理锚点 (红线) 具有极高的余弦相似度,但与记忆锚点 (蓝线) 则不然。

机制总结:

  1. 嵌入: 小初始化 + 结构化标签 = 推理 Token 的独特嵌入。
  2. 注意力 1: 对上下文取平均。由于推理嵌入是独特的并且与投影矩阵对齐,它们的值被保留并传播。
  3. 注意力 2: 识别“键”并将它与聚合的“锚点”信息结合以计算结果 (\(Key + \sum Anchors\)) 。

记忆任务在这个过程中失败了,因为它们的嵌入在早期阶段从未充分分离以发挥作用。

6. 对比: “懒惰”的大初始化

如果我们使用标准 (大) 初始化会发生什么?

在高维空间中,从大尺度分布中抽取的随机向量往往彼此正交 (垂直) 。

Figure 14. Characteristic of embedding space of PrOntoQA and TinyStories with initialization rates 0.3 vs 0.5.

图 14 (顶行,\(\gamma=0.3\)) 展示了大初始化的情况。嵌入空间由正交性定义——仅仅出于偶然,每个 Token 都与其他 Token 不同。

  • 优点: 这对记忆很有利。如果每个输入都是唯一且正交的,那么很容易将输入 A 映射到输出 A,而不会干扰输入 B。
  • 缺点: 这对推理来说很糟糕。模型将“10”和“11”视为完全不相关的实体。它学不到它们之间的关系

这解释了为什么图 3 中的大初始化模型记住了数据,但在推理测试集上失败了。它们学会了利用权重的正交性将特定输入映射到输出,而不是学习加法规则。

7. 现实世界的影响

研究人员使用 PrOntoQA (推理) 和 TinyStories (叙事) 数据集在 GPT-2 上验证了这些发现。

Figure 7. Reasoning bias of GPT-2 in real language tasks.

图 7 显示,在小初始化 (\(\gamma=0.8\)) 下,PrOntoQA 数据集的嵌入 (B 中的左热图) 形成了具有相关性的丰富结构,而 TinyStories 的嵌入 (右热图) 在很大程度上保持不相关。

面板 A 显示了“推理偏置”指标 (\(\Delta L\)) 。随着 \(\gamma\) 增加 (初始化变得更小) ,模型越来越倾向于最小化推理损失而不是记忆损失。

8. 结论

这篇论文为一个强有力的想法提供了理论和实证基础: 我们可以通过初始化来控制 LLM 的“认知风格”。

  • 小初始化允许数据的标签分布塑造嵌入空间。如果任务具有结构 (逻辑、数学、语法) ,这种结构就会被烙印在模型上,从而促进推理和泛化。
  • 大初始化迫使模型进入正交状态。这有利于记忆任意映射,但阻碍了底层规则的发现。

对于学生和从业者来说,这表明如果你的目标是训练能够推理和泛化而不是产生幻觉或记忆训练数据的模型,那么关注初始化规模不仅仅是一个技术细节——这是一个基本的架构选择。

关键要点

  1. 初始化是一个偏置旋钮: 小权重 = 推理偏置;大权重 = 记忆偏置。
  2. 标签塑造嵌入: 在小权重状态下,Token 标签的统计分布驱动了嵌入空间的几何形状。
  3. 早期动态至关重要: 模型的“个性” (推理者 vs. 记忆者) 是在训练早期决定的,取决于嵌入能够多么容易地将自己区分开来。

本文基于论文 “An Analysis for Reasoning Bias of Language Models with Small Initialization” (2025)。