Manifold Contrastive Learning with Variational Lie Group Operators

Kion Fallah, Alec Helbling, Kyle A. Johnsen, Christopher John Rozell

Abstract: Self-supervised learning of deep neural networks has become a prevalent paradigm for learning representations that transfer to a variety of downstream tasks. Similar to proposed models of the ventral stream of biological vision, it is observed that these networks lead to a separation of category manifolds in the representations of the penultimate layer. Although this observation matches the manifold hypothesis of representation learning, current self-supervised approaches are limited in their ability to explicitly model this manifold. Indeed, current approaches often only apply a pre-specified set of augmentations for "positive pairs" during learning. In this work, we propose a contrastive learning approach that directly models the latent manifold using Lie group operators parameterized by coefficients with a sparsity-promoting prior. A variational distribution over these coefficients provides a generative model of the manifold, with samples which provide feature augmentations applicable both during contrastive training and downstream tasks. Additionally, learned coefficient distributions provide a quantification of which transformations are most likely at each point on the manifold while preserving identity. We demonstrate benefits in self-supervised benchmarks for image datasets, as well as a downstream semi-supervised task. In the former case, we demonstrate that the proposed methods can effectively apply manifold feature augmentations and improve learning both with and without a projection head. In the latter case, we demonstrate that feature augmentations sampled from learned Lie group operators can improve classification performance when using few labels.

Submission Length: Regular submission (no more than 12 pages of main content)

Changes Since Last Submission:

Camera ready changes:
* Format changed to camera ready
* Lie group background reworked into main text
* NNCLR results moved to main text
* Missing $\log$ added to ELBO equation in section 2.2
* Formatting of references updated
* Acknowledgements section added

Code: https://github.com/kfallah/manifold-contrastive

Assigned Action Editor: ~Jake_Snell1

License: Creative Commons Attribution 4.0 International (CC BY 4.0)

Submission Number: 1444