Manifold Contrastive Learning with Variational Lie Group Operators

Published: 29 Feb 2024, Last Modified: 29 Feb 2024Accepted by TMLREveryoneRevisionsBibTeX
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
Supplementary Material: zip
Assigned Action Editor: ~Jake_Snell1
License: Creative Commons Attribution 4.0 International (CC BY 4.0)
Submission Number: 1444