Wasserstein Dependency Measure for Representation Learning

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.

References in corpus (11)

• Auto-Encoding Variational Bayes
• Estimating Mutual Information
• Large Scale GAN Training for High Fidelity Natural Image Synthesis
• Conditional Image Synthesis With Auxiliary Classifier GANs
• Variational Inference with Normalizing Flows
• Density estimation using Real NVP
• Adversarial Feature Learning
• Estimating divergence functionals and the likelihood ratio by convex risk minimization
• Adversarially Learned Inference
• Entropy and information in neural spike trains: Progress on the sampling problem
• Formal Limitations on the Measurement of Mutual Information

Cited by in corpus (7)

• Unsupervised State Representation Learning in Atari
• Fast Task Inference with Variational Intrinsic Successor Features
• A Note on Connecting Barlow Twins with Negative-Sample-Free Contrastive Learning
• Integrating Auxiliary Information in Self-supervised Learning
• Representation Learning for Sequence Data with Deep Autoencoding Predictive Components
• Moreau-Yosida $f$-divergences
• On Study of Mutual Information and its Estimation Methods