[论文解读] On the Pros and Cons of Momentum Encoder in Self-Supervised Visual Representation Learning

Trung Pham, Chaoning Zhang+3 more|arXiv (Cornell University)|Aug 11, 2022

Digital Imaging for Blood Diseases被引用 5

一句话总结

本文研究了动量编码器在自监督视觉表征学习中的作用，并提出仅投影头使用动量（projector-only momentum）作为全网络动量的更高效替代方案。通过分析深层的梯度不稳定性，作者表明，仅对最终投影头应用动量（而非整个编码器）即可在保持性能增益的同时，通过避免主干网络的双重前向传播来降低计算成本。

ABSTRACT

Exponential Moving Average (EMA or momentum) is widely used in modern self-supervised learning (SSL) approaches, such as MoCo, for enhancing performance. We demonstrate that such momentum can also be plugged into momentum-free SSL frameworks, such as SimCLR, for a performance boost. Despite its wide use as a fundamental component in modern SSL frameworks, the benefit caused by momentum is not well understood. We find that its success can be at least partly attributed to the stability effect. In the first attempt, we analyze how EMA affects each part of the encoder and reveal that the portion near the encoder's input plays an insignificant role while the latter parts have much more influence. By monitoring the gradient of the overall loss with respect to the output of each block in the encoder, we observe that the final layers tend to fluctuate much more than other layers during backpropagation, i.e. less stability. Interestingly, we show that using EMA to the final part of the SSL encoder, i.e. projector, instead of the whole deep network encoder can give comparable or preferable performance. Our proposed projector-only momentum helps maintain the benefit of EMA but avoids the double forward computation.

研究动机与目标

• 理解动量编码器在自监督学习中的经验优势，超越现有动机。
• 探究动量带来的性能增益是否源于稳定性，特别是深层网络层中的稳定性。
• 通过仅将动量隔离到投影头而非整个编码器，降低动量带来的计算开销。
• 验证仅投影头动量在显著减少前向计算量的同时，仍能达到相当或更优的性能。

提出的方法

• 将动量编码器分解为骨干网络的各个模块，并分析各层之间的梯度波动。
• 监控整体损失相对于每个模块输出的梯度，以量化不稳定性。
• 提出一种仅投影头动量机制，其中仅最终的多层感知机（MLP）投影头使用指数移动平均（EMA），而骨干网络权重保持共享。
• 在最终层（第4层块和投影头）应用批量归一化（BN），以进一步稳定训练。
• 在CIFAR-10和CIFAR-100上比较全EMA（如BYOL）、仅投影头EMA与基线模型的线性评估准确率。
• 可视化训练过程中滤波器权重的变化，以展示有无EMA时的稳定性差异。

实验结果

研究问题

• RQ1 自监督学习中动量编码器带来性能增益的主要机制是什么？
• RQ2 动量在编码器网络不同阶段的影响如何变化？
• RQ3 能否在降低计算成本的同时保留动量带来的性能优势？
• RQ4 是否可以仅对投影头应用动量而不损失性能？
• RQ5 批量归一化与仅投影头动量如何协同作用以提升训练稳定性？

主要发现

• 动量带来的性能增益主要源于训练稳定性的提升，尤其是在编码器的最终层。
• 最终层（如第4层块和投影头）的梯度波动显著更高，表明其稳定性较低，而动量有助于缓解这一问题。
• 仅投影头动量在CIFAR-100上实现了88.45%的线性评估准确率，与基线BYOL性能相当，同时避免了对骨干网络的双重前向传播。
• 在第4层块输出和投影头处应用批量归一化可进一步提升仅投影头EMA的性能，达到88.45%的准确率——几乎与完整BYOL基线持平。
• 可视化结果表明，无动量时滤波器呈现不稳定的尖峰（如在第~10k步左右），而EMA则带来持续平滑的学习曲线。
• 对编码器早期阶段应用动量影响甚微，证实最终层是性能增益的主要来源。