(CVPR 2021)自监督学习新思路!基于蒸馏(distillation loss)的自监督学习算法
今天给大家介绍一篇CVPR 2021关于自监督的文章叫S2-BNN [1],来自CMU,HKUST和IIAI。它探讨了如果网络规模在非常小的情况下(比如efficient networks或者binary neural networks),什么样的自监督训练策略和方法是最好的。本文发现基于小网络的前提下,基于蒸馏(distillation learning)的自监督学习得到的模型性能远远强于对比学习(contrastive learning),同时他们还发现同时使用蒸馏和对比学习效果反而不如单独使用蒸馏损失,这也是一个非常有意思的发现。
论文:https://arxiv.org/abs/2102.08946
代码:https://github.com/szq0214/S2-BNN
本文的初衷是:如果网络模型的规模比较小,那么它的表达能力也会非常有限,从而会造成这类网络的预测概率变得不那么确信(confident),这时候作者就产生一个疑问:最常使用的对比学习在这种类型的网络上面还能取得较好的结果吗?
基于这个疑问:作者首先使用默认参数的MoCo V2作为自监督学习基准算法(baseline),在使用ReActNet作为主干网的时候在ImageNet上得到46.9%。由于该主干网是一个二值化网络,因此作者提出需要调整优化器、学习率策略、数据增强策略来匹配二值网络的特性。作者尝试使用Adam和SGD两种优化器来训练主干网,在线性评估(linear evaluation)阶段使用不同学习率来寻找最佳的参数设计。作者发现使用使用SGD训练的网络,在线性评估阶段学习率比较大的时候性能比较高,学习率降低时精度严重下降,而Adam训练的网络刚好相反,学习率降低时精度反而上升,同时Adam训练的网络在取得最佳精度的设置时结果明显优于SGD的最佳结果。因此作者首先提出了一个基于MoCo V2的更强的baseline+,性能为52.5%。
接下来作者提出并比较了如下三种策略:
1. 单纯使用对比学习(使用增强后的MoCo V2作为对比学习算法);
2. 对比学习损失加上知识蒸馏损失(文章使用cross-entropy loss作为蒸馏损失)训练模型;
3. 只使用知识蒸馏损失(cross-entropy loss)训练模型。
蒸馏损失函数表达:
作者这边使用的是soft版本的cross-entropy loss作为蒸馏损失函数,具体实现细节可以去看他们的代码,损失函数表达形式如下:
如何获取teacher呢?
这里作者讨论了两种方案来得到teacher网络,一种是在线(online)的方式同时训练teacher和student,如下图所示。另一种是离线(offline)的方式,即先使用自监督方法训好teacher,然后固定teacher的权重来蒸馏目标网络,这也是本文中采用的策略。为什么使用offline的策略呢?其主要好处是效率高,由于teacher只需要训练一次,后面可以重复使用,从而使student的训练更加高效。其次就是teacher在蒸馏过程中权重都是freeze的,产生的监督信号也跟精确更稳定,对于student的收敛也会有帮助。
下面是两种策略的算法示意图:
核心实验结果:
作者在ImageNets-1k上进行了实验,核心结果如下:
表中① ② ③分别对应了三种训练策略,本文增强之后的结果比原始的baseline高了5.6%,加上蒸馏损失(distillation loss)之后结果提升到5.6%,去掉对比损失(contrastive loss)精度进一步提升了5.5%,达到61.5%。可以看到只使用知识蒸馏的策略相比其他两个方案,对于性能的提升还是非常可观的。
为什么同时使用蒸馏损失和对比损失效果反而没有单独使用蒸馏损失效果好?
直观地说,蒸馏损失会迫使 student去模仿teacher网络的预测概率输出,而对比学习倾向于从数据本身中发现潜在模式(patterns)。如下图所示,在二值化网络场景中,对比损失学习细粒度表示的能力相比蒸馏损失相对较弱,学到的表达在语义层面也更模糊。因此,由于优化空间的差异而导致两者结合使用可能并不是最佳的解决方案。
同时期一些基于蒸馏的自监督学习方法:
最近基于知识蒸馏的自监督方法有不少,包括跟本文同时期的SEED [2] (发表于ICLR 2021, 两者投稿相隔一个月,可以认为是同时期的工作) 以及后续比较有名的FAIR的DINO [3] 等等, SEED基本上跟这篇文章的方法是类似的,只是在student的选取上一个选择的是小规模网络,本文选择的是efficient的二值化网络,但是本身训练方法上没有大的差别。DINO跟他们两者的主要区别在于DINO的teacher的权重在训练过程是不固定的。
更多分析和消融实验请阅读论文原文。
References:
[1] Zhiqiang Shen, Zechun Liu, Jie Qin, Lei Huang, Kwang-Ting Cheng, and Marios Savvides. "S2-bnn: Bridging the gap between self-supervised real and 1-bit neural networks via guided distribution calibration." In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 2165-2174. 2021.
[2] Zhiyuan Fang, Jianfeng Wang, Lijuan Wang, Lei Zhang, Yezhou Yang, and Zicheng Liu. "Seed: Self-supervised distillation for visual representation." arXiv preprint arXiv:2101.04731 (2021).
[3] Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. "Emerging properties in self-supervised vision transformers." arXiv preprint arXiv:2104.14294 (2021).