On the Pros and Cons of Momentum Encoder in Self-Supervised Visual Representation Learning

TL;DR

This paper investigates the role of momentum (EMA) in self-supervised learning, revealing that applying momentum to the encoder's final layers or projector can improve stability and performance without increasing computation.

Contribution

It is the first to analyze how EMA affects different parts of the encoder and proposes a projector-only momentum approach for efficiency and effectiveness.

Findings

EMA's benefit is mainly due to stability effects.

Applying EMA to the projector yields comparable or better results.

Projector-only momentum reduces computation while maintaining performance.

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.