Atlas Based Representation and Metric Learning on Manifolds

Abstract: We explore the use of a topological manifold, represented as a collection of charts, as the target space of neural network based representation learning tasks. This is achieved by a simple adjustment to the output of an encoder's network architecture plus the addition of a maximal mean discrepancy (MMD) based loss function for regularization. Most algorithms in representation and metric learning are easily adaptable to our framework and we demonstrate its effectiveness by adjusting SimCLR (for representation learning) and standard triplet loss training (for metric learning) to have manifold encoding spaces. Our experiments show that we obtain a substantial performance boost over the baseline for low dimensional encodings. In the case of triplet training, we also find, independent of the manifold setup, that the MMD loss alone (i.e. keeping a flat, euclidean target space but using an MMD loss to regularize it) increases performance over the baseline in the typical, high-dimensional Euclidean target spaces. Code for reproducing experiments is provided at https://github.com/ekorman/neurve .

• The paper introduces a novel atlas-based representation that leverages a collection of charts and MMD regularization to systematically encode data on manifolds.
• The approach adapts frameworks like SimCLR and triplet loss by mapping data into low-dimensional manifold spaces, which enhances efficiency and visualization.
• Experiments on MNIST, FashionMNIST, CIFAR10, CUB-200, and Stanford Cars demonstrate significant performance improvements and reduced computational demands.

Atlas Based Representation and Metric Learning on Manifolds

This paper introduces a novel approach to representation and metric learning by utilizing topological manifolds through a collection of charts as target spaces instead of traditional Euclidean spaces. The process entails adjusting the output of an encoder's network while integrating a maximal mean discrepancy (MMD) loss for regularization. The approach can be seen as an extension of prior work which employed manifolds as latent spaces within autoencoders, offering promising benefits.

Theoretical Framework

Traditional representation learning approaches typically encode data into high-dimensional Euclidean spaces, which can be unnecessarily large and inefficient, particularly if the manifold hypothesis holds—that data is well represented by a manifold. According to the Whitney and Nash embedding theorems, embedding such manifolds may require dimensions beyond practicality unless their topology is utilized effectively.

This paper proposes replacing the Euclidean space with manifold encoding spaces for two learning algorithms: SimCLR for representation learning and standard triplet loss training for metric learning. The proposed method uses an encoder network outputting chart embeddings and membership probabilities. A scoring function determines which chart's encoding to use, aiming for an optimal manifold representation that yields significant performance improvements over baseline methods, especially in low dimensions.

Methodology

The authors generalize typical learning architectures to incorporate manifold encoding by forming an atlas—a collection of charts each with a local encoding mapping. A maximal mean discrepancy (MMD) loss is introduced to regularize encoding spaces, ensuring data distribution across the manifold is uniform and optimal.

The experimental framework validates this approach using modified versions of SimCLR and triplet loss, showing that mapping input data into manifold spaces offers improvement in low-dimensional representation scenarios. Moreover, even in high-dimensional spaces, MMD regularization contributes to enhanced performance without necessitating manifold encoding.

Experimental Results

Experiments conducted on MNIST, FashionMNIST, and CIFAR10 datasets demonstrated a substantial boost in low-dimensional manifolds over standard SimCLR baselines. Importantly, using 2D and 4D manifolds resulted in marked improvements, substantiating the practical advantages of manifold encoding such as reduced storage, improved visualization, and faster computations.

In metric learning tasks, representation benefits were consistent across benchmark datasets CUB-200 and Stanford Cars, further validating the theoretical underpinning of the proposal. Supplementing high-dimensional triplet loss with MMD regularization alone yielded notable performance gains, underscoring the regularization's independent efficacy.

Implications and Future Work

The findings suggest a shift towards manifold-based encoding can bring theoretical insights into practical algorithmic enhancements for low-dimensional scenarios. As a consequence, future research can explore broader applicability across diverse datasets and algorithm families beyond SimCLR and triplet training—such as MoCo, BYOL, and AMDIM.

Conclusion

The framework described advocates for a paradigm shift from Euclidean-centric representation learning to manifold-based approaches, offering manifold advantages highlighted by empirical improvements and theoretical interest. The exploration of alternative geometries and regularization techniques unlocks new avenues for advancing representation efficiency, visualization capabilities, and computation scalability. As the support for standardized networks expands, manifold learning tools can potentially refine manifold representation algorithms across varied AI applications.