CVPR 2024 | 知识蒸馏中的Logit标准化：辅助logit-based KD算法稳定涨点
太长不看的一句话总结：传统知识蒸馏默认学生/教师网络的温度是全局一致的，这种设置迫使学生模仿教师的 logit 的具体值，而非其关系，论文方法提出 logit 标准化，解决了这个问题。
论文链接：
开源代码：
项目主页：
https://sunsean21.github.io/logit-stand-KD.html
知识蒸馏（Knowledge Distillation）
该部分是给不熟悉的读者科普的，熟悉的读者可以看一下符号表达就跳过。
2015 年，Hinton #ref_1 注意到深度学习模型变得越来越大，率先想到是否可以利用一个训练好的大模型（俗称 Teacher、教师模型），教授一个小模型（俗称 Student、学生模型）进行学习。
整个过程可以看下面的图 1a。
上面简单的概述内容还是有所局限，大家也可以看一下其他科普文章，例如：
潘小小：【经典简读】知识蒸馏（Knowledge Distillation）经典之作（https://zhuanlan.zhihu.com/p/102038521）。
问题背景与分析
距离 Hinton #ref_1 2015 年提出知识蒸馏已经过去了 9 年，温度这个超参数最早就被设定为教师、学生之间共享的，并且是对所有样本都全局不变的，而这样的设置并没有理论支持。
已有工作中，CTKD #ref_2 引入了对抗学习，针对不同难度的样本选择不一样的温度，但是它仍然让学生和教师共享温度；ATKD #ref_3 则是引入一种锐利度指标，针对性地选择学生和教师之间的温度平衡。但是他们均没有讨论学生和教师的温度来自于哪里，也没有理论性地讨论它们是否可以全局不一致、学生教师之间不共享。
因此针对这个问题，文章做出了三个贡献：
文章基于信息论中的熵最大化理论，推导了含有超参数温度的 softmax 函数表达式。基于推导过程，发现温度并没有明显的约束条件，即没有理论强制学生和教师在全局范围内共享温度（见 1.1.1、1.1.2）
文章发现已有的“共享温度”的设置会导致两个问题：
学生网络被迫输出和教师网络相当的 logit（见1.2）
KL 散度难以真实反映学生的蒸馏效果（见2.3）
文章提出了 logit 标准化，可作为预处理辅助现有基于 logit 的知识蒸馏算法
注：本文初次于 2023 年 3 月投稿至 ICCV 2023，与 2023 年 8 月公开的 NormKD #ref_4 属于 concurrent work。
▲ 图1. logit 标准差辅助的知识蒸馏和传统知识蒸馏的对比
2.1 超参数-温度的来源
这部分推导了带有超参数温度的 softmax 函数，如果只想看结论，可直接跳到 1.3 小节。
2.1.1 教师网络的温度
基于 Edwin #ref_5 1957 年提出的最大熵理论（Maximum-Entropy Principle），可以证明出分类任务的 softmax 函数是求解条件熵最大化问题的唯一解形式，也就是说，对于一个训练好的教师网络，面对以下两个条件，其概率密度 应该使得熵处于最大值，两个条件分别为
变量向量 需要求和为 1。也就是说，它满足概率密度的形式。 原始 logit 向量 的期望为正确的 logit 项。也就是说，它需要预测正确。
其数学表达为以下式子：
针对这个求解问题，可以利用拉格朗日乘子法，引入拉格朗日乘子 （条件1）和 （条件2），将条件优化变为单一表达式：
对目标函数求导后取零，即可得到优化问题的解的形式。于是我们对 求偏导，得到：
对其取 0，我们得到
2.1.2 学生网络的温度
学生 logit 向量的期望需要等于教师 logit 的期望。也就是说，需要学生学习到教师的知识。
类似地引入拉格朗日乘子 （条件1）、（条件2）、（条件3），可以将条件优化变为单一表达式：
对 求偏导，得到：
对其取 0，并且设 ，我们得到
问题1：学生和教师之间是否可以取不同的温度？
问题2：蒸馏过程中是否可以对不同样本取不同的温度？
2.2 共享温度的弊端（1/2）
其中， 是标准差函数。
假设我们设定教师和学生的温度相同，即 ，那么公式（14）就会变为：
由此可以看出经典知识蒸馏问题中共享温度的设定，最终会强制学生和教师 logit 之间存在一个固定的差 ，但是考虑到学生网络和教师网络之间的能力差距，学生很难生成与教师 logit 的均值相当的 logit #ref_6 #ref_7 #ref_8（见图 2 横轴）。
我们继续考虑教师和学生的温度相同的情况下的标准差，将公式（15）进行简化，可以得到：
由此可以看出经典知识蒸馏问题中共享温度的设定，最终也会强制学生和教师输出的 logit 标准差一致，考虑到其二者的能力差距，学生同样很难生成与教师 logit 的标准差相当的 logit（见图 2 纵轴）。
2.3 小节
我们从本节可以得到以下结论：
学生、教师网络的温度没有明显的约束条件，不必一定全局共享，可以人为指定
在温度共享的情况下，学生、教师网络之间有一个 logit 范围的强制性匹配
基于学生、教师网络的能力差距， 上述强制性匹配很可能限制学生网络的蒸馏效果
提出方法：Logit标准化
3.1 算法
3.2 优势
均值为零
之前的工作 #ref_1 #ref_3 常常有学生/教师 logit 的均值为 0 的假设，但是如图 2 所示，其几乎是不可能真实实现的，而基于提出的 logit 标准化函数，logit 均值会自动变为 0。
标准差为
单调性
有界性，上下界为
3.3 Toy Case：共享温度的弊端（2/2）
▲ 图3. 一个 toy case，展示了“温度共享”带来的KL散度无法准确评估学生网络性能的情况
实验结果
4.1 数值结果
▲ 表1. CIFAR-100 上面的蒸馏结果，此表为学生、教师的网络结构不同的情况
4.2 消融实验
4.3 可视化结果
logit 范围：
从图 4 第一行中，可以看出经典知识蒸馏算法可以迫使学生产生与教师 logit 相仿的 logit 范围，但是其预测的最大 logit 值没有达到教师的值（12），预测产生了错误。而文章的方法可以使原始 logit 摆脱“强制性匹配”（见图 4b 第一行），而其标准化后的 logit 结果则与教师的几乎完美匹配（见图 4c 第一行）。
从图 4 第二行中可以从全局角度得出类似的结论。
▲ 图4. 第一行：一组学生、教师 logit 的柱状图可视化案例；第二行：学生、教师 logit 之间的平均差值的热度图
特征可视化：
4.4 教师-学生鸿沟问题
4.5 蒸馏ViT
参考文献
更多阅读
#投 稿 通 道#
让你的文字被更多人看到
如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。
总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。
PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个，让知识真正流动起来。
📝 稿件基本要求：
• 文章确系个人原创作品，未曾在公开渠道发表，如为其他平台已发表或待发表的文章，请明确标注
• 稿件建议以 markdown 格式撰写，文中配图以附件形式发送，要求图片清晰，无版权问题
• PaperWeekly 尊重原作者署名权，并将为每篇被采纳的原创首发稿件，提供业内具有竞争力稿酬，具体依据文章阅读量和文章质量阶梯制结算
📬 投稿通道：
• 投稿邮箱：[email protected]
• 来稿请备注即时联系方式（微信），以便我们在稿件选用的第一时间联系作者
• 您也可以直接添加小编微信（pwbot02）快速投稿，备注：姓名-投稿
△长按添加PaperWeekly小编
🔍
现在，在「知乎」也能找到我们了
进入知乎首页搜索「PaperWeekly」
点击「关注」订阅我们的专栏吧
微信扫码关注该文公众号作者