4.3 自己教自己：自蒸馏

4.3 自己教自己：自蒸馏

本节摘要：自蒸馏让一个模型自己当自己的师傅：深层教浅层、上一代教下一代、历史快照教当前版本。它不需要外部教师，却常能带来实打实的收益。本节讲三代形态——深监督式、代际式（born-again）、快照式——各自的机制、成本与适用现场，并用一条完整案例走完代际自蒸馏的全流程。

第四章师资变体的第二种：师资来源不是别人，是模型自己。听起来像左口袋倒右口袋，但收益真实存在——它在分类、检测、检索多个行当都被反复验证过。

三代形态

深监督式：一个模型之内，深层教浅层。 训练时给网络的不同深度接出多个输出头，浅层的头在损失里向深层的头（以及真实标签）同时学习。深层在任何时刻都比浅层更接近答案，这个"深度差"就是天然的师生差。训练完成后浅层头全部丢弃，只留主干——等于让网络的浅层在训练期接受了额外辅导。实现代价几乎为零，是性价比最高的一种自蒸馏。

代际式：上一代教下一代。 先把模型正常训练到收敛，作为第一代教师；再用它的软化输出训练一个同结构甚至不同结构的学生作为第二代。born-again 网络的研究显示，第二代的精度常高于第一代，重复几代还能继续涨——同结构的学生学到的不只是教师的知识，还包括"以教师的输出为地形"的更好优化路径。

快照式：历史快照教当前版本。 训练过程中按学习率周期存若干快照，当前版本向历史快照的集成输出学习。一次训练产出师生两代，算力开销约一点几倍，介于深监督式与代际式之间。

快照式的核心是一个"快照池"：每个学习率周期结束存一份权重，当前版本向池子里所有成员的平均输出学习，池子随训练滚动更新——

# 快照池的最小实现：滚动集成当教师
import copy
import torch
import torch.nn.functional as F

class SnapshotPool:
    def __init__(self, max_snapshots=3):
        self.snaps = []                          # 历史快照列表
        self.max = max_snapshots

    def update(self, model):
        if len(self.snaps) >= self.max:
            self.snaps.pop(0)                    # 滚动淘汰最老快照
        self.snaps.append(copy.deepcopy(model).eval())

    @torch.no_grad()
    def ensemble_logits(self, x):
        return torch.stack([s(x) for s in self.snaps]).mean(0)

def snapshot_distill_step(model, pool, x, y, T=4.0, alpha=0.8):
    s_logits = model(x)
    if len(pool.snaps) == 0:                     # 池子还空：退化为直接训练
        return F.cross_entropy(s_logits, y)
    t_logits = pool.ensemble_logits(x)           # 历史快照的平均意见
    soft = F.kl_div(F.log_softmax(s_logits / T, dim=-1),
                    F.softmax(t_logits / T, dim=-1),
                    reduction="batchmean") * T * T
    hard = F.cross_entropy(s_logits, y)
    return alpha * soft + (1 - alpha) * hard

# 输出示例（余弦退火三个周期后的池状态）：
# 快照数: 3 | 成员验证分: [89.1, 90.2, 89.8] | 集成验证分: 90.9
# 集成分始终高于任何单个快照——这正是当前版本能持续学的"教材"。
​

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

class SelfDistillNet(nn.Module):
    """深监督式自蒸馏：主干加深层辅助头，训练后辅助头丢弃。"""
    def __init__(self):
        super().__init__()
        self.block1 = nn.Sequential(nn.Conv2d(3, 64, 3, padding=1), nn.ReLU())
        self.rest = nn.Sequential(
            nn.Conv2d(64, 128, 3, padding=1), nn.ReLU(),
            nn.Conv2d(128, 256, 3, padding=1), nn.ReLU(),
            nn.AdaptiveAvgPool2d(1), nn.Flatten(),
        )
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.head_deep = nn.Linear(256, 10)     # 深层主头（保留）
        self.head_shallow = nn.Linear(64, 10)   # 浅层辅导头（训练后丢弃）

    def forward(self, x, return_aux=False):
        f1 = self.block1(x)                     # 浅层特征：B x 64 x H x W
        feat = self.rest(f1)                    # 深层特征：B x 256
        if return_aux:
            return self.head_deep(feat), self.head_shallow(self.pool(f1).flatten(1))
        return self.head_deep(feat)

def self_distill_step(model, x, y, T=4.0, gamma=0.3):
    deep_logits, shallow_logits = model(x, return_aux=True)
    loss_deep = F.cross_entropy(deep_logits, y)
    # 浅层头向深层头学：同模型内的师生差
    loss_shallow = F.kl_div(
        F.log_softmax(shallow_logits / T, dim=-1),
        F.softmax(deep_logits.detach() / T, dim=-1),
        reduction="batchmean") * T * T
    # 浅层头也学真实标签，防它只当"复读机"
    loss_shallow_hard = F.cross_entropy(shallow_logits, y)
    loss = loss_deep + gamma * (loss_shallow + 0.5 * loss_shallow_hard)
    loss.backward()
    return loss.item()

# 输出说明：深层头 detach 后当师傅；gamma 控制辅导开销占比，
# 部署时只跑 head_deep，浅层头直接删掉，推理零开销。

m = SelfDistillNet()
x = torch.randn(4, 3, 32, 32); y = torch.randint(0, 10, (4,))
self_distill_step(m, x, y)
print("前向输出形状:", m(x).shape)
# 输出示例：
# 前向输出形状: torch.Size([4, 10])
​

图 4-2 三代自蒸馏的师资来源

一条完整案例：代际自蒸馏连传两代

背景。 检索业务的 embedding 模型已训练到 89.3%（召回口径），预算不允许引入更大的教师，业务方想知道还有没有免费的增长空间。

操作。 走代际自蒸馏。第一代（89.3%）对全量训练数据离线打分（温度 3，检索任务分布平缓用低温）；第二代从随机初始化训练，损失为硬标签交叉熵加向第一代软化输出的 KL（alpha 0.8）。第二代收敛后同样打分，再教第三代。

结果。 第二代 90.1%，第三代 90.4%，第四代 90.3% 开始回落。两代净增 1.1 个点，全程没有引入任何外部模型。

解读。 收益来源正如 3.3 的机制清单：第一代的软目标为第二代提供了更平滑的地形与更高的信息密度，哪怕师生结构完全相同。增益代际递减且第四代回落——软目标的暗知识在第二、三代基本榨干，继续传只会放大教师自身的误差（错误继承机制压过正则机制）。工程纪律：代际自蒸馏一般最多传两三代，逐代盯验证集，回落即停。

变式。 深监督式与代际式可以叠用（训练第二代时同时开深层辅导头）；快照式适合配余弦退火的学习率周期，每个谷底存快照；若嫌代际式训练时间翻倍，第一代可以选择训练到 90% 进度（不彻底收敛）就开教——实验显示终点成绩几乎不掉、总时间省两成。

⚠️ 常见坑：代际传多了误差滚雪球。软目标里的错误每传一代被放大一分，第二代赚的是地形红利，第四代以后开始还误差的债。逐代盯验证集、见顶即停是铁律。

💡 关键直觉：自蒸馏证明了一件事——蒸馏收益的来源不是"师傅比徒弟大"，而是"软目标比硬标签好"。师生差可以来自深度差、代际差、时间差，唯独不必来自参数量差。

本节要点回顾

• 三代形态：深监督式（零部署开销）、代际式（born-again，每代可涨）、快照式（一次训练出师生）。
• 师生差来源：深度差、代际差、时间差，不必是参数量差。
• 代际纪律：最多传两三代，逐代盯验证集，回落即停，防误差滚雪球。
• 省时变体：第一代训练到九成即开教，终点几乎无损、总时间省两成。
• 收益同源：与普通蒸馏共享 3.3 的三个机制，代价结构不同而已。