大语言模型预训练中隐藏层蒸馏的研究

# 大型语言模型预训练中隐藏层蒸馏的研究
来源: https://arxiv.org/html/2605.11513
Maxime Guigon Google DeepMind guimax@google\.com &Lucas Dixon Google DeepMind ldixon@google\.com &Michaël E\. Sander Google DeepMind michaelsander@google\.com

###### 摘要

知识蒸馏（KD）是训练大型语言模型（LLM）的关键工具，但绝大多数研究集中在仅依赖输出 logits 的方法上，忽略了教师模型中间表示中的语义信息。虽然隐藏层蒸馏（HLD）在编码器架构中显示出潜力，但其在大规模纯解码器预训练中的应用仍未得到充分探索。通过计算量受控的实验，我们将 HLD 与基于 logits 的 KD 和自监督基线进行了基准测试，使用 Gemma3 3.4B 作为教师模型，使用 123M 和 735M 的学生模型在 C4 数据集上训练，训练数据量高达 168B tokens。我们的实验表明，HLD 在下游评估任务上并未始终优于标准 KD。然而，我们证明，在所有共享超参数配置下，HLD 相对于 KD 都能获得系统性的困惑度（perplexity）增益，这表明可以提取潜在信号，但要使其在 LLM 预训练中发挥更重要的作用，可能需要取得突破。

## 1 引言

大型语言模型（LLM）在 numerous 自然语言处理（NLP）任务中取得了显著的成功，这一现象很大程度上由“缩放定律”（scaling laws）的经验观察所驱动，该定律将模型大小和训练 token 数量与性能提升联系起来（Hoffmannet al., 2022 (https://arxiv.org/html/2605.11513#bib.bib20); Rosenfeldet al., 2020 (https://arxiv.org/html/2605.11513#bib.bib21); Hestnesset al., 2017 (https://arxiv.org/html/2605.11513#bib.bib22)）。然而，这种规模不断扩大的趋势带来了显著的实践挑战，包括飙升的部署成本、增加的推理延迟以及巨大的能源消耗。因此，减少 LLM 的计算足迹，特别是在推理阶段，已成为一个关键的研究领域。

本文聚焦于知识蒸馏（KD）（Hintonet al., 2015 (https://arxiv.org/html/2605.11513#bib.bib1)），这是一种稳健且被广泛采用的范式，用于通过将较大的“教师”模型的知识转移到较小、更高效的“学生”模型中，训练后者以复现教师的行为。在过去十年的研究中，KD 被证明是一种持久且高效的模型压缩技术，在计算机视觉（Ahnet al., 2019 (https://arxiv.org/html/2605.11513#bib.bib23)）、语音处理（Tan and Wang, 2021 (https://arxiv.org/html/2605.11513#bib.bib24)）以及包括纯编码器 LLM（Sanhet al., 2019 (https://arxiv.org/html/2605.11513#bib.bib8)）和涵盖预训练和后训练阶段的各种现代纯解码器 LLM（Gemma Team, 2025 (https://arxiv.org/html/2605.11513#bib.bib18); DeepSeek-AI, 2024 (https://arxiv.org/html/2605.11513#bib.bib19)）在内的多种架构中均显示出效用。尽管如此，当前的蒸馏方法仍难以缩小学生与教师之间的性能差距（Penget al., 2025 (https://arxiv.org/html/2605.11513#bib.bib4); Sanhet al., 2019 (https://arxiv.org/html/2605.11513#bib.bib8); Wanget al., 2020 (https://arxiv.org/html/2605.11513#bib.bib9); Liuet al., 2024b (https://arxiv.org/html/2605.11513#bib.bib3)），而教师模型通常表现出内部冗余（Durraniet al., 2020 (https://arxiv.org/html/2605.11513#bib.bib25); Dalviet al., 2020 (https://arxiv.org/html/2605.11513#bib.bib26)），预示着无损压缩的可能性。此外，在大规模、高成本的 LLM 预训练背景下，学术界对蒸馏技术的探索仍然有限，主要集中于经典的 logits 级 KD（Penget al., 2025 (https://arxiv.org/html/2605.11513#bib.bib4); Busbridgeet al., 2025 (https://arxiv.org/html/2605.11513#bib.bib5)）。

隐藏层蒸馏（HLD）是 KD 的早期扩展，最初由 Romeroet al. (2015 (https://arxiv.org/html/2605.11513#bib.bib2))（FitNets）为 CNN 引入，它利用教师的隐藏状态，已被广泛应用于较小规模的编码器-解码器语言模型，显示出相比 KD 的潜在优势（Jiaoet al., 2020 (https://arxiv.org/html/2605.11513#bib.bib6); Lianget al., 2023b (https://arxiv.org/html/2605.11513#bib.bib7); Sanhet al., 2019 (https://arxiv.org/html/2605.11513#bib.bib8); Wanget al., 2020 (https://arxiv.org/html/2605.11513#bib.bib9), 2021 (https://arxiv.org/html/2605.11513#bib.bib10); Yuet al., 2025 (https://arxiv.org/html/2605.11513#bib.bib11); Zuoet al., 2022 (https://arxiv.org/html/2605.11513#bib.bib12); Dasgupta and Cohn, 2025 (https://arxiv.org/html/2605.11513#bib.bib13); Sunet al., 2019 (https://arxiv.org/html/2605.11513#bib.bib14); Lianget al., 2023a (https://arxiv.org/html/2605.11513#bib.bib15)）。尽管具有这种潜力，HLD 的应用从未扩展到当前最先进 LLM 的规模。

鉴于此，我们提出专门针对预训练阶段的纯解码器模型进行 HLD 探索。评估 LLM 对微小的超参数变化非常敏感，这些变化可能会不成比例地影响最终性能；因此，我们建立了严格的基线，确保任何性能提升归因于所评估的蒸馏方法，而非伪影。我们的实验建立在开源 NanoDo 代码库（Liuet al., 2024a (https://arxiv.org/html/2605.11513#bib.bib39)）、C4 数据集（Raffelet al., 2020 (https://arxiv.org/html/2605.11513#bib.bib27)）以及 Gemma 系列开源权重模型（Gemma Team, 2025 (https://arxiv.org/html/2605.11513#bib.bib18)）之上。

TeacherEmbedLayer 1Layer 2Layerkk⋮\\vdotsLayerLLStudentEmbedLayer 1Layermm⋮\\vdotsLayerNNxxxxDe\-embedDe\-embedLKD\\mathcal\{L\}\_\{\\mathrm\{KD\}\}RegressorLemb\\mathcal\{L\}\_\{\\mathrm\{emb\}\}hkTh\_\{k\}^\{T\}hmSh\_\{m\}^\{S\}

图 1: HLD 概述。学生接收两个训练信号：LKD\\mathcal\{L\}\_\{\\mathrm\{KD\}\} 匹配教师的输出 logits，而 L_emb\\mathcal\{L\}\_\{\\mathrm\{emb\}\} 通过一个学习到的回归器将学生隐藏状态 h_m^S\\mathbf\{h\}\_\{m\}^\{S\} 与教师隐藏状态 h_k^T\\mathbf\{h\}\_\{k\}^\{T\} 对齐。更具体地说，我们做出了以下贡献：

- •我们系统地评估了专门应用于纯解码器 LLM 预训练的 HLD，通过使用 gemma 3.4B 教师模型，在 8B 和 168B tokens 上训练 123M 和 735M 的 gemma 学生模型（Gemma Team, 2025 (https://arxiv.org/html/2605.11513#bib.bib18)）。
- •我们建立了一个严格的评估框架，该框架基于严格的计算等价性（FLOPs）而非 token 数量来基准测试方法，明确考虑了十亿以下参数学生模型中解除嵌入（de-embedding）和损失计算不可忽视的成本。
- •我们通过广泛的计算量匹配实验以及在 Wikitext 103、HellaSwag、WinoGrande、LAMBADA、PIQA 和 ARC-E 上的评估证明，与 KD 相比：
    - –(i) 隐藏层对齐和 logits 蒸馏的联合优化（HLDC）的表现与 KD 持平；
    - –(ii) 顺序优化（HLDF）在所有共享超参数配置下在 C4 困惑度方面带来了适度的提升，而下游评估的性能保持相当。

## 2 背景和相关工作

#### Transformers。

Transformer（Vaswaniet al., 2017 (https://arxiv.org/html/2605.11513#bib.bib37)）是用于建模序列数据的主要神经网络架构。当前的生成式 LLM 使用纯解码器版本。给定输入字符串 s，分词器生成一系列离散的 token 索引 T=[t_1,...,t_n]\\mathbf\{T\}=\[t\_\{1\},\\dots,t\_\{n\}\] \\in \{0,...,V-1\}^n\\in\\\{0,\\dots,V\-1\\\}^\{n\}，其中 V 是词汇表大小。这些索引被映射为连续的嵌入 H^0=[h_1^0,...,h_n^0]\\mathbf\{H\}^\{0\}=\[\\mathbf\{h\}\_\{1\}^\{0\},\\dots,\\mathbf\{h\}\_\{n\}^\{0\}\] \\in \\mathbb\{R\}^\{d\_\{emb\}\\times n\}。骨干网络由 D 个残差层组成，更新规则为 H^{k+1}=H^k+f_{layer}^{k+1}(H^k)\\mathbf\{H\}^\{k\+1\}=\\mathbf\{H\}^\{k\}\+f\_\{layer\}^\{k\+1\}\(\\mathbf\{H\}^\{k\}\)。最后，解除嵌入层将最终表示 H^D\\mathbf\{H\}^\{D\} 投影到 logits Z\\in\\mathbb\{R\}^\{V\\times n\}。下一个 token 的概率分布为 p_n=\\text\{softmax\}\(\\mathbf\{z\}\_\{n\}/\\tau\)，其中 \\tau 是温度参数。

#### 自监督训练。

在没有教师模型的情况下，训练 Transformer 的传统方法是在全面的文本语料库上进行自监督学习。对于面向生成任务的纯解码器模型，优化目标通常采用因果语言建模损失。如果 \\bm\{\\theta\} 代表模型参数，则单个输入序列的损失定义为：

L_{\\text\{data\}\}\(\\bm\{\\theta\}\)=\-\\frac\{1\}\{n\-1\}\\sum\_\{i=1\}^\{n\-1\}\\log\\mathbf\{p\}\_\{i\}\[t\_\{i\+1\}\]。 (1)

我们将此目标称为 NLL（负对数似然）。

#### 知识蒸馏。

在存在教师模型的背景下，知识蒸馏（KD）（Hintonet al., 2015 (https://arxiv.org/html/2605.11513#bib.bib1)）框架可应用于任何分类模型。该优化目标旨在将学生模型的输出分布与教师模型对齐。为了量化这些分布之间的差异（此处记为 p 和 q）并计算相关损失，通常在 softmax 函数中使用温度项结合 Kullback-Leibler（KL）散度：

\\operatorname\{KL\}\(p\\,\\\|\\,q\)=\\sum\_\{x\\in\\mathcal\{X\}\}p\(x\)\\log\\frac\{p\(x\)\}\{q\(x\)\}。

这种方法因其相比 NLL 的优越性能而受到认可，因为教师的 logits 包含比 one-hot 真实标签更全面的信息。将学生的输出概率记为 p_i^S\\mathbf\{p\}\_\{i\}^\{S\}，教师记为 p_i^T\\mathbf\{p\}\_\{i\}^\{T\}，相关目标定义如下：

L_{\\text\{logits\}\}\(\\bm\{\\theta\}\_\{S\}\)=\\frac\{\\tau^\{2\}\}\{n\}\\sum\_\{i=1\}^\{n\}\\text\{KL\}\(\\mathbf\{p\}\_\{i\}^\{T\}\|\|\\mathbf\{p\}\_\{i\}^\{S\}\)。

在实践中，使用的损失是 NLL 损失和知识蒸馏损失的组合：

L_{\\text\{KD\}\}\(\\bm\{\\theta\}\_\{S\}\)=\(1\-\\alpha\)\\mathcal\{L\}\_\{\\text\{data\}\}\(\\bm\{\\theta\}\_\{S\}\)\+\\alpha\\mathcal\{L\}\_\{\\text\{logits\}\}\(\\bm\{\\theta\}\_\{S\}\)，(2)

#### 开源研究中的蒸馏。

关于纯解码器模型蒸馏技术的开源文献在大规模模型方面相对稀缺。Peng et al. (2025 (https://arxiv.org/html/2605.11513#bib.bib4)) 做出了类似的观察，这促使他们使用透明、开源权重的 9B 和 32B 参数教师模型，以及从 330M 到 6.8B 参数规模的学生模型，对设计空间进行了彻底的探索。他们观察到，与 NLL 相比，知识蒸馏始终带来改进，并探索超参数空间以识别最佳配置。同样，Busbridge et al. (2025 (https://arxiv.org/html/2605.11513#bib.bib5)) 提出了一种蒸馏缩放定律，以根据计算预算以及学生与教师模型之间的分配来预测蒸馏模型的性能。他们将教师和学生规模从 143M 扩大到 12.6B 参数，在多达 Chinchilla 最优 token 预算（Hoffmannet al., 2022 (https://arxiv.org/html/2605.11513#bib.bib20)）16 倍的训练数据上进行训练。

#### 专有开发中的蒸馏。

与这些学术努力并行，主要行业实验室已广泛采用蒸馏来训练最先进的模型（Yanget al., 2025 (https://arxiv.org/html/2605.11513#bib.bib17); Gemini Team, 2025 (https://arxiv.org/html/2605.11513#bib.bib16); Gemma Team, 2025 (https://arxiv.org/html/2605.11513#bib.bib18); DeepSeek-AI, 2024 (https://arxiv.org/html/2605.11513#bib.bib19)）。虽然具体的训练管道和超参数通常属于专有的，但这些发布的高性能——其中一些是开源权重模型——为蒸馏在预训练阶段的有效性提供了强有力的实证证据。

#### 隐藏层蒸馏。

由 Romero et al. (2015 (https://arxiv.org/html/2605.11513#bib.bib2))（FitNets）引入的 HLD 采用两阶段训练策略。第一阶段使用学习到的回归器将学生的中间隐藏激活与教师的对齐。第二阶段然后在整个架构上应用标准知识蒸馏。

#### 针对 Transformers 的 HLD。

从 Transformer 模型的潜在表示中提取价值的当代研究对初始框架引入了各种修改，产生了多样化的损失函数。这些实验变体的例子包括注意力图（Jiaoet al., 2020 (https://arxiv.org/html/2605.11513#bib.bib6)）和值关系（Wanget al., 2020 (https://arxiv.org/html/2605.11513#bib.bib9), 2021 (https://arxiv.org/html/2605.11513#bib.bib10)）的对齐，使用不同层分配方案的多层匹配策略（Yuet al., 2025 (https://arxiv.org/html/2605.11513#bib.bib11)），预先训练任务感知过滤器以匹配学生和教师隐藏维度（Lianget al., 2023b (https://arxiv.org/html/2605.11513#bib.bib7)），应用非学习的变换不变技术以调和不同的隐藏激活维度（Dasgupta and Cohn, 2025 (https://arxiv.org/html/2605.11513#bib.bib13)），直接 RMS 匹配（Sunet al., 2019 (https://arxiv.org/html/2605.11513#bib.bib14); Zuoet al., 2022 (https://arxiv.org/html/2605.11513#bib.bib12)）或相同宽度教师和学生对归一化激活的余弦距离（Sanhet al., 2019 (https://arxiv.org/html/2605.11513#bib.bib8)）。

值得注意的是，当前的实现通常倾向于简单的线性回归器，而不是 Romero et al. (2015 (https://arxiv.org/html/2605.11513#bib.bib2)) 最初提出的多层模型。此外，早期文献中描述的两阶段训练协议通常被单阶段过程取代，该过程最小化一个聚合目标函数：

L_{HLD}=\\beta\\mathcal\{L\}\_\{data\}\+\\alpha\\mathcal\{L\}\_\{logits\}\+\\gamma\\mathcal\{L\}\_\{emb\}。 (3)

不幸的是，大多数现有研究集中在后训练阶段以及编码器或编码器-解码器模型上。据我们所知，Liang et al. (2023b (https://arxiv.org/html/2605.11513#bib.bib7)) 是唯一在纯解码器模型上实验 HLD 的研究，且仅限于持续预训练。他们使用 Radford et al. (2019 (https://arxiv.org/html/2605.11513#bib.bib38)) 的 GPT-2_12（120M 参数）作为教师，GPT-2_6（82M 参数）作为学生，后者初始化为教师层的一个子集。据我们所知，我们的工作是第一项针对因果 LLM 预训练期间 HLD 的开源研究。

## 3 提出方法

为了研究原始的 HLD 公式以及 Transformer 文献中的当代方法，我们评估以下两种方法。

#### 顺序优化（HLDF）。

我们适配了 Romero et al. (2015 (https://arxiv.org/html/2605.11513#bib.bib2)) 的