Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Invariant Representations via Wasserstein Correlation Maximization

Keenan Eikenberry, Lizuo Liu, Yoonsang Lee

TL;DR

This work proposes Wasserstein correlation maximization as a geometry-aware alternative to mutual information for unsupervised representation learning. By optimizing a baseline objective with WC_p and extending encoders to probabilistic augmented encoders via Markov-Wasserstein kernels, the authors obtain latent representations that preserve input topology and geometry while achieving approximate invariance to chosen augmentations. The framework includes a rigorous theory connecting adapted Wasserstein dependence to barycenters and Markov kernels, and practical, scalable computations via sliced Wasserstein distances. Empirical results on MNIST, CIFAR10, and STL10 with simple networks demonstrate both structure preservation and augmentation invariance, and they show that invariance can be imparted to pretrained models with only a few additional layers. The work provides a principled OT-based approach to invariant, geometry-preserving representations with publicly available code for replication.

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 .

Invariant Representations via Wasserstein Correlation Maximization

TL;DR

This work proposes Wasserstein correlation maximization as a geometry-aware alternative to mutual information for unsupervised representation learning. By optimizing a baseline objective with WC_p and extending encoders to probabilistic augmented encoders via Markov-Wasserstein kernels, the authors obtain latent representations that preserve input topology and geometry while achieving approximate invariance to chosen augmentations. The framework includes a rigorous theory connecting adapted Wasserstein dependence to barycenters and Markov kernels, and practical, scalable computations via sliced Wasserstein distances. Empirical results on MNIST, CIFAR10, and STL10 with simple networks demonstrate both structure preservation and augmentation invariance, and they show that invariance can be imparted to pretrained models with only a few additional layers. The work provides a principled OT-based approach to invariant, geometry-preserving representations with publicly available code for replication.

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 .
Paper Structure (24 sections, 9 theorems, 63 equations, 4 figures, 4 tables)

This paper contains 24 sections, 9 theorems, 63 equations, 4 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Wasserstein Correlation Maximization
  4. Baseline Objective
  5. Augmented Encoders
  6. General Objective
  7. Computational Objective
  8. Theory
  9. Experiments
  10. Structure Preservation
  11. Invariance Evaluation
  12. Conclusion
  13. Background
  14. Wasserstein Distance
  15. Sliced Wasserstein Distance
  16. ...and 9 more sections

Key Result

Theorem 2.1

Let $\gamma\in \Pi_p(\mu,\nu)$ with $\gamma =\int_X(\delta_x\otimes F_x)\mu(dx)$ for some $\mu$-almost surely unique $F:X\to\mathcal{P}_p(Y)$, and let $C_{\nu}:X\to\mathcal{P}_p(Y)$ be the constant Markov kernel defined by $C_{\nu}(x)=\nu$. Then, adapted Wasserstein dependence can be written in term

Figures (4)

  • Figure 1: Latent codes of rotation-augmented encoder for original MNIST data
  • Figure 2: Latent codes of rotation-augmented encoder for original MNIST data plus $90$ degree rotations
  • Figure 3: Latent codes of non-augmented encoder for original MNIST data plus $90$ degree rotations
  • Figure 4: Latent codes of noise-augmented encoder for original (i.e., non-noisy) STL10 DINOv2 features

Theorems & Definitions (32)

  • Theorem 2.1
  • Corollary 2.2
  • Proposition 2.3
  • Definition A.1
  • Definition A.2
  • Definition A.3
  • Definition A.4
  • Proposition A.5
  • proof
  • Definition A.6
  • ...and 22 more