Representation Learning Dynamics of Self-Supervised Models

TL;DR

This paper investigates the learning dynamics of self-supervised models, revealing issues like dimension collapse and proposing orthogonality constraints, with theoretical analysis supported by numerical experiments.

Contribution

It introduces a novel theoretical framework for SSL learning dynamics, including orthogonality constraints and Grassmannian manifold analysis, extending beyond existing generalisation bounds.

Findings

Dimension collapse occurs in naive SSL dynamics.

Orthogonality constraints prevent trivial solutions.

Theoretical dynamics deviate from neural tangent kernel approximations.

Abstract

Self-Supervised Learning (SSL) is an important paradigm for learning representations from unlabelled data, and SSL with neural networks has been highly successful in practice. However current theoretical analysis of SSL is mostly restricted to generalisation error bounds. In contrast, learning dynamics often provide a precise characterisation of the behaviour of neural networks based models but, so far, are mainly known in supervised settings. In this paper, we study the learning dynamics of SSL models, specifically representations obtained by minimising contrastive and non-contrastive losses. We show that a naive extension of the dymanics of multivariate regression to SSL leads to learning trivial scalar representations that demonstrates dimension collapse in SSL. Consequently, we formulate SSL objectives with orthogonality constraints on the weights, and derive the exact (network width independent) learning dynamics of the SSL models trained using gradient descent on the Grassmannian manifold. We also argue that the infinite width approximation of SSL models significantly deviate from the neural tangent kernel approximations of supervised models. We numerically illustrate the validity of our theoretical findings, and discuss how the presented results provide a framework for further theoretical analysis of contrastive and non-contrastive SSL.