[论文解读] The Mechanism of Prediction Head in Non-contrastive Self-supervised Learning

Zixin Wen, Yuanzhi Li|arXiv (Cornell University)|May 12, 2022

Stochastic Gradient Optimization Techniques被引用 6

一句话总结

本文揭示了可训练预测头如何在非对比自监督学习中防止维数坍塌的机制。通过实证与理论分析，识别出两个关键效应——替代与加速：预测头通过允许更强特征替代较弱特征并加速其学习，使神经网络能够学习多样化特征，从而即使在损失函数景观中存在平凡坍塌解的情况下，也能避免陷入此类解。

ABSTRACT

Recently the surprising discovery of the Bootstrap Your Own Latent (BYOL) method by Grill et al. shows the negative term in contrastive loss can be removed if we add the so-called prediction head to the network. This initiated the research of non-contrastive self-supervised learning. It is mysterious why even when there exist trivial collapsed global optimal solutions, neural networks trained by (stochastic) gradient descent can still learn competitive representations. This phenomenon is a typical example of implicit bias in deep learning and remains little understood. In this work, we present our empirical and theoretical discoveries on non-contrastive self-supervised learning. Empirically, we find that when the prediction head is initialized as an identity matrix with only its off-diagonal entries being trainable, the network can learn competitive representations even though the trivial optima still exist in the training objective. Theoretically, we present a framework to understand the behavior of the trainable, but identity-initialized prediction head. Under a simple setting, we characterized the substitution effect and acceleration effect of the prediction head. The substitution effect happens when learning the stronger features in some neurons can substitute for learning these features in other neurons through updating the prediction head. And the acceleration effect happens when the substituted features can accelerate the learning of other weaker features to prevent them from being ignored. These two effects enable the neural networks to learn all the features rather than focus only on learning the stronger features, which is likely the cause of the dimensional collapse phenomenon. To the best of our knowledge, this is also the first end-to-end optimization guarantee for non-contrastive methods using nonlinear neural networks with a trainable prediction head and normalization.

研究动机与目标

• 理解为何使用随机梯度下降训练的神经网络在非对比自监督学习中能避免平凡坍塌解，尽管此类解在全局上是最优的。
• 研究预测头在打破对称性、促进学习多样化且不坍塌的表征中的作用。
• 提供一个理论框架，解释预测头在非对比对比学习式训练中引入的隐式偏差。
• 表征在身份初始化预测头存在下，特征学习的动力学过程，特别是非对角线项的上升与下降。

提出的方法

• 实证研究仅更新非对角线项的单位矩阵初始化预测头的训练动态。
• 将训练过程理论建模为四个阶段：学习更强特征、替代效应、加速效应与收敛。
• 定义并分析替代效应：一个神经元中的学习使预测头能够替代另一个神经元的学习。
• 定义并分析加速效应：被替代的特征加速了弱特征的学习，防止其在优化过程中被忽略。
• 使用简化线性模型，推导具有可训练预测头和批量归一化非线性神经网络的优化保证。
• 利用张量幂方法界分析不同学习率和初始化条件下梯度更新与收敛性。

实验结果

研究问题

• RQ1 为何使用SGD训练的神经网络在非对比自监督学习中能避免维数坍塌，即使坍塌解是全局最优的？
• RQ2 当预测头可能学习恒等映射而使其失效时，它如何仍能促进多样化特征的学习？
• RQ3 预测头中的非对角线项在训练过程中起什么作用？为何它们遵循一致的上升与下降轨迹？
• RQ4 预测头的替代与加速效应如何协同作用，防止特征遗忘并实现完整表征学习？
• RQ5 我们能否为具有可训练预测头和归一化机制的非对比自监督学习提供端到端的优化保证？

主要发现

• 当预测头以单位矩阵初始化且仅训练非对角线项时，网络仍能学习到具有竞争力的表征，表明头的结构——而不仅仅是其参数——具有关键作用。
• 训练过程中观察到预测头中非对角线项呈现一致的上升与下降轨迹，表明存在动态的特征替代过程。
• 替代效应指：一个神经元中的强特征可通过更新预测头，替代另一个神经元的学习。
• 加速效应指：被替代的特征可加速弱特征的学习，防止其在优化过程中被忽略。
• 替代与加速效应共同作用，使网络能够学习所有特征而非仅关注最强特征，从而解释了维数坍塌的避免。
• 本文首次为具有可训练预测头和归一化机制的非线性神经网络在非对比自监督学习中提供了端到端优化保证。