Wasserstein Dependency Measure for Representation Learning

By Sherjil Ozair, Corey Lynch, Yoshua Bengio, Aäron van den Oord, Sergey Levine, Pierre Sermanet. 2019

140 citations recorded by the source corpus.

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.

Subfields: Representation & Self-Supervised Learning

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Representation & Self-Supervised Learning

Citations in the corpus

On Mutual Information Maximization for Representation Learning (2019, methodology): “…issues with the current estimators, one may investigate other notions of divergence between p ( x, y ) and p ( x ) p ( y ) , for example using the Wasserstein distance instead (Ozair et al., 2019), which could enable a natural approach for enforcing smoothness in the resulting estimators.”
InfoMax-GAN: Improved Adversarial Image Generation via Information Maximization and Contrastive Learning (2020, methodology): “We conjecture this could be related to the Wasserstein Dependency Measure [50], a variant of mutual information which replaces the KL divergence term with Wasserstein distance, as measured using encoders from the class of 1-Lipschitz functions.”
Language Conditioned Imitation Learning Over Unstructured Data (2020, background): “The latter provides an upper bound on how all well the various methods can learn language conditioned control, independent of a difﬁcult perception problem (which might be improved upon independently with self-supervised representation learning methods (e.g. [41, 33, 34])).”
Understanding the Limitations of Variational Mutual Information Estimators (2019, background): “Therefore, it would be helpful to acknowledge these limitations and consider alternative measurements of information that are more suited for modern machine learning applications (Ozair et al., 2019; Tschannen et al., 2019).”
Decoupled Contrastive Learning (2021, background): “[22] proposes a Wasserstein distance to prevent the encoder from learning any other differences between unpaired samples.”
Joint Hand Motion and Interaction Hotspots Prediction from Egocentric Videos (2022, methodology): “We don’t use KLD (Kullback-Leibler Divergence) metric [10] as it is known to be sensitive to the tail of the distributions [4,63,91].”