理解投影器在知识蒸馏中的作用
01论文信息
论文标题:Understanding the Role of the Projector in Knowledge Distillation
论文作者:Roy Miles、Krystian Mikolajczyk
论文发表单位:Imperial College London(帝国理工学院)
论文期刊\会议:AAAI 2024
论文代码:未开源
02论文主要贡献
- 重新解释投影器的作用:论文指出,投影层并不只是用于对齐教师与学生的特征维度。其梯度更新会持续累积学生特征的自相关信息以及教师—学生特征的互相关信息,因此一个简单的可学习投影器就能隐式编码历史样本之间的关系。
- 揭示投影器与归一化之间的耦合关系:不同归一化方式会改变投影矩阵的训练轨迹和最终解。合适的归一化能够减少投影矩阵奇异值向零收缩,降低维度坍缩和信息丢失,从而显著改善学生模型性能。
- 使用 LogSum 缓解教师—学生容量差距:论文引入简单的 LogSum 软最大距离,减弱已经较好对齐样本对损失的影响,使优化更加关注难以对齐的特征,在教师与学生容量差距较大时尤其有效。
- 提出低成本通用蒸馏方案:通过“投影器 + 归一化 + LogSum 距离”三个组件,在 CIFAR100、ImageNet、COCO2017以及数据高效 Transformer 训练中取得与先进蒸馏方法相当或更好的结果,同时避免显式构造关系矩阵、特征核和大型记忆库。
03论文创新点
- 从训练动力学解释投影器:由均方误差下的投影矩阵更新公式证明,投影权重天然受到学生自相关矩阵和师生互相关矩阵控制,为投影器能够学习关系知识提供了理论解释。
- 用奇异值分析投影坍缩:论文追踪投影矩阵在不同归一化方式下的奇异值变化,将学生性能下降与投影空间中的维度收缩、信息丢失联系起来。
- 发现大投影网络并非一定更好:更复杂的 MLP 投影器虽然具有更高容量,却可能逐渐使其输入和输出特征去相关,导致投影器学习到无法传回学生骨干网络的知识。因此,大规模实验优先采用简单的线性投影器。
- 将容量差距转化为距离函数设计问题:LogSum 通过软最大机制重新分配不同特征误差的梯度权重,不要求学生严格匹配所有教师特征,因而能够缓解大模型教师对小模型学生造成的过强约束。
04方法
4.1整体框!架
分别送入学生网络和冻结的教师网络提取特征;学生特征先经过线性投影器映射到教师特征空间,该投影器不仅对齐通道维度,还能在训练中积累师生特征的关系信息;随后,投影后的学生特征与教师特征分别进行无仿射归一化,以统一特征尺度并减少投影器发生维度坍缩;最后通过 LogSum 距离比较两者,利用软最大机制重点关注较难匹配的特征,从而缓解教师与学生之间较大的容量差距,并将梯度反向传播给学生网络和投影器。
给定输入xxx,教师网络和学生网络分别输出表示:
Zt=ft(x),Zs=fs(x). Z_t=f_t(x),\qquad Z_s=f_s(x).Zt=ft(x),Zs=fs(x).
学生表示首先经过可学习投影器WpW_pWp,将其映射到教师表示空间。随后,教师特征和投影后的学生特征分别进行归一化,最后通过 LogSum 距离计算特征蒸馏损失。
完整蒸馏流程包含三个核心组件:
- 线性投影器:对齐特征维度,并隐式积累关系信息;
- 特征归一化:稳定梯度并抑制投影矩阵发生维度坍缩;
- LogSum 距离:以软最大方式处理各特征误差,缓解较大的师生容量差距。
4.2投影器为什么能够编码关系知识
论文首先考虑无偏置线性投影器和简单的ℓ2\ell_2ℓ2特征距离:
D(Zs,Zt;Wp)=12∥ZsWp−Zt∥22. D(Z_s,Z_t;W_p)=\frac{1}{2}\left\|Z_sW_p-Z_t\right\|_2^2.D(Zs,Zt;Wp)=21∥ZsWp−Zt∥22.
对投影矩阵WpW_pWp求梯度,可得到连续形式的更新方向:
W˙p=−∂D∂Wp=−ZsTZsWp+ZsTZt. \dot{W}_p =-\frac{\partial D}{\partial W_p} =-Z_s^{\mathrm T}Z_sW_p+Z_s^{\mathrm T}Z_t.W˙p=−∂Wp∂D=−ZsTZsWp+ZsTZt.
定义学生特征自相关矩阵和师生互相关矩阵:
Cs=ZsTZs,Cst=ZsTZt, C_s=Z_s^{\mathrm T}Z_s,\qquad C_{st}=Z_s^{\mathrm T}Z_t,Cs=ZsTZs,Cst=ZsTZt,
则更新公式可以写为:
W˙p=Cst−CsWp. \dot{W}_p=C_{st}-C_sW_p.W˙p=Cst−CsWp.
该公式说明,投影器的每次更新都由两类关系信息决定:
- CsC_sCs描述学生表示内部不同维度之间的相关性;
- CstC_{st}Cst描述教师与学生表示之间的跨空间相关性。
随着小批次训练不断进行,WpW_pWp会累积之前样本产生的相关性信息。因此,不需要额外建立关系矩阵、Gram 矩阵或记忆库,投影器本身就能充当一个可学习的关系编码器。
4.3归一化如何影响投影器
如果学生特征经过白化,使其满足
Cs=I, C_s=I,Cs=I,
那么投影器在固定点处满足
Cst−CsWp=0⟹Wp=Cst. C_{st}-C_sW_p=0 \quad\Longrightarrow\quad W_p=C_{st}.Cst−CsWp=0⟹Wp=Cst.
此时,投影矩阵直接编码教师与学生之间的互相关关系。对于一般归一化方式,CsC_sCs不再等于单位矩阵,投影器的固定点和训练轨迹都会发生改变。
考虑学习率αp\alpha_pαp和权重衰减η\etaη,投影器更新为
Wp←(1−η)Wp+αpW˙p. W_p\leftarrow(1-\eta)W_p+\alpha_p\dot{W}_p.Wp←(1−η)Wp+αpW˙p.
当η=αp\eta=\alpha_pη=αp时,该更新形式与关系特征的移动平均相似,这也解释了简单投影器为什么能够承担类似动量编码器或记忆库的功能。
论文进一步比较无归一化、ℓ2\ell_2ℓ2归一化、GroupNorm 和 BatchNorm。实验发现,性能更好的归一化方式会保留更多非零奇异值;较差的归一化会使更多奇异值收缩到零,相当于将输入投影到更低维空间,造成信息丢失。
4.4LogSum蒸馏距离
当教师远强于学生时,学生很难使所有特征完全对齐。如果仍使用普通均值距离,已经接近的特征和难以对齐的特征会共同影响优化,学生可能因追求无法实现的精确匹配而损害下游性能。
论文采用 LogSum 距离:
D(Zs,Zt;Wp)=log∑i∣ZsWp−Zt∣iα, D(Z_s,Z_t;W_p) =\log\sum_i \left|Z_sW_p-Z_t\right|_i^{\alpha},D(Zs,Zt;Wp)=logi∑∣ZsWp−Zt∣iα,
其中,iii遍历特征误差元素,α\alphaα是平滑系数。对数求和形成一种软最大函数:误差较大的元素获得更强关注,而已经较好匹配的元素贡献会被相对弱化。
论文发现该方法对α\alphaα较为鲁棒,在不同教师—学生组合中,α=4∼5\alpha=4\sim5α=4∼5通常效果最好,后续大规模实验主要采用α=4\alpha=4α=4。
4.5总体训练目标
学生模型同时优化原任务损失和特征蒸馏损失:
L=Ltask+D(Zs,Zt;Wp). \mathcal{L} =\mathcal{L}_{\text{task}} +D(Z_s,Z_t;W_p).L=Ltask+D(Zs,Zt;Wp).
论文未在该公式中额外引入蒸馏权重。实际实现时,先冻结教师模型,再联合训练学生骨干和投影器。
4.7不同任务下的归一化方式
- 图像分类和 Transformer 蒸馏:使用不带 affine 参数的 BatchNorm,并设置ϵ=10−4\epsilon=10^{-4}ϵ=10−4;
- 目标检测:由于检测训练的 batch size 较小,论文不在 batch 维归一化,而是沿特征图的高度和宽度维度进行空间归一化;
- 投影层选择:主要大规模实验采用线性投影器;CIFAR100 实验采用隐藏维度为 1024 的 MLP 投影器。
{-4}$;
2.目标检测:由于检测训练的 batch size 较小,论文不在 batch 维归一化,而是沿特征图的高度和宽度维度进行空间归一化;
3.投影层选择:主要大规模实验采用线性投影器;CIFAR100 实验采用隐藏维度为 1024 的 MLP 投影器。