Wasserstein Dependency Measure for Representation Learning

TL;DR

This paper introduces the Wasserstein dependency measure, a new approach for representation learning that overcomes limitations of mutual information methods by using Wasserstein distance, leading to more complete representations especially in high mutual information tasks.

Contribution

The paper proposes a novel Wasserstein dependency measure for representation learning, addressing mutual information limitations and improving representation completeness in complex tasks.

Findings

Wasserstein dependency measure outperforms mutual information-based methods.

The approach yields more complete representations in high mutual information scenarios.

Empirical results show significant improvements on real-world tasks.

Abstract

Mutual information maximization has emerged as a powerful learning objective for unsupervised representation learning obtaining state-of-the-art performance in applications such as object recognition, speech recognition, and reinforcement learning. However, such approaches are fundamentally limited since a tight lower bound of mutual information requires sample size exponential in the mutual information. This limits the applicability of these approaches for prediction tasks with high mutual information, such as in video understanding or reinforcement learning. In these settings, such techniques are prone to overfit, both in theory and in practice, and capture only a few of the relevant factors of variation. This leads to incomplete representations that are not optimal for downstream tasks. In this work, we empirically demonstrate that mutual information-based representation learning approaches do fail to learn complete representations on a number of designed and real-world tasks. To mitigate these problems we introduce the Wasserstein dependency measure, which learns more complete representations by using the Wasserstein distance instead of the KL divergence in the mutual information estimator. We show that a practical approximation to this theoretically motivated solution, constructed using Lipschitz constraint techniques from the GAN literature, achieves substantially improved results on tasks where incomplete representations are a major challenge.