Invariant Representations via Wasserstein Correlation Maximization

2025-05-16

AI Summary

This paper addresses the challenge of jointly preserving geometric fidelity and enhancing robustness in unsupervised representation learning. We propose a novel framework based on Wasserstein correlation: by maximizing the Wasserstein correlation between the input data distribution and the encoded distribution, our autoencoder preserves the topological and geometric structure of the original data during dimensionality reduction. To our knowledge, this is the first work to leverage Wasserstein correlation for learning representations that are approximately invariant to specified augmentations (or augmentation sets). We introduce a Markov–Wasserstein kernel to construct an augmentation encoder, enabling efficient acquisition of invariance without fine-tuning—i.e., with the pretrained backbone frozen. We provide theoretical guarantees establishing the optimality and stability of the proposed metric. Empirical results demonstrate that even simple feedforward networks trained under this framework yield highly invariant and geometrically faithful representations.

Technology Category

• Machine Learning: Representation Learning
• Computer Vision: Representation Learning for Vision
• Search and Optimization: Learning to Search

Application Category

• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs
• Search and Retrieval-Augmented AI: Web query analysis, representation and understanding
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

This work investigates the use of Wasserstein correlation -- a normalized measure of statistical dependence based on the Wasserstein distance between a joint distribution and the product of its marginals -- for unsupervised representation learning. Unlike, for example, contrastive methods, which naturally cluster classes in the latent space, we find that an (auto)encoder trained to maximize Wasserstein correlation between the input and encoded distributions instead acts as a compressor, reducing dimensionality while approximately preserving the topological and geometric properties of the input distribution. More strikingly, we show that Wasserstein correlation maximization can be used to arrive at an (auto)encoder -- either trained from scratch, or else one that extends a frozen, pretrained model -- that is approximately invariant to a chosen augmentation, or collection of augmentations, and that still approximately preserves the structural properties of the non-augmented input distribution. To do this, we first define the notion of an augmented encoder using the machinery of Markov-Wasserstein kernels. When the maximization objective is then applied to the augmented encoder, as opposed to the underlying, deterministic encoder, the resulting model exhibits the desired invariance properties. Finally, besides our experimental results, which show that even simple feedforward networks can be imbued with invariants or can, alternatively, be used to impart invariants to pretrained models under this training process, we additionally establish various theoretical results for optimal transport-based dependence measures. Code is available at https://github.com/keenan-eikenberry/wasserstein_correlation_maximization .

Problem

• Maximizing Wasserstein correlation for unsupervised representation learning
• Creating invariant encoders preserving input distribution properties
• Extending pretrained models with augmentation-invariant features

Innovation

• Uses Wasserstein correlation for unsupervised representation learning
• Maximizes correlation to preserve input distribution properties
• Defines augmented encoder with Markov-Wasserstein kernels