Representation Learning Dynamics of Self-Supervised Models
Pascal Esser Satyaki Mukherjee Debarghya Ghoshdastidar
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.