跨越鸿沟: 双空间知识蒸馏如何统一教师与学生大模型

当前的人工智能时代是由“缩放定律”定义的。我们已经看到，增加大型语言模型 (LLM) 的参数数量始终能带来更好的泛化和推理能力。然而，这种对智能的追求伴随着昂贵的代价。像 LLaMA-70B 或 GPT-4 这样的模型体量巨大，这使得它们在现实场景中的部署成本极其高昂且运行缓慢。

这引发了人们对 知识蒸馏 (Knowledge Distillation, KD) 的极大兴趣。KD 的前提非常优雅: 利用一个庞大、智能的“教师”模型来训练一个更小、更快的“学生”模型。目标是在不损失太多性能的情况下，将教师的知识压缩到一个轻量级的包中。

然而，来自北京交通大学的研究人员指出了目前我们执行这种蒸馏操作中存在的一个关键缺陷。在他们的论文《面向大型语言模型的双空间知识蒸馏》 (Dual-Space Knowledge Distillation for Large Language Models) 中，他们认为标准方法忽略了教师和学生之间根本性的“空间差异”。

在这篇文章中，我们将深入探讨这项研究。我们将探讨为什么现有的方法无法有效地传递知识，提出的 双空间知识蒸馏 (DSKD) 框架如何通过统一模型的输出空间来解决这个问题，以及一种新颖的注意力机制如何让这一过程在模型说着不同“语言” (词表) 时依然有效。

现状: 白盒知识蒸馏

要理解这项创新，我们必须先了解基准方法。在标准的“白盒”KD 中，学生模型试图模仿教师的概率分布。

当 LLM 处理一段文本序列时，它会计算每一个可能的下一个 token 的概率。学生模型通常同时在两个目标上进行训练:

  1. 标准学习: 从数据集中预测正确的下一个词 (交叉熵损失) 。
  2. 蒸馏: 最小化其预测概率分布与教师分布之间的差异。

在数学上，标准语言建模损失如下所示:

标准交叉熵损失方程。

在这里，模型优化的是正确 token \(x_i^*\) 的概率。

同时，KD 损失强制学生模型 \(q_\theta\) 去匹配教师模型 \(p\):

标准知识蒸馏损失方程。

\(\mathcal{D}\) 代表距离函数，通常是 Kullback-Leibler (KL) 散度。其理念很简单: 如果教师认为“dog”的可能性是 80%，而“cat”的可能性是 20%，学生应该学习这种分布，从而捕捉到教师推理的细微差别。

隐藏的问题: 空间差异

研究人员认为，上述方程掩盖了一个重大问题。学生模型和教师模型是不同的模型。它们拥有不同的架构、不同的隐藏层维度，关键是，拥有不同的预测头 。

预测头是神经网络的最后一层，负责将内部隐藏状态转换为词表上的概率分布。由于学生和教师使用它们各自的预测头，它们的输出分布存在于不同的“空间”中。

作者假设，最小化来自不同空间的分布之间的距离是低效的。它迫使分布看起来相似，但并不一定迫使内部表示 (模型的“思想”) 对齐。

可视化失败案例

为了证明这一点，作者进行了一项模拟实验。他们创建了一个具有 2D 隐藏状态的合成教师和学生模型。他们运行了标准 KD (最小化 KL 散度) ，并将训练前后的内部隐藏状态进行了可视化。

一张展示模拟结果的组合图像。子图 (b) 显示，在使用不同预测头的情况下，学生和教师的表示仍然截然不同。子图 (c) 显示，共享预测头可以使它们完美对齐。子图 (d) 显示，共享预测头的损失收敛得更快。

仔细观察上面的 图 1 :

  • (b) KD 后 (不同头) : 这代表了目前的标准方法。即使经过训练，学生的表示 (红色三角形) 和教师的表示 (蓝色星形) 在结构上仍然是截然不同的。它们没有重叠。知识传递是不完整的。
  • (c) KD 后 (共享头) : 当作者强制模型共享同一个预测头 (统一输出空间) 时，表示完美对齐了。
  • (d) 损失曲线: 橙色线显示，与标准方法 (蓝线) 相比，共享输出空间能带来更快的收敛速度和更低的最终损失。

这项模拟证实了假设: 空间差异限制了教师和学生之间的相似性。

解决方案: 双空间知识蒸馏 (DSKD)

作者提出了一个新的框架: 双空间知识蒸馏 (DSKD) 。 其核心思想是将一个模型的隐藏状态投影到另一个模型的表示空间中，从而有效地让它们共享预测头。

DSKD 不是仅在其各自的空间中比较输出，而是在两个统一的空间中执行蒸馏: 学生空间和教师空间 。

1. 学生空间中的蒸馏

首先，我们希望将教师的知识带入学生的世界。

教师的隐藏状态 (\(\mathbf{h}^t\)) 通常比学生的维度更大。我们使用一个可训练的线性投影器 \(\mathcal{P}^{t \to s}\) 将教师的隐藏状态缩小到学生的维度 (\(d\)) 。

显示将教师隐藏状态投影到学生维度的方程。

一旦投影完成，我们将这些“学生尺寸”的教师状态通过学生的预测头 (\(\mathbf{W}^s\)) 传递。这生成了一个概率分布 \(\mathbf{p}^{t \to s}\)。

显示使用学生预测头生成教师概率的方程。

由于投影器是随机初始化的，我们需要对其进行训练，以确保它保留了教师的知识。我们添加了一个损失项，以确保投影后的教师状态仍然可以预测正确的真实 token:

投影教师状态的交叉熵损失方程。

现在，学生 (自然地) 和教师 (通过投影) 都在使用同一个预测头 (\(\mathbf{W}^s\)) 生成分布。它们处于同一个输出空间中。我们现在可以在学生空间中计算 KD 损失:

学生空间中的知识蒸馏损失方程。

2. 教师空间中的蒸馏

DSKD 是对称的。我们也希望推动学生去理解教师的世界。

我们使用另一个投影器 \(\mathcal{P}^{s \to t}\) 将学生的隐藏状态扩展到教师的维度 (\(D\)) 。

显示将学生隐藏状态投影到教师维度的方程。

然后，我们将这些投影后的学生状态通过教师的预测头 (\(\mathbf{W}^t\)) 传递，得到分布 \(\mathbf{q}^{s \to t}\)。

显示使用教师预测头生成学生概率的方程。

由于教师的预测头已经是固定且预训练好的，我们这里不需要辅助损失。我们直接计算教师空间中的 KD 损失 (具体来说是 KL 散度) :

教师空间中的知识蒸馏损失方程。

3. 统一目标

最终的训练目标结合了来自两个空间的蒸馏损失，以及用于训练投影器的辅助损失。

DSKD 总损失方程，包括学生空间 KD、教师空间 KD 和投影器交叉熵的总和。

通过最小化这个组合损失，DSKD 确保学生不仅在最终概率上模仿教师，而且在内部表示结构上也模仿教师，从而最大化两个模型之间的相似性。

解决词表不匹配: 跨模型注意力机制

还剩下一个主要障碍。上述方法假设学生和教师共享相同的词表 (因此也共享相同的分词器) 。如果它们共享词表，它们的序列长度就是相同的。

但在现代 LLM 领域，这种情况很少见。你可能想将知识从 Qwen (不同词表) 蒸馏到 GPT-2，或者从 Mistral 蒸馏到 TinyLlama。如果分词器不同，序列 I like AI 对于学生来说可能是 3 个 token，而对于教师来说可能是 4 个 token。你不能简单地投影隐藏状态，因为序列长度 (\(n\) 和 \(m\)) 不匹配。

为了解决这个问题，作者引入了 跨模型注意力 (Cross-Model Attention, CMA) 机制。

通过注意力对齐 Token

作者没有使用简单的线性投影，而是将对齐问题视为一个注意力任务。

  1. 查询 (Query，学生) : 学生的嵌入和目标 token 充当查询 (\(Q\)) 。
  2. 键与值 (Key & Value，教师) : 教师的嵌入和隐藏状态充当键 (\(K\)) 和值 (\(V\)) 。

定义跨模型注意力的查询、键和值矩阵的方程。 定义源自教师的键和值矩阵的方程。

通过计算学生查询和教师键之间的注意力，模型学习到一个 注意力矩阵 (\(\mathbf{a}^{t \to s}\)) 。这个矩阵代表了对齐关系——它告诉我们教师序列中的哪些 token 对应于学生序列中的特定 token。

从教师到学生的注意力矩阵方程。

然后，我们使用这个矩阵来计算教师值的加权和。结果是一个教师隐藏状态序列，它与学生的序列长度完美地——逐 token 地——对齐了。

显示使用注意力矩阵计算对齐的教师隐藏状态的方程。

这个对齐后的表示 \(\tilde{\mathbf{h}}\) 现在可以直接插入上一节描述的 DSKD 框架中。作者还执行了反向操作 (学生到教师的对齐) 以支持双空间蒸馏。

从学生到教师的注意力矩阵方程。 显示计算对齐的学生隐藏状态的方程。

这种机制使得 DSKD 成为一个通用框架，适用于任何一对 LLM，无论它们的词表如何。

实验结果

作者在指令遵循基准 (如 Dolly、Self-Instruct 和 Vicuna-Evaluation) 上评估了 DSKD。他们测试了两种场景: 具有相同词表的模型和具有不同词表的模型。

场景 1: 相同词表 (GPT-2)

在第一组实验中，他们将 GPT2-1.5B (教师) 蒸馏到 GPT2-120M (学生) 。他们将 DSKD 与标准 KD 方法 (如普通 KL、反向 KL (RKL) 和 Jensen-Shannon (JS) 散度) 进行了比较。

表 1 显示 Rouge-L 分数。DSKD 在所有指标上均优于 GPT-2 的基线。

如 表 1 所示，添加“w/ DSKD”组件在所有距离函数上都一致地提高了性能。例如，将 DSKD 与 JS 散度结合使用，平均 Rouge-L 分数从 16.98 提高到了 18.37——这是一个显著的提升。

场景 2: 相同词表 (LLaMA)

他们扩大了实验规模，使用 LLaMA2-7B 作为教师，TinyLLaMA-1.1B 作为学生。

表 2 显示 LLaMA 实验的 Rouge-L 分数。DSKD 显示出优于基线的显著收益。

表 2 证实了这一趋势。DSKD 取得了实质性的收益。值得注意的是，使用自适应 KL (AKL) 配合 DSKD 带来了巨大的性能提升，将平均分数从 23.30 提高到了 26.65。

场景 3: 不同词表

表 1 和表 2 的底部部分 (标记为“Different Vocabularies”) 展示了跨模型注意力 (CMA) 机制的威力。

  • Qwen \(\to\) GPT-2: DSKD-CMA 优于现有的跨分词器方法，如 MinED 和 ULD。
  • Mistral \(\to\) TinyLLaMA: DSKD-CMA-AKL 取得了 26.99 的平均分，超过了最佳基线 (MinED) 5 分以上。有趣的是，这一结果甚至比 LLaMA2 \(\to\) TinyLLaMA 的蒸馏效果更好，这可能是因为 Mistral 是比 LLaMA2 更强的教师模型。

分析: 为什么它有效？

为了验证“双空间”概念是成功的驱动力，作者进行了消融研究。他们比较了:

  1. 异空间 (Diff. Space): 标准 KD (不同预测头) 。
  2. 学生空间 (Student Space): 仅投影教师 \(\to\) 学生。
  3. DSKD: 学生和教师空间两者兼有。

表 3 显示消融研究结果。仅学生空间的 KD 优于标准 KD，但 DSKD 是最好的。

表 3 显示，简单地将蒸馏移入学生空间就能带来巨大的提升。增加教师空间 (DSKD) 则进一步优化了性能。

此外，作者测量了训练后学生和教师的表示结构之间的实际几何距离。

图 3 显示，与普通 KD 相比，DSKD 显著减小了教师和学生表示结构之间的距离。

图 3 验证了最初的假设: 与普通 KD (绿色框) 相比，DSKD (橙色框) 导致表示结构之间的距离要小得多。学生确实在像教师一样“思考”。

最后，他们请 GPT-4 对生成的回复质量进行评判。

图 2 显示 GPT-4 的胜率。DSKD 在大多数情况下胜过普通 KD。

如 图 2 所示，与基线模型相比，DSKD 训练的模型在大多数情况下更受 GPT-4 青睐。

结论

论文《双空间知识蒸馏》强调了我们在压缩大型语言模型方式中一个微妙但关键的低效之处。通过将教师和学生视为处于独立的向量空间中，标准知识蒸馏无法完全对齐它们的表示。

DSKD 通过显式投影隐藏状态以统一输出空间来解决这个问题。当与跨模型注意力结合时，它提供了一个强大的框架，即使模型使用完全不同的分词器也能有效工作。

随着我们继续依赖大型 LLM，像 DSKD 这样的技术将变得至关重要。它们使我们能够提取庞大模型的高级推理能力，并将其注入到紧凑、可部署的学生模型中，从而使人工智能对每个人都更加易用且高效。

附录: 实验细节

对于那些对实现这些结果的具体超参数感兴趣的人，作者提供了他们的训练配置，并指出基于验证集表现，蒸馏温度 \(\tau\) 被设定为 2.0。

表 4 显示了训练配置，如学习率和批量大小。

总体损失函数平等地平衡了任务学习和蒸馏:

总损失方程，平等加权交叉熵和 KD 损失。