3. 中间层蒸馏：隐藏状态对齐、注意力迁移与表征匹配实战

前面我们聊了 Logit 层面的蒸馏，说白了就是让 Student 模仿 Teacher 的最终输出。但只学结果，不学过程，总觉得差点意思。中间层蒸馏，就是要把 Teacher 的“思考过程”也教给 Student。

我个人习惯把中间层蒸馏分成三个层次：隐藏状态对齐、注意力矩阵迁移、表征匹配。它们分别对应 Transformer 里的不同组件，难度和收益也各不相同。

3.1 隐藏状态对齐：用 MSE 拉近特征空间

最直接的做法，就是让 Student 的中间层输出，在数值上尽量接近 Teacher 的对应层。损失函数通常用 MSE（均方误差）。

为什么需要投影层？因为 Student 的隐藏层维度通常比 Teacher 小。比如 Teacher 是 768 维，Student 是 384 维，直接算 MSE 维度不匹配。加一个线性层，相当于做了一次“翻译”。

具体实现时，你需要决定对齐哪些层。常见策略有两种：

策略	做法	适用场景
均匀对齐	每 k 层对齐一次，比如 Teacher 每 2 层对齐 Student 每 1 层	层数差异大时（如 Teacher 12 层，Student 6 层）
关键层对齐	只对齐最后几层或特定层（如第 6、9、12 层）	计算资源有限，或只想保留高层语义

代码实现其实不复杂：

class HiddenStateLoss(nn.Module):
    def __init__(self, teacher_dim, student_dim):
        super().__init__()
        self.proj = nn.Linear(student_dim, teacher_dim)
        self.mse = nn.MSELoss()
    
    def forward(self, h_teacher, h_student):
        # h_teacher: [batch, seq_len, teacher_dim]
        # h_student: [batch, seq_len, student_dim]
        h_student_proj = self.proj(h_student)
        return self.mse(h_teacher, h_student_proj)

3.2 注意力矩阵迁移：让 Student 学会“看哪里”

Transformer 里，注意力矩阵决定了模型“关注”哪些位置。如果 Student 能学到 Teacher 的注意力模式，那效果往往比单纯对齐隐藏状态更好。

注意力迁移的核心思想：让 Student 的注意力矩阵尽量接近 Teacher 的注意力矩阵。损失函数通常用 KL 散度或 MSE。

我建议用 KL 散度，因为它对概率分布更敏感：

def attention_transfer_loss(A_teacher, A_student):
    # A_teacher, A_student: [batch, num_heads, seq_len, seq_len]
    # 对每个 head 和每个位置计算 KL 散度
    loss = F.kl_div(
        F.log_softmax(A_student, dim=-1),
        F.softmax(A_teacher, dim=-1),
        reduction='batchmean'
    )
    return loss

实际操作中，有几个细节要注意：

• 多头对齐：Teacher 有 12 个头，Student 只有 8 个头，怎么对齐？可以取 Teacher 的前 8 个头，或者用线性映射把 8 个头映射到 12 个头。
• 层选择：通常对齐中间几层的注意力矩阵效果最好。太浅的层注意力比较分散，太深的层注意力已经高度抽象。
• 温度参数：给 softmax 加一个温度 T，让注意力分布更平滑。T 越大，分布越均匀，梯度越稳定。

3.3 Representation Matching：更高阶的表征对齐

隐藏状态对齐和注意力迁移，都是“点对点”的匹配。Representation Matching 更进一步，它关注的是表征空间的结构关系。

说白了，就是让 Student 的表征空间和 Teacher 的表征空间在“几何结构”上相似。常用的方法有：

方法	核心思想	损失函数
CCA（典型相关分析）	最大化 Teacher 和 Student 表征的线性相关性	1 - CCA 相关系数
CKA（中心核对齐）	比较两个表征空间的相似性矩阵	1 - CKA 相似度
对比学习	让同一句话的表征靠近，不同句子的表征远离	InfoNCE 损失

我个人比较推荐 CKA，因为它对线性变换不敏感，而且计算相对简单。CKA 的核心是计算两个表征的 Gram 矩阵，然后比较它们的相似度：

def cka_loss(H_teacher, H_student):
    # H: [batch, seq_len, dim]
    # 计算 Gram 矩阵
    G_teacher = H_teacher @ H_teacher.transpose(-2, -1)
    G_student = H_student @ H_student.transpose(-2, -1)
    
    # 中心化
    G_teacher = G_teacher - G_teacher.mean(dim=-1, keepdim=True)
    G_student = G_student - G_student.mean(dim=-1, keepdim=True)
    
    # 计算 CKA
    numerator = (G_teacher * G_student).sum()
    denominator = torch.sqrt((G_teacher**2).sum() * (G_student**2).sum())
    cka = numerator / denominator
    
    return 1 - cka

3.4 三种方法的对比与选择

这三种方法不是互斥的，实际项目中往往组合使用。我整理了一个对比表，方便你根据场景选择：

方法	计算开销	效果提升	实现难度	推荐场景
隐藏状态对齐（MSE）	低	中	低	快速 baseline，资源受限
注意力迁移	中	高	中	分类、匹配任务
Representation Matching	高	最高	高	生成任务、语义相似度

我的建议是：先用 MSE 对齐隐藏状态作为基础，然后加上注意力迁移，最后根据任务需求决定是否上 Representation Matching。不要一上来就全上，否则调参会让你怀疑人生。

3.5 整体流程与 SVG 框架图

下面这张图展示了中间层蒸馏的完整流程。你可以看到，Teacher 和 Student 的每一层都可以对齐，但实际中我们只选关键层。

从图中可以看到，三种对齐方式可以同时进行，也可以只选一种。我个人习惯是：先对齐中间层（注意力迁移），再对齐高层（MSE 或 CKA）。这样 Student 既能学到局部的注意力模式，又能学到全局的语义表征。