Rethinking Uniformity in Self-Supervised Representation Learning

Xianghong Fang, Jian Li, Xiangchu Feng, Benyou Wang

Keywords: Collapse Analysis, Wasserstein Distance, Self-Supervised Representation Learning

Abstract: Self-supervised representation learning has achieved great success in many machine learning tasks. While many research efforts focus on learning better representations by preventing the model from the \emph{collapse} problem, less attention has been drawn to analyzing the collapse degrees of representations. In this paper, we present a formal study of collapse analysis via the \emph{uniformity} metric, which measures how uniformly learned representations distribute on the surface of the unit hypersphere. We fundamentally find that \textit{representation that obeys zero-mean isotropic Gaussian distribution is with the ideal uniformity} since its $l_2$-normalized form uniformly distributes on the surface of the unit hypersphere. Therefore, we propose to use the Wasserstein distance between the distribution of learned representations and the ideal distribution as a quantifiable metric of \emph{uniformity}. Moreover, we design five desirable constraints for ideal uniformity metrics, based on which we find that the proposed uniformity metric satisfies all constraints while the existing one does not. Synthetic experiments also demonstrate the proposed uniformity metric is capable to deal with the dimensional collapse while the existing one is insensitive. Furthermore, we impose the proposed \emph{uniformity} metric as an auxiliary loss term for various existing self-supervised methods, which consistently improves the downstream performance.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics

Submission Guidelines: Yes

Please Choose The Closest Area That Your Submission Falls Into: Unsupervised and Self-supervised learning