4. Logits蒸馏实战

Logits蒸馏，说白了就是让“好学生”模仿“好老师”的答题思路。老师模型输出的是一个概率分布，而不仅仅是“猫”或“狗”的硬标签。学生模型要学的，就是这种“猫有点像狗，但更像猫”的软知识。

我个人习惯把这种蒸馏叫做“暗知识迁移”。你想想看，一个训练好的BERT-large，它知道“苹果”和“梨”都是水果，但更知道“苹果”和“橘子”在颜色上更接近。这种细腻的语义关系，硬标签根本表达不出来。

4.1 教师模型与学生模型的选择策略

选模型这件事，我踩过不少坑。先说说我的经验总结：

选择维度	教师模型	学生模型
参数量级	大（300M~1B+）	小（10M~100M）
架构差异	同族或异族均可	推荐同族（如BERT→TinyBERT）
任务匹配	预训练+微调完成	从零或预训练开始
推理速度	不要求实时	目标延迟内

具体到选择策略，我一般分三步走：

1. 先定学生——根据部署环境的算力、内存、延迟要求，确定学生模型的上限。比如移动端只能用TinyBERT，那就锁定它。
2. 再找老师——在学生能承受的蒸馏成本内，选最大的预训练模型。BERT-large蒸馏到TinyBERT，是我在项目中验证过的最稳组合。
3. 检查对齐——教师和学生的输出层维度必须一致（vocab_size或num_labels）。不一致时，需要加一个线性映射层对齐。

4.2 Logits蒸馏的数学原理

Logits蒸馏的核心损失函数长这样：

L_distill = KL_div(softmax(teacher_logits / T), softmax(student_logits / T)) * T²

这里的 T 是温度参数。T越大，概率分布越平滑，学生能学到更多“次优类”的信息。T=1时就是标准的softmax，T=5时分布几乎变成均匀分布。

嗯，这里要注意：T² 这个缩放因子不能丢。因为梯度会随着T增大而缩小，乘以T²才能让梯度量级保持一致。我刚开始做蒸馏时忘了加这个，结果温度调了半天都没效果。

最终的损失函数通常是蒸馏损失和硬标签损失的加权和：

L_total = α * L_distill + (1 - α) * L_CE

α一般取0.5~0.9。我个人习惯先设0.7，然后根据验证集表现微调。

4.3 基于HuggingFace Trainer的实现

HuggingFace的Trainer类非常灵活，我们只需要自定义一个 DistillationTrainer，重写 compute_loss 方法即可。

先看完整代码：

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

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):
        # 教师模型不更新梯度
        with torch.no_grad():
            teacher_outputs = self.teacher_model(**inputs)
            teacher_logits = teacher_outputs.logits

        # 学生模型前向
        student_outputs = model(**inputs)
        student_logits = student_outputs.logits

        # 蒸馏损失
        loss_distill = 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_ce = F.cross_entropy(student_logits, inputs['labels'])

        # 加权融合
        loss = self.alpha * loss_distill + (1 - self.alpha) * loss_ce

        return (loss, student_outputs) if return_outputs else loss

这段代码的核心逻辑就三步：

1. 用 torch.no_grad() 冻结教师模型，只做推理
2. 分别计算蒸馏损失和交叉熵损失
3. 按α权重融合

4.4 训练脚本与参数配置

训练参数的设置直接影响蒸馏效果。这是我常用的配置：

training_args = TrainingArguments(
    output_dir='./distill_output',
    per_device_train_batch_size=32,
    per_device_eval_batch_size=64,
    learning_rate=5e-5,
    num_train_epochs=6,
    logging_steps=100,
    evaluation_strategy='steps',
    eval_steps=500,
    save_steps=1000,
    fp16=True,  # 混合精度加速
    dataloader_num_workers=4,
)

trainer = DistillationTrainer(
    teacher_model=teacher,
    model=student,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    tokenizer=tokenizer,
    temperature=4.0,
    alpha=0.7,
)

trainer.train()

几个关键参数的经验值：

• temperature：4.0~8.0。任务越复杂，温度可以越高。文本分类我用4.0，NER我用6.0。
• alpha：0.5~0.9。如果学生模型从零开始训练，α设大一点（0.8~0.9）；如果学生已经预训练过，α可以小一点（0.5~0.7）。
• learning_rate：比正常微调小一点。我一般用5e-5，比标准微调的2e-5略大，因为蒸馏损失更平滑。
• epochs：4~8轮。蒸馏收敛比正常训练快，但也不要太少。我见过3轮就过拟合的案例。

4.5 知识蒸馏流程总览

下面这张图是我每次做蒸馏都会贴在工位上的流程：

整个流程其实就一句话：教师模型生成软标签，学生模型同时拟合软标签和硬标签，反向传播只更新学生参数。

4.6 实战效果与指标对比

我在GLUE的MRPC任务上做过一组对比实验，结果如下：

模型	参数量	准确率	推理延迟（ms）	模型大小
BERT-base（教师）	110M	88.9%	12.3	440MB
TinyBERT（无蒸馏）	14.5M	82.1%	2.1	58MB
TinyBERT（Logits蒸馏）	14.5M	86.7%	2.1	58MB

看到没？蒸馏后的TinyBERT只比教师低了2.2个点，但推理速度快了6倍，模型大小只有1/7。这就是Logits蒸馏的价值——用极小的代价换取接近大模型的效果。