The Edge of Orthogonality: A Simple View of What Makes BYOL Tick

TL;DR

This paper provides a simple mathematical explanation for why BYOL, a self-supervised learning method, works by highlighting the role of orthogonal projections and orthonormalization, and introduces new predictor variants that outperform standard BYOL.

Contribution

It offers a minimalistic mathematical framework based on orthogonality to explain BYOL's success and proposes four new predictor variants with improved performance.

Findings

Orthogonal projection approximates the optimal predictor in BYOL.

Stop-gradient acts as an orthonormalization mechanism.

New closed-form predictors outperform standard BYOL at 100 and 300 epochs.

Abstract

Self-predictive unsupervised learning methods such as BYOL or SimSiam have shown impressive results, and counter-intuitively, do not collapse to trivial representations. In this work, we aim at exploring the simplest possible mathematical arguments towards explaining the underlying mechanisms behind self-predictive unsupervised learning. We start with the observation that those methods crucially rely on the presence of a predictor network (and stop-gradient). With simple linear algebra, we show that when using a linear predictor, the optimal predictor is close to an orthogonal projection, and propose a general framework based on orthonormalization that enables to interpret and give intuition on why BYOL works. In addition, this framework demonstrates the crucial role of the exponential moving average and stop-gradient operator in BYOL as an efficient orthonormalization mechanism. We use these insights to propose four new \emph{closed-form predictor} variants of BYOL to support our analysis. Our closed-form predictors outperform standard linear trainable predictor BYOL at $100$ and $300$ epochs (top-$1$ linear accuracy on ImageNet).