[论文解读] Investigating the Benefits of Projection Head for Representation Learning

Yihao Xue, Eric Gan+3 more|arXiv (Cornell University)|Mar 18, 2024

Educational Games and GamificationPsychology被引用 3

一句话总结

本文对对比学习中投影头为何能提升表示学习提供了理论分析，表明其可缓解深层特征专业化问题，并增强分布偏移下的鲁棒性。研究揭示，预投影表示因隐式层间特征加权而更具泛化性，并证明固定重加权头可达到与投影头相当的性能，提供了一种更具可解释性的替代方案。

ABSTRACT

An effective technique for obtaining high-quality representations is adding a projection head on top of the encoder during training, then discarding it and using the pre-projection representations. Despite its proven practical effectiveness, the reason behind the success of this technique is poorly understood. The pre-projection representations are not directly optimized by the loss function, raising the question: what makes them better? In this work, we provide a rigorous theoretical answer to this question. We start by examining linear models trained with self-supervised contrastive loss. We reveal that the implicit bias of training algorithms leads to layer-wise progressive feature weighting, where features become increasingly unequal as we go deeper into the layers. Consequently, lower layers tend to have more normalized and less specialized representations. We theoretically characterize scenarios where such representations are more beneficial, highlighting the intricate interplay between data augmentation and input features. Additionally, we demonstrate that introducing non-linearity into the network allows lower layers to learn features that are completely absent in higher layers. Finally, we show how this mechanism improves the robustness in supervised contrastive learning and supervised learning. We empirically validate our results through various experiments on CIFAR-10/100, UrbanCars and shifted versions of ImageNet. We also introduce a potential alternative to projection head, which offers a more interpretable and controllable design.

研究动机与目标

• 理解对比自监督学习中投影头经验成功背后的理论机制。
• 研究为何未经直接优化的预投影表示能实现更好的下游性能。
• 分析数据增强与网络非线性对特征分布与泛化性的影响。
• 将发现扩展至监督对比学习与监督学习，证明其提升了鲁棒性并减少了类别坍塌。
• 提出并验证一种标准投影头的更可解释替代方案：固定重加权头。

提出的方法

• 对使用对比损失训练的线性模型进行理论分析，揭示因优化偏差导致的层间渐进特征加权。
• 推导预投影表示更具优势的条件，尤其在数据增强扭曲有用特征时。
• 将分析扩展至非线性网络，表明低层可学习到高层或投影头输出中缺失的特征。
• 在CIFAR-10/100、UrbanCars及分布偏移的ImageNet变体上进行实证验证，比较预投影与后投影表示。
• 引入并评估一种固定重加权头作为可学习投影头的替代方案，性能相当。
• 使用半合成数据集（MNIST在CIFAR-10上）及在ImageNet上的微调与OOD基准测试鲁棒性。

Figure 1: Weights of features in a two-layer fully connected ReLU network trained with CL. Left : With all features having equal strength, those that are more disrupted by augmentation have smaller/zero weights. Right : With augmentation treating all features equally. features with the largest/small

实验结果

研究问题

• RQ1 为何预投影表示在未被直接优化的情况下仍比后投影表示泛化性更好？
• RQ2 优化的隐式偏差如何导致对比学习中层间渐进特征加权？
• RQ3 在数据增强损害有用特征，或特征过强/过弱时，投影头在何种情况下能提升鲁棒性？
• RQ4 低层是否能学习到投影头输出中不存在的特征，这对泛化性有何影响？
• RQ5 固定重加权头能否实现与可学习投影头相当的性能，是否可作为更具可解释性的替代方案？

主要发现

• 当数据增强扭曲有用特征，或预训练期间特征过强/过弱时，预投影表示比后投影表示更具鲁棒性且泛化性更优。
• 优化的隐式偏差导致层间渐进特征加权，使深层更专注于预训练目标，从而使低层更标准化且专业化程度更低。
• 非线性激活放大特征专业化，使低层能学习到投影头输出中完全缺失的特征，从而提升泛化性。
• 在分布偏移的ImageNet基准上，预投影表示在分布外准确率上优于后投影表示与原始模型，展现出对背景偏差更强的鲁棒性。
• 固定重加权头性能与标准投影头相当，验证了特征重加权是核心机制的理论洞察，并提示了一种更具可解释性的设计替代方案。
• 在监督对比学习与监督学习中，预投影表示减少类别坍塌并提升对分布偏移的少样本适应能力，尤其在UrbanCars与ImageNet上表现显著。

Figure 2: Left: Weights of features in a two-layer fully connected ReLU network trained using SCL. The subclass feature is not represented post-projection but is represented pre-projection. Right: As a result, the four subclasses are only separable in pre-projection representations.