模型蒸馏

本页通过示例介绍如何让 SentenceTransformer 模型更快、成本更低、体积更小。这些轻量模型在下游任务中能够达到原始模型性能的 97.5%～100%。

知识蒸馏

知识蒸馏是把教师模型的知识转移给学生模型的过程。它可以用于把句子嵌入扩展到新的语言，参见论文通过知识蒸馏使单语言句子嵌入支持多语言。传统用法则是用速度较慢但效果好的教师模型，训练运行速度更快的学生模型。

学生模型模仿教师模型，从而在保持较高速度的同时取得较好的效果。

我们实现了两种构建学生模型的方案：

1. model_distillation.py：使用 TinyBERT 或 BERT-Small 等轻量 Transformer 模型模仿较大的教师模型。训练使用 EmbedDistillLoss。它可以包含一个可学习的投影，因此教师和学生不必具有相同的嵌入维度。训练后，学生仍保留自身原有的较小输出维度。
2. model_distillation_layer_reduction.py：从教师模型中仅保留某些层，例如只保留4层，并使用 MSELoss 训练。由于减少层数后的学生继承了教师的隐藏层宽度，因此不需要投影，输出维度保持一致，可作为同维度的直接替代。

第二种方案通常效果更好，因为它保留了教师模型的大部分权重。第一种方案则需要从头调整学生模型的全部权重。

速度与性能之间的取舍

更小的模型速度更快，但在下游任务中的效果通常会稍有下降。为了直观展示这种取舍，下面给出 stsb-roberta-base 模型采用不同层数时的原文结果：

降维

预训练模型默认输出的嵌入维度为768（base 模型）或1024（large 模型）。但如果要存储数百万条嵌入，所需内存或存储空间就可能相当大。

dimensionality_reduction.py 展示了一个简单例子：使用主成分分析（PCA），把嵌入降低到所需维度。原文示例将768维降至128维，使存储需求缩减为原来的六分之一，STS benchmark 数据集上的得分仅由85.44小幅降至84.96。

这种降维方法很容易应用于已有模型。甚至可以降至32维，将存储需求缩减为原来的二十四分之一，此时得分降至81.82。

量化

量化模型使用整数而非浮点数执行部分或全部运算。这样既能得到更紧凑的模型，也能利用许多硬件平台上的高性能向量化运算。

对于在 CPU 上运行的模型，原文给出的结果是模型体积可缩小40%，推理速度也会提升；依 CPU 而异，提速幅度为15%～400%。原文同时注明，在该段写作时，PyTorch 尚不支持 GPU 上的模型量化。

示例参见 model_quantization.py。