The Hidden Pitfalls of the Cosine Similarity Loss

TL;DR

This paper reveals fundamental issues with cosine similarity loss in deep learning, showing it can cause points to grow in magnitude and hinder training, and proposes a simple initialization fix to improve convergence.

Contribution

The paper uncovers the gradient vanishing problems of cosine similarity loss in SSL and introduces cut-initialization to enhance training stability and speed.

Findings

Gradient of cosine similarity vanishes in certain settings.

Optimizing cosine similarity can cause points to grow in magnitude.

Cut-initialization improves convergence speed in SSL methods.

Abstract

We show that the gradient of the cosine similarity between two points goes to zero in two under-explored settings: (1) if a point has large magnitude or (2) if the points are on opposite ends of the latent space. Counterintuitively, we prove that optimizing the cosine similarity between points forces them to grow in magnitude. Thus, (1) is unavoidable in practice. We then observe that these derivations are extremely general -- they hold across deep learning architectures and for many of the standard self-supervised learning (SSL) loss functions. This leads us to propose cut-initialization: a simple change to network initialization that helps all studied SSL methods converge faster.