自蒸馏：用自身深层指导浅层或历史模型指导当前

text 输入 → 浅层特征 → [辅助分类器1] ⇠ 蒸馏损失 ⇢ 深层特征 → 主分类器 ↓ …… ↓ 深层特征 → [辅助分类器k] ⇠ 蒸馏损失 ⇢ 主分类器

损失函数由两部分组成：各个辅助分类器与真实标签的交叉熵损失，以及辅助分类器输出与主分类器软标签之间的KL散度损失。最终目标可写为：

$$
\mathcal{L} = \sum_{i=1}^{K} \left( \alpha \cdot \mathcal{L}_{CE}(p_i, y) + \beta \cdot \mathcal{L}_{KL}(p_i, p_{final}) \right) + \mathcal{L}_{CE}(p_{final}, y)
$$

其中 $p_i$ 是第 $i$ 个辅助分类器的预测概率，$p_{final}$ 是最终分类器的输出，$y$ 是真值标签，$\alpha$ 和 $\beta$ 为平衡系数。

这种方式迫使浅层学会提取更具语义的特征，而不是仅仅捕捉局部纹理。实验表明，即使训练结束后移除辅助分类器，主模型的性能依然能得到提升。

### 无额外分类器的特征蒸馏

为了防止引入辅助分类器带来的结构侵入，近几年的工作倾向于**直接对齐中间层特征**。一种典型做法是计算深层特征图与浅层特征图之间的相似度矩阵，然后用均方误差或KL散度拉近它们。

例如，对于第 $l$ 层的特征图 $F_l \in \mathbb{R}^{C \times H \times W}$，可以先转化为格拉姆矩阵或注意力图，然后让浅层特征模仿深层特征的分布：

$$
\mathcal{L}_{feat} = \sum_{l} \| \text{Attn}(F_l) - \text{Attn}(F_{L}) \|_2^2
$$

其中 $\text{Attn}(\cdot)$ 表示将特征映射为注意力图的操作（如空间平均后的Softmax归一化）。这种约束使得浅层开始关注与深层相似的关键区域，增强了信息传递效率。

## 历史模型指导当前：穿越时间的教师

另一种自蒸馏范式不依赖网络深度，而是在时间维度上利用**训练历史**。其核心假设是：模型在早期训练阶段可能恰好探索到更优的权值空间，用那时的“自己”来指导后期训练，能抑制过拟合和遗忘。

### 时序平均教师

一种经典实现是**指数移动平均（EMA）**。在训练过程中维护一组参数的滑动平均 $\theta_{EMA}$：

$$
\theta_{EMA} \leftarrow \gamma \theta_{EMA} + (1-\gamma) \theta
$$

$\theta$ 为当前模型参数，$\gamma$ 通常设为0.999或更高。这个平均模型本身就是历史的聚合，更能表征训练轨迹中的稳定特征。然后将当前模型的预测与EMA模型的预测进行蒸馏：

$$
\mathcal{L}_{distill} = \text{KL}(p_{\theta} \| p_{\theta_{EMA}})
$$

此方法不但无需存储额外的历史检查点，还能在几乎不增加计算开销的情况下提升泛化性，在半监督和自监督学习中被广泛使用。

### 最优检查点蒸馏

更直接的方式是在训练过程中周期性保存验证集上表现最佳的模型快照，然后以该快照为教师蒸馏后续的迭代。这种方式被称为**SnapShot Distillation**或**历史重放蒸馏**。蒸馏损失通常为：

$$
\mathcal{L} = \mathcal{L}_{CE}(p_\theta, y) + \lambda \cdot \mathcal{L}_{KL}(p_\theta, p_{\theta^*})
$$

其中 $\theta^*$ 是之前保存的最佳参数。由于教师模型固定，这种方式比EMA更稳定，但会引入额外的存储和加载开销。不少研究会在训练后期启动蒸馏，避免早期教师质量不足带来的负面影响。

## 自蒸馏为何有效：理论视角

为什么看似“循环论证”的自蒸馏能生效？目前学界给出了几个解释：

- **正则化效应**：深层提供的软标签或特征分布约束，相当于一种强力的正则化项，防止模型在有限数据上陷入尖锐极小值，使损失曲面更加平滑。
- **多视角一致性**：深层与浅层相当于不同抽象层次上的“专家”，强制它们达成一致可以看作是一种多尺度数据增强，促使模型学习内在不变性。
- **暗知识提取**：Logits之间的相对大小（即使对于错误类别）包含了类别间的相似性信息，深层能更好地捕捉这种暗知识，用于纠正浅层的预测偏差。
- **自我纠正能力**：历史教师往往在验证集上误差更低，用它指导当前模型，相当于不断将模型拉回可能泛化更好的参数区域。

## 动手实现：一个简单的自蒸馏示例

以下为基于PyTorch的深层指导浅层自蒸馏的伪代码实现，使用一个带辅助分类器的简单卷积网络。

```python
import torch
import torch.nn as nn
import torch.nn.functional as F

class SelfDistillCNN(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 32, 3)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(32, 64, 3)
        self.fc = nn.Linear(64 * 6 * 6, 128)
        self.classifier = nn.Linear(128, num_classes)
        # 浅层辅助分类器（接在conv1之后）
        self.aux_classifier = nn.Linear(32 * 14 * 14, num_classes) 
        
    def forward(self, x):
        x1 = self.pool(F.relu(self.conv1(x)))   # 浅层特征
        out_aux = self.aux_classifier(x1.view(x1.size(0), -1))
        
        x2 = self.pool(F.relu(self.conv2(x1)))
        x2 = x2.view(x2.size(0), -1)
        feat = F.relu(self.fc(x2))
        out_final = self.classifier(feat)
        return out_aux, out_final, feat

def self_distillation_loss(out_aux, out_final, target, T=3.0, alpha=0.5):
    # 硬标签损失
    ce_loss = F.cross_entropy(out_aux, target) + F.cross_entropy(out_final, target)
    # 软标签蒸馏：深层指导浅层
    soft_teacher = F.softmax(out_final / T, dim=1)
    soft_student = F.log_softmax(out_aux / T, dim=1)
    distill_loss = F.kl_div(soft_student, soft_teacher, reduction='batchmean') * (T * T)
    return alpha * ce_loss + (1 - alpha) * distill_loss

在训练循环中，只需调用该损失函数即可：

model = SelfDistillCNN()
optimizer = torch.optim.Adam(model.parameters())

for data, target in dataloader:
    optimizer.zero_grad()
    out_aux, out_final, _ = model(data)
    loss = self_distillation_loss(out_aux, out_final, target)
    loss.backward()
    optimizer.step()