4.4 师徒过招：对抗蒸馏

4.4 师徒过招：对抗蒸馏

本节摘要：对抗蒸馏在普通的模仿损失之外引入一个判别器，让它判断"这份输出是教师打的还是学生打的"；学生的目标升级为骗过裁判。学生因此不再逐点抄答案，而是抓教师输出分布的"神韵"。本节讲对抗损失的构造、与模仿损失的配合、训练稳定性的四个技巧，以及哪些任务值得上这套更贵的教法。

普通蒸馏是"照着描"，对抗蒸馏是"描到裁判看不出真假"。师资关系从师徒单向传授变成师徒加裁判的三方博弈，这是第四章师资变体的第三种。

为什么要加一个裁判

逐点模仿损失（KL、MSE）有个先天局限：它奖励"每个点都贴近"，却不惩罚"整体气质不像"。就像临帖，笔画都对，气韵全无。判别器补的就是这一环：它见过真迹（教师输出），学生对假迹（自己输出）做小扰动，只要气韵有差，裁判就能分辨。学生为骗过裁判，被迫捕捉那些逐点损失度量不到的分布细节——多模态结构、长尾的分布形状、样本之间的联合气质。

损失由三项构成：任务损失（对真实标签）、模仿损失（对教师输出的经典蒸馏项）、对抗损失（判别器判学生输出为真时的负对数）。前两项稳住下限，第三项向上限发力。

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

class OutputDiscriminator(nn.Module):
    """轻量判别器：分辨一份 logits 来自教师还是学生。"""
    def __init__(self, dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim, 256), nn.LeakyReLU(0.2),
            nn.Linear(256, 128), nn.LeakyReLU(0.2),
            nn.Linear(128, 1),
        )
    def forward(self, logits):
        return self.net(logits)

def adversarial_kd_step(student, teacher, disc, x, y, T=4.0, alpha=0.5, lam=0.1):
    teacher.eval()
    s_logits = student(x)
    with torch.no_grad():
        t_logits = teacher(x)
    # 学生侧：任务 + 模仿 + 骗过裁判
    task = F.cross_entropy(s_logits, y)
    imitate = F.kl_div(F.log_softmax(s_logits / T, dim=-1),
                       F.softmax(t_logits / T, dim=-1),
                       reduction="batchmean") * T * T
    fool = F.binary_cross_entropy_with_logits(
        disc(s_logits),
        torch.ones(s_logits.size(0), 1, device=x.device))
    (task + alpha * imitate + lam * fool).backward()

    # 裁判侧：教师输出判真，学生输出判假
    d_loss = (F.binary_cross_entropy_with_logits(disc(t_logits.detach()),
              torch.ones(x.size(0), 1, device=x.device)) +
              F.binary_cross_entropy_with_logits(disc(s_logits.detach()),
              torch.zeros(x.size(0), 1, device=x.device)))
    d_loss.backward()

# 输出说明：学生与裁判各自 backward 一次，用两套优化器分别 step。
# lam 是对抗项权重，起点 0.05 到 0.2；调大了训练极易不稳。
​

训练稳定性四技巧

对抗成分让蒸馏训练从"爬山"变成"拔河"，稳定性技巧必须备齐：

1. 权重从最小起步：lam 从 0.05 开始，确认模仿损失不恶化再逐步加；对抗项权重超 0.3 是绝大多数翻车现场的原因。
2. 裁判弱于师生：判别器容量刻意做小（一两层即可），裁判太强会把学生逼到模仿教师的高频噪声。
3. 标签平滑：裁判的真假标签加 0.1 平滑，防止裁判过度自信后梯度消失。
4. 熔断机制：监控模仿损失，若对抗项开启后它连续数轮上升，立即衰减 lam——保住模仿下限比追求气韵重要。

熔断机制值得给一段可直接搬的实现——它是对抗蒸馏唯一必须自动化的部分，靠人盯日志既累又慢半拍：

# 模仿损失熔断：对抗项伤到主干时自动衰减 lam
import collections

class LamCircuitBreaker:
    def __init__(self, lam, patience=3, decay=0.5, floor=0.01):
        self.lam, self.patience = lam, patience
        self.decay, self.floor = decay, floor
        self.best = float("inf")
        self.bad = collections.deque(maxlen=patience)

    def step(self, imitate_loss):
        self.best = min(self.best, imitate_loss)
        self.bad.append(imitate_loss > self.best * 1.05)  # 比最优差 5% 记一次
        if len(self.bad) == self.patience and all(self.bad):
            self.lam = max(self.lam * self.decay, self.floor)
            self.bad.clear()                # 衰减后重新观察
            print(f"熔断触发，lam 降至 {self.lam:.3f}")
        return self.lam

cb = LamCircuitBreaker(lam=0.1)
# 输出示例（模仿损失连续劣化的三步后）：
# 熔断触发，lam 降至 0.050
# 再恶化一次会降到 0.025，floor 0.01 兜底——对抗项被压到
# "只锦上添花、绝不雪中送炭"的位置。
​

一条完整案例：文本生成模型的学生怎么教

背景。 团队有一个自回归文本生成大模型，要压一个端上可跑的学生。文本生成的输出是逐 token 的分布序列，逐点 KL 只能教"下一个词的概率排布"，教不会句式的气韵——生成的句子用词都对、节奏生硬。

操作。 在常规 token 级蒸馏之上叠对抗项：判别器输入为一个句子的 token 级分布序列特征（过一层均值池化），判断是教师生成还是学生生成。lam 从 0.05 起步，每 10 轮评估一次，模仿损失无恶化则升到 0.1、0.15 封顶。判别器只有两层，且真假标签做 0.1 平滑。

结果。 纯 token 蒸馏的学生，自动评测分与教师差 4.7；加对抗项后差距缩到 3.1；人工盲评里"生硬感"维度从 2.9 提到 3.8（5 分制）。训练全程未触发熔断。

解读。 收益集中在"气韵"维度而非事实维度，印证了对抗项的作用位置：它教的不是下一个词对不对，而是整句分布的形态像不像师傅。也注意代价：这套训练比普通蒸馏多一个组件、约四成的训练时间、以及每个变体都要重新调 lam——它是第四章教法里工程成本最高的之一。

变式。 判别器除了看输出分布，还可以看特征（2.2 的 hint 特征上做对抗，气韵检查前移到中间层）；样本级对抗（判别器分辨整批样本的真实性）比 token 级更稳，适合初试；若只是想给分布"补形"，成本更低的做法是 4.5 的多教师或关系蒸馏，先排除便宜方案再上对抗。

⚠️ 常见坑：把对抗蒸馏当普通蒸馏的超参直接迁移。lam、判别器容量、标签平滑三个变量彼此联动，单变量扫描的做法在这里容易误判——建议每次只动一个，并始终以模仿损失不恶化为约束条件。

💡 关键直觉：模仿损失教"对不对"，对抗损失教"像不像"。任务只在乎对错时别用对抗；任务的验收里存在"气质"成分（生成、风格、节奏）时，裁判才有用武之地。

本节要点回顾

• 三方结构：学生、教师、判别器；学生目标从贴近教师升级为骗过裁判。
• 损失三件套：任务损失保下限、模仿损失保主干、对抗损失攻上限，lam 从 0.05 起步。
• 稳定四招：权重最小起步、裁判刻意做弱、真假标签平滑、模仿损失熔断。
• 收益位置：分布在"气质"维度的任务（生成、风格）收益最大，纯判别任务通常不值得。
• 成本意识：多一个组件、四成额外训练时间、三个联动超参，先穷尽便宜方案。