Transformer——Q114 推导蒸馏小型Transformer的知识传递损失函数

墨顿

于 2025-05-24 00:54:22 发布

阅读量881

点赞数 7

CC 4.0 BY-SA版权

文章标签： transformer 深度学习人工智能架构变体高效架构

本文链接：https://blog.youkuaiyun.com/pzccool/article/details/148172233

该问题归类到Transformer架构问题集——架构变体——高效架构。请参考LLM数学推导——Transformer架构问题集。

1. 问题背景：大模型瘦身的迫切需求

GPT-3 拥有 1750 亿参数，在文本生成、知识问答等任务中表现惊艳，但庞大的参数量如同沉重的枷锁，不仅训练成本高昂，而且难以部署到手机、车载终端等资源受限设备上。知识蒸馏技术就像一把精巧的手术刀，旨在将大型教师模型蕴含的知识，精准地 “移植” 到小型学生模型中，让后者在保持高性能的同时大幅缩减体积。

然而，Transformer 的知识体系极为复杂，其不仅存在于最终输出的预测结果中，更蕴含在每一层的隐藏状态、注意力矩阵里。如何设计一个精准的损失函数，让学生模型既能模仿教师模型的输出结论，又能复刻其内部推理过程？这需要我们从数学原理出发，深入剖析知识传递的本质。

2. 技术原理：多维度知识传递的数学解构

知识蒸馏的损失函数通常由三部分构成：软标签损失、特征匹配损失、注意力蒸馏损失。每一部分都对应着 Transformer 不同层面的知识，下面我们从数学原理角度，详细解析其设计逻辑。

2.1 软标签损失：为何选择 KL 散度度量分布差异？

教师模型输出的 logits 经过 softmax 后，会形成一个概率分布 $p(\mathbf{y}|\mathbf{x}) = \text{softmax}(T \cdot \text{Logits}_t)$ ，其中 T 是温度参数，T > 1 时会软化概率分布，使得各个类别之间的概率差异不那么尖锐。学生模型的输出分布为 $q(\mathbf{y}|\mathbf{x}) = \text{softmax}(T \cdot \text{Logits}_s)$ 。

我们希望衡量这两个概率分布的差异，从而引导学生模型逼近教师模型的输出分布。KL 散度（Kullback-Leibler Divergence）正是为此而生。从信息论角度来看，KL 散度表示用分布 q 来近似分布 p 时，平均额外需要的信息量。

其数学公式为： $\text{KL}(p || q) = \sum_{y} p(y) \log \frac{p(y)}{q(y)}$

这里的对数运算，本质上是在计算两个概率值的比例关系。当 p(y) 和 q(y) 越接近时， $\frac{p(y)}{q(y)}$ 越趋近于 1， $\log \frac{p(y)}{q(y)}$ 趋近于 0，整个求和结果也就越小。反之，如果两个分布差异很大，某些类别的概率比值会显著偏离 1，导致 KL 散度增大。

在知识蒸馏中，我们除以 $T^2$ 是因为温度缩放会放大 logits 之间的差异，进行归一化处理后，能使梯度更加稳定，避免训练过程中出现剧烈波动。

2.2 特征匹配损失：MSE 为何适合度量隐藏状态差异？

Transformer 每一层的隐藏状态 $\mathbf{H}_t^l$ （教师模型第 l 层输出），记录了输入序列在该层的语义编码，蕴含着丰富的上下文信息。学生模型对应的隐藏状态 $\mathbf{H}_s^l$ 需要尽可能与教师模型对齐。

隐藏状态是连续的向量，我们需要一种度量方法来衡量两个向量之间的距离。均方误差（MSE）是最直观的选择，它通过计算对应元素差值的平方和，来反映两个向量的差异程度： $\mathcal{L}_{\text{feat}} = \sum_{l=1}^{L} \frac{1}{n} \|\mathbf{H}_t^l - \mathbf{H}_s^l\|_2^2$

平方运算的作用是将所有差值都转化为正数，并且放大较大的差值，使得模型更加关注差异明显的部分。求和操作则综合考虑了整个向量的所有元素，取平均（除以 n）是为了消除序列长度的影响，让损失具有可比性。

2.3 注意力蒸馏损失：如何捕捉注意力模式的差异？

教师模型的注意力矩阵 $\mathbf{A}_t^l$ 展示了第 l 层中每个 token 对其他 token 的关注程度，这是 Transformer 理解上下文依赖关系的核心机制。学生模型的注意力矩阵 $\mathbf{A}_s^l$ 需要学习这种聚焦模式。

如果将注意力矩阵视为概率分布（每一行元素之和为 1），那么 KL 散度依然是合适的度量方法。它能够衡量学生模型和教师模型在注意力分配上的差异，促使学生模型学会像教师模型那样关注关键信息。如果更关注绝对值差异，也可以使用 MSE： $\mathcal{L}_{\text{attn}} = \sum_{l=1}^{L} \frac{1}{n^2} \text{KL}(\mathbf{A}_t^l || \mathbf{A}_s^l)$

这里除以 $n^2$ 是因为注意力矩阵是 $n \times n$ 的方阵，进行归一化后可以使损失值处于合理范围，便于训练过程中的优化。

2.4 总损失函数：知识的融合与平衡

将上述三部分损失加权求和，得到最终的知识传递损失函数： $\mathcal{L}_{\text{KD}} = \alpha \mathcal{L}_{\text{soft}} + \beta \mathcal{L}_{\text{feat}} + \gamma \mathcal{L}_{\text{attn}}$

超参数 $\alpha, \beta, \gamma$ 起到平衡不同类型知识的作用。例如，如果教师模型层数远多于学生模型，那么在特征匹配损失中，可以让 $\beta$ 随着层数增加而衰减，避免强行对齐导致的信息失真。

3. LLM 中的实战：小模型的逆袭之路

案例 1：DistilBERT 蒸馏 BERT 教师模型是 12 层的 BERT-Base，学生模型是 6 层的 DistilBERT。在蒸馏过程中：
- 软标签损失让 DistilBERT 模仿 BERT 的掩码语言模型输出分布，学习词汇之间的语义关联；
- 特征匹配损失对齐每一层的隐藏状态（学生第 l 层匹配教师第 2l 层），捕捉不同层次的语义表示；
- 注意力蒸馏损失使得 DistilBERT 学会像 BERT 那样关注关键 token，例如在命名实体识别任务中，准确捕捉人名、地名等实体。最终 DistilBERT 体积减少 40%，推理速度提升 60%，同时保留了 BERT 95% 的性能。
案例 2：GPT-2 蒸馏到对话机器人 教师模型是 768M 参数的 GPT-2，学生模型是轻量化的对话模型。通过知识蒸馏：
- 软标签损失让学生模型生成的回复概率分布更合理，避免出现逻辑混乱的回答；
- 特征匹配损失使学生模型能够记住对话历史中的关键信息，例如用户之前提到的偏好、问题；
- 注意力蒸馏损失帮助学生模型聚焦对话中的重要轮次，在多轮交互中保持上下文一致性。
案例 3：跨模态蒸馏（T5 到文本生成模型） 教师模型是处理图像 + 文本的多模态 T5，学生模型是纯文本的小型 Transformer。通过注意力蒸馏，学生模型能够学习到教师模型在生成图像描述时的注意力分配模式，即使在没有图像输入的情况下，也能生成更具画面感和细节的文本。

4. 优缺点：蒸馏技术的双刃剑

优点：
- 轻量化部署：大幅减少模型参数量和计算量，适合在移动端、嵌入式设备上运行；
- 知识复用：无需大量标注数据，直接继承教师模型在大规模预训练中学习到的通用知识；
- 泛化能力增强：软标签包含了教师模型的不确定性信息，有助于学生模型在面对未知数据时表现更稳健。
缺点：
- 教师依赖问题：如果教师模型存在偏见或错误，学生模型也会继承这些缺陷；
- 层匹配难题：当学生和教师模型层数、结构不同时，特征对齐需要复杂的跨层映射，容易引入额外误差；
- 信息衰减风险：如果损失函数设计不合理，学生模型可能只能学到教师模型的表面特征，无法掌握深层知识。

5. 优化策略：让蒸馏更精准高效

动态权重调整：根据训练阶段动态调整损失权重。初期侧重特征匹配损失，帮助学生模型建立基础语义表示；后期加大软标签损失权重，细化输出结果。例如采用余弦退火策略动态调整权重值。
自适应层选择：通过计算各层隐藏状态的信息熵、激活值方差等指标，自动筛选出信息丰富的关键层进行重点蒸馏，提高知识传递效率。
自监督辅助蒸馏：在损失函数中加入学生模型的自监督任务（如掩码语言模型损失），增强学生模型的自主学习能力，减少对教师模型的过度依赖。

6. 代码示例：PyTorch 实现知识蒸馏

import torch
import torch.nn as nn
import torch.nn.functional as F

class DistillationLoss(nn.Module):
    def __init__(self, teacher_model, student_model, temperature=2.0, alpha=1.0, beta=1.0, gamma=1.0):
        super().__init__()
        self.teacher = teacher_model
        self.student = student_model
        self.T = temperature
        self.alpha = alpha
        self.beta = beta
        self.gamma = gamma

    def forward(self, inputs, labels):
        # 教师模型输出（假设包含logits、hidden_states、attentions）
        with torch.no_grad():  # 冻结教师模型参数
            teacher_outputs = self.teacher(**inputs, output_hidden_states=True, output_attentions=True)
            teacher_logits = teacher_outputs.logits
            teacher_hiddens = teacher_outputs.hidden_states
            teacher_attns = teacher_outputs.attentions
        
        # 学生模型输出
        student_outputs = self.student(**inputs, output_hidden_states=True, output_attentions=True)
        student_logits = student_outputs.logits
        student_hiddens = student_outputs.hidden_states
        student_attns = student_outputs.attentions

        # 1. 软标签损失（KL散度）
        teacher_soft = F.softmax(teacher_logits / self.T, dim=-1)
        student_soft = F.softmax(student_logits / self.T, dim=-1)
        loss_soft = F.kl_div(student_soft.log(), teacher_soft, reduction='batchmean') * (self.T ** 2)

        # 2. 特征匹配损失（MSE）
        loss_feat = 0.0
        num_layers = min(len(teacher_hiddens), len(student_hiddens))
        for th, sh in zip(teacher_hiddens[:num_layers], student_hiddens[:num_layers]):
            loss_feat += F.mse_loss(th, sh)
        loss_feat /= num_layers

        # 3. 注意力蒸馏损失（KL散度）
        loss_attn = 0.0
        num_attn_layers = min(len(teacher_attns), len(student_attns))
        for ta, sa in zip(teacher_attns[:num_attn_layers], student_attns[:num_attn_layers]):
            ta_avg = ta.mean(dim=1)  # 平均多头注意力
            sa_avg = sa.mean(dim=1)
            loss_attn += F.kl_div(sa_avg.log(), ta_avg, reduction='batchmean')
        loss_attn /= num_attn_layers

        # 总损失
        total_loss = self.alpha * loss_soft + self.beta * loss_feat + self.gamma * loss_attn
        return total_loss

# 示例用法
if __name__ == "__main__":
    # 假设教师和学生模型都是Transformer模型
    teacher = nn.TransformerEncoder(nn.TransformerEncoderLayer(512, 8), num_layers=12)
    student = nn.TransformerEncoder(nn.TransformerEncoderLayer(256, 4), num_layers=6)
    distill_loss = DistillationLoss(teacher, student)
    
    inputs = {
        "src": torch.randn(2, 100, 512),  # batch=2, seq_len=100, embed_dim=512
        "mask": torch.zeros(2, 100, dtype=torch.bool)
    }
    loss = distill_loss(inputs, None)  # 此处labels未使用，仅用于蒸馏
    print(f"Distillation loss: {loss.item()}")