📚 课程目录第3章 Logits蒸馏实战

3. Logits蒸馏实战：基于HuggingFace Trainer实现BERT-small蒸馏BERT-base

好，咱们进入实战环节。前面讲了那么多蒸馏的理论，今天直接上手干——用HuggingFace的Trainer，把BERT-base的知识蒸馏到BERT-small里。说白了，就是让一个小模型学会大模型的“思考方式”。

我个人习惯，做实验前先把整个流程在脑子里过一遍。蒸馏不是玄学，它是有章可循的。咱们今天要做的，就是调两个关键参数：蒸馏温度T和损失权重α。这两个值怎么配，直接决定你蒸馏出来的模型是“学霸”还是“学渣”。

  核心目标：让BERT-small的Logits分布，尽可能逼近BERT-base的Logits分布。温度T控制分布的“软硬”程度，α控制“学老师”和“学课本”的比例。

3.1 蒸馏损失函数：不只是KL散度

先看损失函数。蒸馏的损失分两块：

硬损失（Hard Loss）：学生模型输出与真实标签的交叉熵。这是常规的监督学习损失。
软损失（Soft Loss）：学生模型输出与教师模型输出的KL散度。这是蒸馏的核心。

总损失就是两者的加权和：

L = α * L_hard + (1 - α) * L_soft

这里α是个0到1之间的数。α=1就是纯监督学习，α=0就是纯模仿老师。我刚开始做蒸馏时，总以为α越小越好，结果模型连基本分类都做不好。后来才明白，α=0.7左右是个不错的起点。

软损失的计算，需要先对Logits做温度缩放：

soft_logits = logits / T

温度T越高，Softmax后的分布越平滑，信息越丰富。T=1就是标准的Softmax，T=0.5会让分布更尖锐，T=5会让分布更平滑。

我的经验：温度T一般设在2-8之间。T太低，学生学不到老师对“错误类别”的排序信息；T太高，所有类别概率都差不多，学生也学不到东西。

3.2 基于Trainer的蒸馏实现

HuggingFace的Trainer非常强大，但默认不支持蒸馏。我们需要自定义一个Trainer，重写compute_loss方法。

来看代码：

import torch
import torch.nn.functional as F
from transformers import Trainer

class DistillationTrainer(Trainer):
    def __init__(self, teacher_model=None, temperature=4.0, alpha=0.7, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.teacher_model = teacher_model
        self.temperature = temperature
        self.alpha = alpha

    def compute_loss(self, model, inputs, return_outputs=False):
        # 学生模型前向传播
        student_outputs = model(**inputs)
        student_logits = student_outputs.logits

        # 教师模型前向传播（不计算梯度）
        with torch.no_grad():
            teacher_outputs = self.teacher_model(**inputs)
            teacher_logits = teacher_outputs.logits

        # 硬损失：学生与真实标签
        labels = inputs.get("labels")
        loss_hard = F.cross_entropy(student_logits, labels)

        # 软损失：学生与教师的KL散度（温度缩放后）
        loss_soft = F.kl_div(
            F.log_softmax(student_logits / self.temperature, dim=-1),
            F.softmax(teacher_logits / self.temperature, dim=-1),
            reduction="batchmean"
        ) * (self.temperature ** 2)

        # 总损失
        loss = self.alpha * loss_hard + (1 - self.alpha) * loss_soft

        return (loss, student_outputs) if return_outputs else loss

注意这行：loss_soft * (self.temperature ** 2)。为什么要乘T²？因为温度缩放后，梯度会缩小T倍，乘回来可以保持梯度量级一致。这个细节，很多教程都不提，但实际训练中影响很大。

我曾经踩过的坑：忘记乘T²，结果软损失的梯度太小，模型根本不学老师。调了半天参数，最后发现是这里的问题。嗯，从此以后我写蒸馏代码，第一件事就是检查这个乘数。

3.3 温度与α的调参实验

调参是蒸馏最磨人的部分。我建议你按下面的表格来设计实验：

实验编号	温度T	α值	预期效果
1	1.0	1.0	纯监督学习，无蒸馏效果
2	4.0	0.7	推荐起点，平衡软硬损失
3	8.0	0.5	高温度，强调类别间关系
4	2.0	0.3	低温度，更依赖老师
5	4.0	0.9	高α，更依赖真实标签

启动训练很简单：

from transformers import AutoModelForSequenceClassification, TrainingArguments

# 加载教师和学生模型
teacher = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=3)
student = AutoModelForSequenceClassification.from_pretrained("bert-small-uncased", num_labels=3)

# 训练参数
training_args = TrainingArguments(
    output_dir="./distill_results",
    num_train_epochs=3,
    per_device_train_batch_size=16,
    learning_rate=2e-5,
    logging_steps=100,
    save_steps=500,
    evaluation_strategy="steps",
)

# 蒸馏Trainer
trainer = DistillationTrainer(
    teacher_model=teacher,
    temperature=4.0,
    alpha=0.7,
    model=student,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
)

trainer.train()

你想想看，跑完这5组实验，你就能画出温度和α对模型性能的影响曲线。我一般会关注两个指标：验证集准确率和学生-教师Logits的KL散度。准确率高说明学得好，KL散度低说明学得像。

3.4 蒸馏流程全景图

下面这张图，把整个蒸馏流程串起来了。我建议你保存下来，每次做实验前看一眼，思路会清晰很多。

3.5 实验结果分析与避坑

跑完实验，你会看到类似这样的结果：

T=1, α=1：准确率最低，因为没有蒸馏效果。这就是个基线。
T=4, α=0.7：准确率最高，软硬损失平衡得好。我大部分项目都用这个组合起步。
T=8, α=0.5：准确率略低，但模型泛化能力可能更好。适合小数据集。
T=2, α=0.3：如果教师模型很强，这个组合能让学生学得最像老师。

避坑指南：我曾经在一个文本分类任务上，把T设到10，α设到0.1，结果学生模型完全学废了——输出概率几乎均匀分布，啥都分不清。后来发现，温度太高，软标签的信息被过度平滑，学生根本抓不住重点。所以，T不要超过8，α不要低于0.3，这是我的经验红线。

另外，别忘了监控训练过程中的KL散度。如果KL散度下降很快，说明学生在快速模仿老师；如果KL散度不降反升，说明训练出了问题——可能是学习率太大，或者温度设置不合理。

嗯，到这里，Logits蒸馏的实战部分就讲完了。你拿着上面的代码和实验设计，直接跑起来就行。调参是个体力活，但也是最能积累经验的地方。多跑几组，你就能找到感觉了。