[论文解读] DisCo: Remedy Self-supervised Learning on Lightweight Models with Distilled Contrastive Learning

Yuting Gao, Jiaxin Zhuang+5 more|arXiv (Cornell University)|Apr 19, 2021

Domain Adaptation and Few-Shot Learning参考文献 46被引用 8

一句话总结

本文提出 DisCo，一种知识蒸馏框架，通过使用均方误差（MSE）损失将大型教师模型的最终对比嵌入知识迁移至小型学生模型，从而在轻量级模型上增强自监督学习（SSL）性能。该方法显著提升了线性评估准确率——当使用 ResNet-101 作为教师模型时，EfficientNet-B0 在 ImageNet 上的 Top-1 准确率提升了 22.1%，且推理阶段不增加参数。

ABSTRACT

While self-supervised representation learning (SSL) has received widespread attention from the community, recent research argue that its performance will suffer a cliff fall when the model size decreases. The current method mainly relies on contrastive learning to train the network and in this work, we propose a simple yet effective Distilled Contrastive Learning (DisCo) to ease the issue by a large margin. Specifically, we find the final embedding obtained by the mainstream SSL methods contains the most fruitful information, and propose to distill the final embedding to maximally transmit a teacher's knowledge to a lightweight model by constraining the last embedding of the student to be consistent with that of the teacher. In addition, in the experiment, we find that there exists a phenomenon termed Distilling BottleNeck and present to enlarge the embedding dimension to alleviate this problem. Our method does not introduce any extra parameter to lightweight models during deployment. Experimental results demonstrate that our method achieves the state-of-the-art on all lightweight models. Particularly, when ResNet-101/ResNet-50 is used as teacher to teach EfficientNet-B0, the linear result of EfficientNet-B0 on ImageNet is very close to ResNet-101/ResNet-50, but the number of parameters of EfficientNet-B0 is only 9.4\%/16.3\% of ResNet-101/ResNet-50. Code is available at https://github. com/Yuting-Gao/DisCo-pytorch.

研究动机与目标

• 解决轻量级模型中自监督学习（SSL）性能下降的问题，即模型规模减小后准确率急剧下降。
• 克服现有蒸馏方法的局限性，这些方法通常传递次优知识（如 logits 或特征图），而非最具信息量的最终嵌入。
• 通过聚焦最终对比嵌入作为主要知识源，提升轻量级模型在 SSL 中的泛化能力。
• 识别并解决‘蒸馏瓶颈’——即学生模型 MLP 头隐藏维度受限导致的性能上限问题——通过扩展嵌入维度来缓解。
• 设计一种在下游部署阶段不引入额外参数的方法，确保在实际应用中的高效性。

提出的方法

• 通过使用均方误差（MSE）损失约束学生模型的最终对比嵌入与教师模型的最终对比嵌入一致，实现从大型教师模型向轻量级学生模型的知识迁移。
• 将教师模型最终投影头输出（MLP 之后）作为主要知识信号，因其包含对比学习中最具信息量的表示。
• 引入一种改进的训练流程，在预训练阶段对学生的最终嵌入进行正则化，使其与教师的最终嵌入对齐。
• 通过增加学生模型 MLP 头的隐藏维度来解决‘蒸馏瓶颈’，从而提升知识传输能力。
• 确保预训练阶段使用的额外 MLP 头在微调和推理前被移除，以保持模型效率。
• 将 DisCo 集成到多种 SSL 框架（如 MoCo-V2、SwAV、DINO）中，以验证其在不同架构和学习范式下的兼容性与泛化能力。

实验结果

研究问题

• RQ1 将大型教师模型的最终对比嵌入迁移至轻量级学生模型，是否能显著提升其在自监督学习中的表示质量？
• RQ2 通过 MSE 损失约束学生最终嵌入与教师嵌入一致，是否能带来优于现有蒸馏方法的下游性能？
• RQ3 轻量级模型的性能下降是否由知识迁移瓶颈引起？扩大学生模型 MLP 头维度能否缓解此问题？
• RQ4 DisCo 是否可有效应用于多种 SSL 框架（如 MoCo-V2、SwAV、DINO）而无需修改架构？
• RQ5 将 DisCo 与其他蒸馏技术（如 AT、RKD、KD）结合，是否能进一步提升性能？

主要发现

• 当使用 ResNet-101 作为教师模型时，DisCo 将 EfficientNet-B0 在 ImageNet 上的线性评估 Top-1 准确率提升 22.1%，达到 66.5% 的准确率。
• 当使用 ResNet-50 作为教师模型时，DisCo 将 EfficientNet-B0 的 Top-1 准确率提升 19.7%，达到 66.6% 的准确率。
• 该方法在所有测试的轻量级模型（包括 MobileNet-v3 和 ResNet-18）上均优于 SOTA 方法（如 SEED 和 MoCo-V2）。
• 学生模型 MLP 头维度的扩展使性能单调提升，最高增益出现在 2048 维时。
• 将 DisCo 与其他蒸馏方法（AT、RKD、KD）结合可进一步提升性能，证实其兼容性与有效性。
• DisCo 在多种 SSL 框架（如 MoCo-V2、SwAV、DINO）中均表现优异，展现出强大的泛化能力与兼容性。