📚 课程目录第4章 Logit蒸馏进阶

4. Logit蒸馏进阶：Logit标准化技巧、NAT方法、动态温度调节策略

Logit蒸馏，说白了就是让Student模型去模仿Teacher模型的“软输出”。但很多同学做的时候发现，直接拿Teacher的logits来蒸馏，效果往往不稳定。我刚开始做BERT蒸馏时就踩过这个坑——Teacher的logits分布太“尖”了，Student根本学不动。

这一章，我们专门聊聊Logit蒸馏的三个进阶技巧：Logit标准化、NAT（Neuron Attention Transfer）、动态温度调节。这三个技巧，是我在多个蒸馏项目里反复验证过的，能显著提升Student模型的收敛速度和最终精度。

4.1 Logit标准化技巧

先说说最基础的问题。Teacher模型的logits，不同类别的数值范围可能差很大。比如分类任务里，某个类别的logit可能是10.0，另一个只有0.5。直接拿这些logits去算KL散度，Student模型会“盯着”那些大数值的类别学，忽略小数值的类别。

我个人习惯的做法是：对Teacher的logits做标准化，让它们的分布更“平缓”。

核心思路：将Teacher的logits减去均值，再除以标准差，让分布接近标准正态分布。然后再乘以一个缩放因子，控制“软硬”程度。

具体实现很简单：

def standardize_logits(logits, eps=1e-8):
    # logits: [batch_size, num_classes]
    mean = logits.mean(dim=-1, keepdim=True)
    std = logits.std(dim=-1, keepdim=True)
    standardized = (logits - mean) / (std + eps)
    return standardized

标准化之后，Teacher的logits分布更均匀，Student学起来不会“偏科”。

我的经验：标准化后的logits，配合温度T=4.0左右，效果最好。温度太低（T=1.0）时，标准化反而会丢失Teacher的“暗知识”。

你可能会问：标准化会不会把Teacher的“知识”给洗没了？嗯，这里要注意——标准化只是调整了数值范围，没有改变类别之间的相对排序。Teacher认为“猫”比“狗”更像“老虎”，这个关系在标准化后依然保留。

4.2 NAT（Neuron Attention Transfer）方法

Logit蒸馏只关注输出层，但Teacher模型的“知识”其实藏在中间层。NAT方法的核心是：让Student的中间层神经元激活模式，去模仿Teacher的中间层神经元激活模式。

我在做文本分类蒸馏时试过NAT，效果比纯Logit蒸馏高了1.2个点。为什么？因为中间层的注意力模式，包含了Teacher对输入文本的“理解方式”。

NAT的流程大致如下：

选取对齐层：从Teacher和Student中选取对应的Transformer层（比如第6层对第3层）。
计算注意力图：对每一层，计算所有head的注意力权重矩阵的平均值，得到一个 [seq_len, seq_len] 的注意力图。
计算损失：让Student的注意力图去逼近Teacher的注意力图，通常用MSE损失。

代码实现：

def nat_loss(student_attn, teacher_attn):
    # student_attn, teacher_attn: [batch_size, num_heads, seq_len, seq_len]
    # 先对head维度取平均
    student_avg = student_attn.mean(dim=1)  # [batch_size, seq_len, seq_len]
    teacher_avg = teacher_attn.mean(dim=1)  # [batch_size, seq_len, seq_len]
    # 计算MSE损失
    loss = F.mse_loss(student_avg, teacher_avg)
    return loss

避坑指南：我曾经在NAT里直接用了所有head的注意力图，结果Student模型学崩了。后来发现，不同head关注的信息差异很大，直接平均会丢失信息。建议先对head做聚类，只选那些“关键head”来对齐。

NAT的另一个变体是Neuron Selectivity Transfer，它不关注注意力图，而是关注每个神经元的激活值。说白了，就是让Student的中间层神经元，在相同输入下，激活模式尽量和Teacher一致。

方法	对齐对象	损失函数	适用场景
NAT（原始）	注意力图	MSE	Transformer模型
Neuron Selectivity	神经元激活值	MSE / KL散度	CNN / RNN / Transformer
FitNet	中间层特征	MSE	通用

你想想看，NAT相当于让Student“偷师”Teacher的思考过程，而不仅仅是看Teacher的答案。这就像学数学，光看答案没用，得学会解题思路。

4.3 动态温度调节策略

温度T是Logit蒸馏里的关键超参数。T越大，softmax输出的分布越“软”，Student能学到更多“暗知识”。但T太大，分布太平坦，Student又学不到区分度。

我见过很多同学固定T=4.0从头训到尾，效果其实一般。更好的做法是：动态调节温度。

动态温度调节的核心思想是：训练初期用高温，让Student广泛探索；训练后期用低温，让Student精细拟合。

我常用的几种策略：

线性衰减：T从8.0线性降到2.0，每1000步降0.1。
指数衰减：T = T0 * exp(-k * step)，衰减速度可调。
基于验证集的自适应调节：每500步在验证集上评估一次，如果Student的精度提升变慢，就降低温度。

代码实现（线性衰减）：

def get_dynamic_temperature(step, total_steps, T_start=8.0, T_end=2.0):
    # 线性衰减
    T = T_start - (T_start - T_end) * (step / total_steps)
    return max(T, T_end)

我的习惯：在训练的前20%步数用高温（T=8.0），中间60%步数线性衰减到T=3.0，最后20%步数固定T=3.0。这样Student既能学到全局知识，又能精细调整。

动态温度调节还有一个好处：缓解蒸馏过程中的“过平滑”问题。高温时，Student的logits分布比较均匀，不容易过拟合到Teacher的噪声上；低温时，Student开始关注细节，提升精度。

4.4 本章知识结构图

下面这张图，把Logit蒸馏进阶的三个技巧串起来了。你可以看到，Logit标准化是“预处理”，NAT是“中间层对齐”，动态温度是“训练策略”。三者配合使用，效果最好。

4.5 实战建议与避坑

最后，我总结几个实战中容易踩的坑：

标准化后温度失效：标准化后的logits，温度T的效果会变弱。我建议标准化后，T从4.0开始调，不要用8.0以上。
NAT层数选择：不是所有层都适合对齐。我试过对齐所有12层，效果反而不如只对齐第6层和第10层。建议从中间层和最后几层里选。
动态温度与学习率联动：温度降低时，学习率也要跟着降。否则Student在低温阶段容易震荡。

我曾经踩过的一个大坑：在NAT里同时对齐了注意力图和神经元激活值，结果两个损失互相打架，Student模型直接不收敛了。后来我把两个损失加权求和，权重设为0.5:0.5，才稳定下来。

嗯，这一章的内容就到这里。Logit标准化、NAT、动态温度调节，这三个技巧单独用都能提升效果，但组合起来才是真正的“王炸”。你可以在自己的项目里试试，看看效果如何。