5.2 NLP 领域与预训练模型压缩

5.2 NLP 领域与预训练模型压缩

本节摘要：NLP 是蒸馏的第二个成熟主场，标志性成果是预训练模型压缩——用教师的预训练过程本身当课堂，学生继承全部语言能力后只需轻量微调。本节拆解 DistilBERT 式三件套，走一遍预蒸馏加微调的两阶段流程，再用翻译任务展示序列蒸馏的特别处理。读完能对"大语言模型怎么变小"给出工程级回答。

视觉之后看 NLP。这个行当的蒸馏有独特分量：预训练模型的训练成本以卡月计，蒸馏是少数能把这笔成本"转卖"给小模型的技术。本节先讲预训练压缩的代表作，再看任务内的序列蒸馏。

预训练压缩三件套

DistilBERT 是这个方向的经典样板，它的做法浓缩成三件事：

一、初始化继承。 学生直接从教师身上抽层初始化——每两层取一层，参数不是从零学起，而是站在教师的骨架上。这一步决定了学生起点就在"会说话"的水平。

二、三合一蒸馏损失。 语言建模损失（掩码语言模型）、教师隐层表示的余弦对齐（每层取一个代表点）、教师 logits 的软化蒸馏（温度固定用 1 的变体，NLP 分布本就平缓）。三路损失互相牵制：只学 logits 会丢语言能力，只对齐隐层会丢任务判断。

三、语料级预蒸馏。 学生不拿微调后的教师学任务，而是拿预训练阶段的原始教师、在海量无标注语料上完成整个预训练过程的浓缩版。学生由此继承教师的语言能力本身，而不是某个任务的判断——之后微调任何下游任务都受益。

# 预蒸馏一步：掩码语言建模 + 隐层对齐 + logits 蒸馏 三合一
import torch
import torch.nn.functional as F

def pretrain_distill_step(student, teacher, batch, T=1.0,
                          w_lm=1.0, w_hid=1.0, w_logit=1.0):
    teacher.eval()
    x, attn, labels, mask_pos = batch  # labels 为 MLM 标签，mask_pos 为掩码位置
    s_out = student(x, attention_mask=attn, output_hidden_states=True)
    with torch.no_grad():
        t_out = teacher(x, attention_mask=attn, output_hidden_states=True)
    # 分量一：语言建模损失（只在掩码位置算）
    lm = F.cross_entropy(
        s_out.logits[mask_pos].view(-1, s_out.logits.size(-1)),
        labels[mask_pos].view(-1))
    # 分量二：隐层余弦对齐（取各层中间点，师生按深度等比取）
    n = min(len(s_out.hidden_states), len(t_out.hidden_states))
    hid = 0.0
    for i in range(1, n):
        j = round(i * (len(t_out.hidden_states) - 1) / (n - 1))
        h_s = F.normalize(s_out.hidden_states[i], dim=-1)
        h_t = F.normalize(t_out.hidden_states[j], dim=-1)
        hid = hid + (1 - (h_s * h_t).sum(-1)).mean()
    # 分量三：logits 蒸馏（掩码位置）
    logit = F.kl_div(
        F.log_softmax(s_out.logits[mask_pos] / T, dim=-1),
        F.softmax(t_out.logits[mask_pos] / T, dim=-1),
        reduction="batchmean") * T * T
    loss = w_lm * lm + w_hid * hid + w_logit * logit
    loss.backward()
    return {"total": loss.item(), "lm": lm.item(),
            "hidden": hid.item() if torch.is_tensor(hid) else hid,
            "logits": logit.item()}

# 输出说明：三项分量分 monitored，任何一项占比异常都提示配比要调；
# NLP 分布平缓，蒸馏温度从 1 起步而不是 4。
​

一条完整案例：六层学生接手十二层教师的活

背景。 团队的文本三分类（质检好评差评中性）服务用的是十二层预训练模型微调版，线上平均延迟 85 毫秒；业务要做实时质检，预算 25 毫秒。

操作。 走两阶段。阶段一预蒸馏：学生六层（从教师抽层初始化），在 30 GB 领域语料上跑三合一蒸馏损失，温度 1，三路权重 1 比 1 比 1，训练预算为教师原预训练成本的十分之一。阶段二任务微调：用业务标注数据（4 万条）按常规方式微调六层学生，微调超参与十二层教师微调完全一致以保对照公平。

结果。 学生测试集 92.4%，教师微调版 92.9%，差距 0.5 个点；线上平均延迟 21 毫秒，达标。同规模直接预训练加微调的对照组 90.7%——预蒸馏带来的 1.7 个点就是"继承语言能力"的市价。

解读。 两个数字值得咀嚼：0.5 的师生差在预训练压缩里属于优秀区间（常规预期 1 到 2 个点），归功于初始化继承——学生不是模仿教师的学生，而是"半个教师本身"；十分之一的预训练成本则说明蒸馏把最贵的那笔账（语言能力的习得）从"自己挣"变成了"继承遗产"。

变式。 预算更紧时跳过阶段一，直接在教师微调版上做任务内蒸馏（见下），损失若干泛化余量；层数抽不满时（教师 12 层学生 4 层），抽取步长加大并在预蒸馏时给深层对齐更高权重；解码式任务（翻译、摘要）走本节末的序列蒸馏变体。

任务内蒸馏：序列任务的特别处理

不走预蒸馏、只压任务模型的路线叫任务内蒸馏。序列任务（翻译、摘要、对话）与分类的差别在于输出是一个分布序列，且推理时存在"曝光偏差"——学生训练时看的是真前文，推理时看的是自己生成的前文，两个分布对不上。处理办法是序列级蒸馏：让学生先自由生成完整序列，教师对这份序列（而不是标准答案）打分，学生学"怎么把自己的草稿改到教师的水平"。

# 序列级蒸馏：教师对学生自己生成的候选序列打分
import torch
import torch.nn.functional as F

def seq_level_kd_step(student, teacher, src_ids, hyp_ids, T=2.0, alpha=0.7):
    """hyp_ids 是学生自回归生成的候选译文（自由解码得来）。"""
    teacher.eval()
    s_out = student(src_ids, decoder_input_ids=hyp_ids[:, :-1])
    with torch.no_grad():
        t_out = teacher(src_ids, decoder_input_ids=hyp_ids[:, :-1])
    # 对"学生自己的草稿"逐 token 对齐分布——而不是对标准答案
    soft = F.kl_div(
        F.log_softmax(s_out.logits / T, dim=-1),
        F.softmax(t_out.logits / T, dim=-1),
        reduction="batchmean") * T * T
    return alpha * soft  # 剩余权重给常规 token 级硬损失（对参考译文）

# 输出说明：这个改动把训练分布与推理分布对齐——
# 学生学到的是"在自主生成路径上的每一步该怎么想"，
# 而不是"在标准答案路径上抄作业"，翻译 BLEU 的提升主要来自这里。
​

一条压缩版案例：某团队的中英翻译模型（教师 BLEU 32.1）压六层学生。纯 token 级蒸馏 29.8；换序列级蒸馏（教师对学生自由生成的候选打分）后 31.2；再把候选集扩到四条取教师分最高者，31.9。曝光偏差的修正贡献了两个 BLEU，这在翻译任务里是量级惊人的收益。

⚠️ 常见坑：NLP 蒸馏照搬视觉的温度经验。语言模型的 token 分布天生平缓（大量同义表达），温度 4 会把分布摊成粥；NLP 的温度从 1 起步，且常只对高置信 token 计损失。

💡 关键直觉：NLP 蒸馏分两条预算档位——预算够就预蒸馏加微调（继承语言能力，泛化最好），预算紧就任务内蒸馏加序列级修正（见效快，泛化略窄）。选档的第一问：未来还会接几个新任务？接得多，就值一笔预蒸馏的学费。

本节要点回顾

• 三件套：抽层初始化、三合一损失（语言建模加隐层对齐加 logits 蒸馏）、语料级预蒸馏。
• 两阶段流程：预蒸馏继承语言能力，再常规微调接任务；师生差可压到 0.5 个点内。
• 成本账：预蒸馏约花教师预训练十分之一的算力，换来学生承接任意下游任务的能力。
• 序列任务：序列级蒸馏修正曝光偏差——教师批改学生的自由草稿，收益可达两个 BLEU。
• 温度差异：NLP 从 1 起步，视觉经验别照搬。