Reproduce “The Edge of Orthogonality: A Simple View of What Makes BYOL Tick” on OpenResearch

About this paper

Pierre H. Richemond, Allison Tam, Yunhao Tang, Florian Strub, Bilal Piot, Felix Hill

Researchers at DeepMind clarified the collapse-prevention mechanisms in BYOL, a self-supervised learning method, by demonstrating through linear algebra that its optimal predictor acts as an orthogonal projection. The work introduces new closed-form predictor variants that surpass standard BYOL performance, revealing how implicit orthonormalization contributes to stable representation learning.

Key insights

• The optimal linear predictor in BYOL acts as an orthogonal projection, with the stop-gradient operation being crucial to prevent collapse by ensuring online representations are not orthogonal to themselves.
• BYOL's interleaved gradient updates and Exponential Moving Average (EMA) steps implicitly enforce approximate orthonormality, which can be interpreted as a trivial retraction in Riemannian Stochastic Gradient Descent.
• Riemannian gradients on the Stiefel manifold inherently contain a term that functions similarly to explicit negative examples in contrastive losses, contributing to the non-collapse property.

Summary by alphaXiv.