[论文解读] Understanding and Improving the Role of Projection Head in Self-Supervised Learning

Kartik Gupta, Thalaiyasingam Ajanthan+2 more|arXiv (Cornell University)|Dec 22, 2022

Domain Adaptation and Few-Shot Learning被引用 12

一句话总结

本文研究了为何在对比自监督学习（SSL）中，尽管训练后会丢弃可学习的投影头，它仍然是必不可少的。本文将SSL重新构建成一个双层优化问题，将投影头视为损失函数的一部分而非网络的一部分，从而在CIFAR-10、CIFAR-100和TinyImageNet上相比使用可学习投影头的标准SimCLR方法，实现了更好的泛化性能。

ABSTRACT

Self-supervised learning (SSL) aims to produce useful feature representations without access to any human-labeled data annotations. Due to the success of recent SSL methods based on contrastive learning, such as SimCLR, this problem has gained popularity. Most current contrastive learning approaches append a parametrized projection head to the end of some backbone network to optimize the InfoNCE objective and then discard the learned projection head after training. This raises a fundamental question: Why is a learnable projection head required if we are to discard it after training? In this work, we first perform a systematic study on the behavior of SSL training focusing on the role of the projection head layers. By formulating the projection head as a parametric component for the InfoNCE objective rather than a part of the network, we present an alternative optimization scheme for training contrastive learning based SSL frameworks. Our experimental study on multiple image classification datasets demonstrates the effectiveness of the proposed approach over alternatives in the SSL literature.

研究动机与目标

• 理解为何在训练后被丢弃的可学习投影头在对比SSL中仍是必要的。
• 研究投影头在对比损失优化过程中对子空间选择的隐式作用。
• 弥补对投影头功能贡献的现有理解空白，超越其作为可学习头的简单角色。
• 提出一种新的优化方案，将投影头视为损失函数的一部分，以提升训练稳定性和泛化能力。
• 通过实证验证，学习投影头相比固定或丢弃的投影头能显著提升性能。

提出的方法

• 将SSL建模为双层优化问题：内层循环通过优化投影头来选择对比损失的最优子空间，外层循环则更新主干网络的特征。
• 将投影头重新构建成损失函数的可学习组成部分，而非固定的网络头，从而实现动态子空间选择。
• 采用交替优化策略：通过梯度下降在对比损失上更新投影头，随后使用更新后的投影头来更新主干网络。
• 将该方法应用于SimCLR和SimSiam，在CIFAR-10、STL-10、TinyImageNet和ImageNet上分析其行为与性能。
• 将所提方法与标准SimCLR和DirectCLR进行比较，评估其在线性评估协议下的泛化性能。
• 分析投影头的秩和零空间，以理解其在特征选择中的隐式作用。

实验结果

研究问题

• RQ1 如果训练后投影头会被丢弃，为何在SSL中仍需要可学习的投影头？
• RQ2 除了作为最终投影层外，投影头在训练过程中还发挥何种功能角色？
• RQ3 是否可以将投影头的行为——特别是其隐式选择特征子空间的能力——形式化并作为损失函数的一部分进行强制执行？
• RQ4 将投影头视为损失函数的一部分，是否能带来比标准可学习头训练更好的泛化性能？
• RQ5 所提方法在多个数据集上的性能与固定投影头或标准SimCLR相比如何？

主要发现

• 投影头充当低秩映射，隐式选择特征子空间以应用对比损失，从而提升训练稳定性和泛化能力。
• 投影头的零空间对泛化有帮助，表明该头学习到了抑制无关特征分量的能力。
• 在标准SSL训练中，可学习投影头无法被固定投影头替代，否则性能会下降。
• 所提出的双层优化方案在CIFAR-10、CIFAR-100和TinyImageNet上优于标准SimCLR，证明了其泛化能力的提升。
• 尽管固定投影头在训练初期表现更优，但可学习投影头最终会超越它，验证了端到端优化的必要性。
• 该方法在线性评估基准上达到了最先进性能，证实投影头的作用并非偶然，而是学习过程中的关键组成部分。