PaperlayerView exact revision on arXiv

Source-linked AI summary

Wasserstein Dependency Measure for Representation Learning

Sherjil Ozair, Corey Lynch, Yoshua Bengio, Aaron van den Oord, Sergey Levine, Pierre Sermanet

arXiv:1903.11780v1cs.LGstat.ML

TL;DR

Mutual-information-based representation learning is limited by the sample requirements of tight lower bounds, especially when mutual information is high. The paper proposes the Wasserstein dependency measure and a practical Lipschitz-regularized approximation; it reports substantially reduced incomplete representations and improved capture of task-salient features. Its Lipschitz approximation may reduce neural-network capacity on complex tasks requiring high capacity.

  • Problem

    Tight mutual-information lower bounds require sample size exponential in the mutual information, limiting their applicability to high-mutual-information prediction tasks.

    Source: §1 · p. 2, §3.1 · p. 3

  • Method

    The paper replaces KL divergence in mutual-information estimation with the Wasserstein dependency measure and approximates it using Lipschitz regularization.

    Source: §1 · p. 2

  • Results

    The practical Wasserstein approximation substantially reduces incomplete representations on tasks where they are a major challenge.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The approximation improves the ability of mutual-information-based representations to capture task-salient features.

    Source: §1 · p. 2

    Lipschitz continuity can significantly reduce neural-network nonlinear capacity, potentially hurting performance on complex tasks where high capacity is essential.

    Source: §5.1 · p. 5

Abstract

from arXiv · show

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.

1. Introduction

Mutual-information-based self-supervised learning can produce incomplete representations, even when low-level shortcuts are addressed. The paper proposes replacing KL divergence with Wasserstein distance and reports that a practical Lipschitz-regularized approximation reduces this problem.

1. Introduction

  • Even after low-level shortcuts are addressed, self-supervised methods can capture only a few underlying factors of variation.Identifying and fixing shortcuts requires domain expertise and may still leave degenerate solutions.Source: §1 · p. 2
  • Mutual-information lower bounds require sample sizes exponential in the mutual information, a concern for high-mutual-information tasks such as predicting successive video frames.The introduction notes that methods in such settings may capture only a few objects while still predicting future frames confidently.Source: §1 · p. 2
  • The paper formally motivates this limitation and empirically demonstrates it on problems where reconstructive models learn complete representations but self-supervised methods struggle.The proposed solution replaces KL divergence with the Wasserstein metric as a training objective.Source: §1 · p. 2
  • A practical Wasserstein approximation substantially reduces incomplete representations and improves the capture of task-salient features.The approximation uses regularization methods designed for generative adversarial networks.Source: §1 · p. 2

2. Background

The background introduces representation learning, mutual information, and Wasserstein distance as the concepts used in the paper. It frames representation learning in terms of downstream utility and the two dependency measures in terms of distributions and their geometry.

2.1. Representation Learning

  • A representation encoder maps an input x to a lower-dimensional representation z intended to make a downstream task easier to learn.The passage identifies improved generalization, sample efficiency, and compute efficiency as possible benefits.Source: §2.1 · p. 2

2.2. Mutual Information

  • Mutual information measures dependency between x and y as the KL divergence between their joint distribution and the product of their marginals.It is zero for independent variables and infinite when the variables are identical.Source: §2.2 · p. 2

2.3. Wasserstein Distance

  • The Wasserstein distance is introduced for distributions on a compact metric space and also has a Kantorovich-Rubenstein dual form over 1-Lipschitz functions.The background identifies couplings of the distributions and defines the relevant function class by a 1-Lipschitz condition.Source: §2.3 · pp. 2–3

3. Mutual Information Estimation and Maximization

The section reviews mutual-information maximization for representation learning and its estimation challenges. It then states a formal sample-complexity limitation: a high-confidence lower bound requires exponentially many samples as mutual information grows.

3. Mutual Information Estimation and Maximization

  • Mutual information has been maximized between inputs and representations or between representations of different data modalities.Recent variational energy-based lower bounds are tractable in high dimensions and estimable with gradient-based optimization.Source: §3 · p. 3
  • Despite these estimation advances, the techniques’ applicability remains limited to low-mutual-information settings.The section links the estimation limitations to those of KL-divergence lower bounds.Source: §3 · p. 3, §3.1 · p. 3

3.1. Formal Limitations in Mutual Information Estimation

  • A high-confidence lower bound on mutual information requires sample size n = exp(I(x; y)).The section derives this requirement from the correspondence between mutual information and KL divergence.Source: §3.1 · p. 3

4. Wasserstein Dependency Measure

The paper replaces KL divergence in mutual-information estimation with a metric-aware Wasserstein dependency measure. It discusses the base metric assumption and argues that Lipschitz-constrained Wasserstein estimation has better generalization properties.

4. Wasserstein Dependency Measure

  • KL divergence can be large even when samples differ little, whereas Wasserstein distance reflects distances between samples in the underlying metric space.The paper proposes Wasserstein distance as a replacement for KL divergence in mutual information and names this measure the Wasserstein dependency measure.Source: §4 · pp. 3–4
  • The Wasserstein dependency measure compares the joint distribution p(x, y) with the product of marginals p(x)p(y).It measures the cost of transforming marginal samples into joint samples, in contrast to KL divergence’s discrimination of distributional differences.Source: §4 · p. 4

4.1. Choice of Metric Space

  • The method assumes a known base metric space; the experiments use a Euclidean metric, treated as a task prior distinct from the data-informed representation metric.The paper describes the base and representation-induced metrics as prior and posterior metrics, respectively.Source: §4.1 · p. 4

4.2. Generalization of Wasserstein Distances

  • Lipschitz neural-network approximation gives Wasserstein distance better generalization properties than the mutual-information lower-bound function class.The stated generalization gap scales with the square root of the network’s Lipschitz constant, bounded at 1 for the Wasserstein function class but unbounded for the mutual-information lower-bound class.Source: §4.2 · p. 4

5. Wasserstein Predictive Coding for Representation learning

The section proposes a predictive-coding objective combining a Lipschitz function class with CPC’s log P exp term, and motivates Lipschitzness as a way to discourage encoders from exploiting isolated differences. It uses gradient penalty to approximate Lipschitz continuity while noting potential capacity costs on complex tasks.

5. Wasserstein Predictive Coding for Representation learning

  • Kantorovich-Rubenstein duality makes Wasserstein distance easier to estimate by enabling gradient-based optimization over neural-network functions.Source: §5 · p. 4
  • The proposed objective is a lower bound on both CPC and the dual Wasserstein dependency measure, retaining the Lipschitz function class and CPC’s log P exp term.Source: §5 · p. 4
  • The log P exp term reduces variance in sample-based gradient estimates and was found to improve performance in practice.Source: §5 · p. 5
  • Constraining the encoder to be Lipschitz bounds learned-representation distances by distances between underlying samples, encouraging it to represent more data components.Source: §5 · p. 5

5.1. Approximating Lipschitz Continuity

  • Gradient penalty (Gulrajani et al., 2017) supports the experiments, but may be ineffective with ReLU networks on complex tasks, where Lipschitz constraints can reduce nonlinear capacity.Source: §5.1 · p. 5

6. Experiments

The experiments evaluate whether WPC improves representation learning over CPC on four paired-image tasks, using linear probes to predict underlying latent factors. WPC generally outperforms CPC, including at high mutual information and under varied data, minibatch, and network conditions.

6.1. Evaluation Methodology

  • Representation quality is measured by training linear classifiers to predict latent factors from the learned representations; experiments focus on datasets with large mutual information.Source: §6.1 · p. 5

6.2. Task Design

  • Four paired-image tasks use spatial or channel-stacked Omniglot, two-view Shapes3D sharing five latent factors, and CelebA image halves labeled by binary attributes.Source: §6.2 · pp. 5–6

6.3. Effect of Mutual Information

  • At three Omniglot characters, exp(I)=137,280 exceeds 50,000 training samples; CPC performance drops significantly, while WPC declines more gradually and remains better.Source: §6.3 · pp. 6–7

6.4. Effect of Dataset Size

  • At about 34.43 nats on the nine-character Omniglot tasks, neither method is near-perfect, though WPC significantly outperforms CPC; Table 2 also reports better WPC for both network classes.Source: §6.4 · p. 7, Table 2 · p. 7

6.5. Effect of Minibatch Size

  • WPC reaches optimal performance at minibatch size 32, and larger minibatches do not improve it; CPC performance increases as minibatches grow.Source: §6.5 · p. 7

6.6. Effect of Neural Network Inductive Bias

  • With CNNs, the WPC–CPC gap widens on StackedMultiOmniglot, whose structure mismatches CNN bias; spatial arrangement suits CNNs, while high mutual information still reduces both methods’ performance.Source: Figure 7 · p. 8, §6.6 · p. 8, Figure 6 · p. 8

6. Experiments

  • Across high-mutual-information, dataset-size, minibatch, and function-class settings, WPC outperforms CPC; its advantage is larger when CNN inductive bias mismatches the data.Source: §6.3 · p. 7, §6.4 · p. 7, §6.5 · p. 7

7. Conclusion

The conclusion reports that Wasserstein predictive coding mitigates limitations of prior mutual-information estimators and presents the Wasserstein dependency measure as an alternative objective. It also identifies optimization of Lipschitz-continuous networks as a remaining challenge.

7. Conclusion

  • WPC mitigates representation-learning problems from prior mutual-information estimators across several problem settings.The conclusion describes this mitigation as substantial, not complete.Source: §7 · p. 8
  • The Wasserstein dependency measure replaces KL divergence with Wasserstein distance as an alternative representation-learning objective.Source: §7 · p. 8
  • Wasserstein predictive coding implements the objective by regularizing existing mutual-information estimators to enforce Lipschitz continuity.Source: §7 · p. 8
  • Optimizing Lipschitz-continuous networks remains challenging, and the authors expect better regularization methods to improve Wasserstein dependency measure representations.Source: §7 · p. 8
Loading 1903.11780v1…