知识蒸馏|使用一个大模型训练另一个大模型的三种方法

235 阅读3分钟

大语言模型(LLM)不仅能从海量文本中学习,它们也能「互相学习」:

·Llama 4 的 Scout 和 Maverick 模型就是在 Llama 4 Behemoth 的辅助下训练出来的。

·Google 的 Gemma 2 和 3 是在内部模型 Gemini 的指导下训练完成的。

这种「互相学习」的过程,主要依赖于知识蒸馏(Distillation)技术。下面这张图展示了目前主流的三种知识蒸馏方式。

image.png

通俗地说,知识蒸馏的目标就是把一个模型中的“知识”迁移给另一个模型。这在传统深度学习中早就很常见。

在 LLM 的训练中,知识蒸馏可以发生在两个阶段:

1.预训练阶段

·将较大的教师模型(Teacher LLM)和较小的学生模型(Student LLM)一同训练。

·Llama 4 就采用了这种方法来训练 Scout 和 Maverick。

2.后训练阶段(微调后再蒸馏)

·先单独训练好教师模型,然后将其知识迁移给学生模型。

·DeepSeek 就采用了这种方式,把 DeepSeek-R1 蒸馏到了 Qwen 和 Llama 3.1。

你也可以在两个阶段都使用蒸馏,Gemma 3 就是这样做的。

三种主流知识蒸馏方法:

1. Soft-label 蒸馏(软标签蒸馏)

image.png

·使用一个已经训练好的教师模型,对整个语料生成 softmax 概率分布(即每个词的概率而不是最终选哪个词)。

·再让一个未训练的学生模型跑同样的语料,也生成自己的 softmax 概率分布。

·然后训练学生模型,让它的输出概率尽可能接近教师模型的输出。

这种方式可以最大限度保留教师模型的推理逻辑与知识,但也有硬性要求:必须能访问到教师模型的权重,才能获取它输出的概率分布。

即便可以访问,还有个大问题:假设词表大小是 10 万,语料是 5 万亿 token。对于每一个 token,都需要生成一个长度为 10 万的概率向量。即便使用 float8 精度,也需要约 5 亿 GB 内存 来存储这些 soft labels!

所以我们来看第二种方法,它解决了这个问题。

2. Hard-label 蒸馏(硬标签蒸馏)

image.png

·同样用一个已经训练好的教师模型,但这次只输出最终选择的单个 token(即 one-hot 编码)。

·学生模型依旧输出 softmax 概率分布。

·然后训练学生模型,使其概率分布尽量接近教师模型最终选择的 token。

DeepSeek 就是采用这种方式,把 DeepSeek-R1 蒸馏到 Qwen 和 Llama 3.1 的。这种方法相比软标签,资源占用低很多,但信息量也少一些。

3. Co-distillation(协同蒸馏)

image.png

·从一个未训练的教师模型和一个未训练的学生模型同时开始。

·两个模型都跑当前的训练 batch,生成各自的 softmax 概率。

·教师模型正常训练,使用 ground-truth 的硬标签来学习。

·学生模型则训练它的 softmax 输出尽可能接近教师模型当前的输出。

Llama 4 就是用这种方法,从 Llama 4 Behemoth 训练出了 Scout 和 Maverick。

当然,刚开始训练时,教师模型的输出概率也不靠谱,所以学生模型的训练同时使用教师模型的 soft labels + 真实的 ground-truth 硬标签。