理解投影器在知识蒸馏中的作用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教师网络和学生网络分别输出表示Ztft(x),Zsfs(x). Z_tf_t(x),\qquad Z_sf_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​∥Zs​Wp​−Zt​∥22​.对投影矩阵WpW_pWp​求梯度可得到连续形式的更新方向W˙p−∂D∂Wp−ZsTZsWpZsTZt. \dot{W}_p -\frac{\partial D}{\partial W_p} -Z_s^{\mathrm T}Z_sW_pZ_s^{\mathrm T}Z_t.W˙p​−∂Wp​∂D​−ZsT​Zs​Wp​ZsT​Zt​.定义学生特征自相关矩阵和师生互相关矩阵CsZsTZs,CstZsTZt, C_sZ_s^{\mathrm T}Z_s,\qquad C_{st}Z_s^{\mathrm T}Z_t,Cs​ZsT​Zs​,Cst​ZsT​Zt​,则更新公式可以写为W˙pCst−CsWp. \dot{W}_pC_{st}-C_sW_p.W˙p​Cst​−Cs​Wp​.该公式说明投影器的每次更新都由两类关系信息决定CsC_sCs​描述学生表示内部不同维度之间的相关性CstC_{st}Cst​描述教师与学生表示之间的跨空间相关性。随着小批次训练不断进行WpW_pWp​会累积之前样本产生的相关性信息。因此不需要额外建立关系矩阵、Gram 矩阵或记忆库投影器本身就能充当一个可学习的关系编码器。4.3归一化如何影响投影器如果学生特征经过白化使其满足CsI, C_sI,Cs​I,那么投影器在固定点处满足Cst−CsWp0⟹WpCst. C_{st}-C_sW_p0 \quad\Longrightarrow\quad W_pC_{st}.Cst​−Cs​Wp​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​αp​W˙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∑​∣Zs​Wp​−Zt​∣iα​,其中iii遍历特征误差元素α\alphaα是平滑系数。对数求和形成一种软最大函数误差较大的元素获得更强关注而已经较好匹配的元素贡献会被相对弱化。论文发现该方法对α\alphaα较为鲁棒在不同教师—学生组合中α4∼5\alpha4\sim5α4∼5通常效果最好后续大规模实验主要采用α4\alpha4α4。4.5总体训练目标学生模型同时优化原任务损失和特征蒸馏损失LLtaskD(Zs,Zt;Wp). \mathcal{L} \mathcal{L}_{\text{task}} D(Z_s,Z_t;W_p).LLtask​D(Zs​,Zt​;Wp​).论文未在该公式中额外引入蒸馏权重。实际实现时先冻结教师模型再联合训练学生骨干和投影器。4.7不同任务下的归一化方式图像分类和 Transformer 蒸馏使用不带 affine 参数的 BatchNorm并设置ϵ10−4\epsilon10^{-4}ϵ10−4目标检测由于检测训练的 batch size 较小论文不在 batch 维归一化而是沿特征图的高度和宽度维度进行空间归一化投影层选择主要大规模实验采用线性投影器CIFAR100 实验采用隐藏维度为 1024 的 MLP 投影器。{-4}$2.目标检测由于检测训练的 batch size 较小论文不在 batch 维归一化而是沿特征图的高度和宽度维度进行空间归一化3.投影层选择主要大规模实验采用线性投影器CIFAR100 实验采用隐藏维度为 1024 的 MLP 投影器。