Dual Perspectives on Non-Contrastive Self-Supervised Learning

TL;DR

This paper analyzes the stop gradient and exponential moving average procedures in non-contrastive self-supervised learning, revealing their role in avoiding collapse through a dual optimization and dynamical systems perspective, supported by empirical evidence.

Contribution

It provides a theoretical understanding of how these procedures prevent collapse, showing they do not optimize the original objective but are asymptotically stable in linear models.

Findings

Stop gradient and EMA avoid collapse without optimizing the original objective.

In linear models, minimizing the original objective leads to collapse.

Equilibria of the associated dynamical systems are explicitly characterized and are asymptotically stable.

Abstract

The {\em stop gradient} and {\em exponential moving average} iterative procedures are commonly used in non-contrastive approaches to self-supervised learning to avoid representation collapse, with excellent performance in downstream applications in practice. This presentation investigates these procedures from the dual viewpoints of optimization and dynamical systems. We show that, in general, although they {\em do not} optimize the original objective, or {\em any} other smooth function, they {\em do} avoid collapse Following~\citet{Tian21}, but without any of the extra assumptions used in their proofs, we then show using a dynamical system perspective that, in the linear case, minimizing the original objective function without the use of a stop gradient or exponential moving average {\em always} leads to collapse. Conversely, we characterize explicitly the equilibria of the dynamical systems associated with these two procedures in this linear setting as algebraic varieties in their parameter space, and show that they are, in general, {\em asymptotically stable}. Our theoretical findings are illustrated by empirical experiments with real and synthetic data.