Source-linked AI summary
The Cramer Distance as a Solution to Biased Wasserstein Gradients
Marc G. Bellemare, Ivo Danihelka, Will Dabney, Shakir Mohamed, Balaji Lakshminarayanan, Stephan Hoyer, Rémi Munos
TL;DR
The paper asks how to retain Wasserstein’s geometric sensitivity while avoiding biased sample gradients that undermine stochastic optimization. It characterizes these properties, proposes the Cramér distance and a Cramér GAN, and reports better practical performance than Wasserstein-based alternatives.
Problem
Wasserstein respects outcome geometry but lacks unbiased sample gradients, creating concerns for sample-based stochastic optimization.
Method
The paper proposes the Cramér distance and extends it to a Cramér GAN using energy distance on transformed variables.
Results
The Cramér distance has all three desired properties, and the paper reports significantly better performance than the related Wasserstein GAN.
Takeaways & Limitations
Unbiased sample gradients allow geometric information to be incorporated into large-scale optimization problems.
Takeaways & Limitations
Whether an unbiased and practical optimization procedure exists for the Wasserstein loss remains an open question.
Abstract
from arXiv · showhide
The Wasserstein probability metric has received much attention from the machine learning community. Unlike the Kullback-Leibler divergence, which strictly measures change in probability, the Wasserstein metric reflects the underlying geometry between outcomes. The value of being sensitive to this geometry has been demonstrated, among others, in ordinal regression and generative modelling. In this paper we describe three natural properties of probability divergences that reflect requirements from machine learning: sum invariance, scale sensitivity, and unbiased sample gradients. The Wasserstein metric possesses the first two properties but, unlike the Kullback-Leibler divergence, does not possess the third. We provide empirical evidence suggesting that this is a serious issue in practice. Leveraging insights from probabilistic forecasting we propose an alternative to the Wasserstein metric, the Cramér distance. We show that the Cramér distance possesses all three desired properties, combining the best of the Wasserstein and Kullback-Leibler divergences. To illustrate the relevance of the Cramér distance in practice we design a new algorithm, the Cramér Generative Adversarial Network (GAN), and show that it performs significantly better than the related Wasserstein GAN.
1 Introduction
KL divergence ignores geometric closeness between outcomes, while Wasserstein incorporates that geometry but can produce biased sample gradients. The paper proposes Cramér distance as a geometrically sensitive alternative with unbiased sample gradients.
- KL divergence measures relative probability but cannot reward predictions that are geometrically close to the correct outcome.This limitation matters when outcome similarity, such as perceptual similarity in images, is important.
- Wasserstein incorporates outcome geometry and applies to distributions with non-overlapping supports, but sample-based estimation yields biased gradients.The paper identifies this bias as a possible reason for tentative deep-learning applications.
- The Cramér distance preserves geometric sensitivity while providing unbiased sample gradients.The paper evaluates the difference between the metrics in categorical estimation, regression, and image generation.
2 Probability Divergences and Metrics
The paper characterizes divergences using sum invariance, scale sensitivity, and unbiased sample gradients. KL has unbiased sample gradients but lacks scale sensitivity, whereas Wasserstein is ideal with respect to sum invariance and scale sensitivity but lacks unbiased gradients.
- Probability Divergences and Metrics: Wasserstein differs from KL by responding to the geometry of possible outcomes rather than only changes in probability.This geometric sensitivity motivates studying divergences beyond KL.
- Properties of a Divergence: Ideal divergences combine scale sensitivity with sum invariance, respectively reflecting outcome scaling and invariance under adding an independent variable.The paper illustrates these properties using shifted and rescaled Dirac distributions.
- Properties of a Divergence: A divergence is viewed as a loss d(P, Qθ) whose minimization seeks parameters θ* representing the target distribution.The paper defines unbiased sample gradients by comparing expected sample-loss gradients with true-loss gradients.
- Properties of a Divergence: Unbiased sample gradients matter because stochastic optimization may converge incorrectly without them, whereas property (U) guarantees Q = P minimizes expected sample loss.In probabilistic forecasting, this property makes the divergence a proper scoring rule.
- Properties of a Divergence: KL has unbiased sample gradients (U) but is not scale sensitive (S), while Wasserstein is ideal under (I, S) but lacks (U).Thus neither divergence simultaneously satisfies all three desired properties.
3 Bias in the Sample Gradient Estimates of the Wasserstein Distance
The paper’s theoretical analysis shows that sample Wasserstein gradients can remain biased and lead optimization toward the wrong minimum. Consequently, naive SGD cannot generally minimize the true Wasserstein loss, although practical unbiased optimization remains an open question.
- Theoretical evidence: Sample Wasserstein gradients are biased even for Bernoulli distributions, with a bias lower bound for every sample size m.This establishes failure of property (U) in a simple estimation setting.
- Theoretical evidence: The minimum of expected sample Wasserstein loss can differ from the minimum of the true Wasserstein loss.The analysis includes examples where a nonzero-entropy distribution is replaced by a zero-entropy solution.
- Implications: Naive stochastic gradient descent therefore cannot in general minimize the Wasserstein loss or converge to the true-loss minimum.The result concerns sample-based SGD and SGD-like methods, not every possible stochastic optimization procedure.
- Wasserstein Bias in the Literature: The authors hypothesize that Wasserstein-gradient bias contributes to instability and poor convergence reported in empirical studies.They connect this hypothesis to prior use of mixed KL-Wasserstein losses and discussion of wrong minima.
- Wasserstein Bias in the Literature: For the dual Wasserstein formulation, unbiasedness appears as m → ∞, but each gradient is substantially more costly and finite-m bias is conjectured to remain.The dual loss is the same as the primal loss, so the reported bias result applies to it as well.
4 The Cramér Distance
The Cramér distance combines Wasserstein-like sensitivity to outcome geometry with unbiased sample gradients. Experiments show that this property matters because sample Wasserstein optimization can converge to poor solutions.
- Definition and Analysis: The Cramér distance is not a metric proper as written, but its square root is a metric.Its square root belongs to the l_p family of metrics.
- Definition and Analysis: The Cramér distance has unbiased sample gradients, unlike the other l_p distances considered.This property supports stochastic optimization whose expected sample loss is minimized at the target distribution.
- Definition and Analysis: The energy distance extends the Cramér distance to multivariate distributions and also possesses sum invariance, scale sensitivity, and unbiased sample gradients.The paper uses the energy distance in the multivariate Cramér GAN construction.
- Comparison to the 1-Wasserstein Metric: The sample Wasserstein minimizer assigns most mass to outcome 0 under the constrained discrete model, despite the exact model being representable.The Cramér solution is relatively close to the 1-Wasserstein solution, while the sample Wasserstein solution is poor.
- Comparison to the 1-Wasserstein Metric: For small sample sizes, stochastic gradient descent fails to find reasonable solutions, and with m = 1 it performs worse than the KL minimizer.Increasing sample size may not be sufficient to guarantee good behavior.
- Ordinal Regression: In ordinal regression, minimizing the Cramér distance yields the lowest RMSE and Wasserstein loss, while KL achieves slightly lower negative log likelihood.The task uses the Year Prediction MSD dataset and evaluates test-set RMSE and sample Wasserstein loss.
5 Cramér GAN
The Cramér GAN replaces the difficult-to-converge Wasserstein critic objective with an energy-distance objective over transformed samples. Experiments report greater diversity and more stable learning than WGAN-GP.
- Cramér GAN: Wasserstein GAN training requires critic convergence in theory, but convergence is rarely achievable and can overfit the empirical training distribution.This motivates a loss that remains useful with imperfect critics.
- Cramér GAN: The Cramér GAN combines the energy distance with a transformation h: R^d → R^k, using k = 256 in the experiments.The generator minimizes the energy distance between transformed real and generated samples, while the critic maximizes it under a gradient-penalty constraint.
- More stable learning: Figure 4 compares approximate Wasserstein distances to the CelebA test set while varying the number of critic updates per generator update.The caption defines Nu as the number of critic updates per generator update.
- Increased diversity: Cramér GAN completions are fairly diverse, whereas WGAN-GP completions are almost deterministic in conditional face completion.The comparison uses right-half completions conditioned on left halves from CelebA 64 × 64 validation images.
- Increased diversity: The paper hypothesizes that Cramér GAN’s increased diversity arises from learning almost deterministic predictions without collapsing generated variation.The hypothesis connects the result to the sample Wasserstein minimizer’s behavior for low-entropy target distributions.
6 Conclusion
The paper concludes that machine-learning losses should incorporate geometry while retaining scale sensitivity, sum invariance, and unbiased estimators. These properties support efficient stochastic optimization.
- Conclusion: Divergences with unbiased estimators can incorporate geometric information while retaining scale sensitivity and sum invariance.The conclusion presents these properties as desirable for machine-learning optimization.
- Conclusion: Correct sample gradients are linked to the convergence properties required for efficient learning.The conclusion frames unbiased gradients as a practical requirement rather than merely a theoretical distinction.
A Proofs
The proofs establish the stated property assignments for KL, total variation, Wasserstein, and related divergences. In particular, the Wasserstein metric lacks unbiased sample gradients.
- Proofs: The KL divergence has property (U), while total variation lacks property (S), and Pinsker’s inequality transfers that failure to KL.The proof also cites prior work for the Wasserstein metric’s properties (I) and (S).
- Proofs: The Wasserstein metric’s lack of property (U) is established in the proof of Theorem 1.This supplies the theoretical basis for the paper’s concern about biased sample gradients.
A.2 Biased estimate
The sample Wasserstein gradient can be biased enough that stochastic optimization converges to a wrong or degenerate minimum, even as sample size increases.
- For Bernoulli distributions, sample Wasserstein gradients differ from true gradients for every θ*∈(0, 1).
- A one-sample Wasserstein gradient estimate is independent of θ and drives descent to 0 or 1 instead of θ*.
- For any sample size m, the minimax gradient bias can remain bounded below by a numerical constant.
- When θ*=1/2 and |θ*−θ|=O(1/√m), the bias is lower-bounded by a constant independent of m.
- The expected sample Wasserstein loss is minimized at the median of the empirical parameter, whereas the true loss is minimized at its mean.
- For low-entropy targets, minimizing the sample Wasserstein loss may produce deterministic solutions, including convergence to 1 instead of θ* near the upper boundary.
A.3 Consistency of the sample 1-Wasserstein gradient
The sample 1-Wasserstein gradient converges to the true gradient as the sample size grows. The proof uses distribution-function convergence and dominated convergence under finite-mean and differentiability assumptions.
- The sample 1-Wasserstein gradient converges to the true gradient as m →∞.
- Uniformly bounded derivatives permit dominated convergence when differentiating with respect to the model parameter.
- The proof represents the 1-Wasserstein distance as the area between the distributions’ cumulative distribution functions.
- The empirical cumulative distribution function converges almost surely to the population cumulative distribution function for every x.
- Finite means are required to keep Cramér distances and expected gradients finite by ensuring sufficiently light distribution tails.
B.2 Image Modelling with PixelCNN
The PixelCNN experiments compare Wasserstein, Cramér, and KL-based training for probabilistic image modelling, where full distributions are unavailable and losses must be optimized from samples. Stochastic minimization of the Wasserstein loss fails, whereas the Cramér distance is readily minimized and achieves lower Wasserstein and Cramér losses, with a trade-off in KL loss.
- Experimental setup: PixelCNN models the joint image distribution on the CelebA 32x32 dataset, which contains 202,599 celebrity-face images.The experiments compare the 1-Wasserstein, Cramér, and KL losses.
- Experimental setup: Because full conditional distributions are almost never available, the experiments require losses that can be optimized from single samples.The sample Wasserstein loss is reported as an upper bound on the true loss, while cross-entropy is reported for KL divergence.
- Results: Minimizing the Wasserstein distance with stochastic gradient fails, while the Cramér distance is as easily minimized as KL.The Cramér-trained model achieves lower Wasserstein and Cramér loss than the alternatives in the reported experiment.
- Results: The Cramér-trained model has higher KL loss than a model trained directly with KL, illustrating a trade-off between optimizing different losses.The paper concludes that Cramér should be preferred to Wasserstein for learning this autoregressive image model.
- Optimization: The Cramér GAN requires four independent samples to estimate gradients with respect to transformation parameters, compared with three samples for its generator loss.The four samples comprise two from the generator and two from the target distribution; the generator-loss estimate uses two generator samples and one target sample.
- Optimization: The training procedure uses a surrogate loss for the critic and also trains the generator to minimize that surrogate loss.The complete procedure is summarized as Algorithm 1.
C.4 Generator Architecture
The generator uses a U-Net conditioned on the left image half and additional Gaussian-noise channels, with independent samples used for the Cramér GAN loss. Generated-image comparisons use matched DCGAN architectures and a larger WGAN-GP minibatch for computational fairness.
- Generator design: The generator architecture is a U-Net previously used for image-to-image translation, without batch normalization or dropout.The generator is conditioned on the left half of the image and extra channels containing Gaussian noise.
- Comparison setup: With the same DCGAN architecture, the paper compares images generated by minimizing raw-image energy distance against images generated with the Cramér GAN loss.The raw-image objective is E(X, Y ), whereas the Cramér GAN objective is E(h(X), h(Y )).
- Generator design: Two independent samples are generated for each input image to compute the Cramér GAN loss.This sampling choice supports the loss computation for the conditional image-generation setting.
- Comparison setup: The Cramér GAN uses minibatch size 64, while WGAN-GP uses 128 for computational fairness.The WGAN-GP minibatch is doubled relative to the Cramér GAN minibatch.
C.5 Critic Architecture
The critic transforms inputs through a deep network and uses distance-based comparisons in the transformed space. The paper connects this construction to energy distance and squared MMD, while reporting that the transformation is important for image generation.
- Critic transformation: The transformation h(x) is a traditional deep convolutional network with 256 outputs and no batch normalization.More outputs are reported to be better, and batch normalization would conflict with the gradient penalty.
- Distance construction: The energy distance can replace the 2-norm with an α-norm; for α in (0, 2), it is zero exactly when the distributions are identical.At α = 2, the energy distance reduces to the squared difference between expectations.
- Distance construction: The energy distance is equivalent to squared maximum mean discrepancy with kernel −∥x − y∥2, so the transformed loss can be viewed as squared MMD.The experiments found ∥x − y∥α distances more stable than Gaussian or Laplacian-kernel distances.
- Image-generation implication: Minimizing the energy distance on raw images does not work well, whereas the learned transformation h(x) is very helpful.The paper illustrates this contrast with generated samples from raw-image energy distance and transformed-space Cramér GAN loss.
- Related approach: A related approach combines generative moment matching with a learned transformation driven to be injective by an additional auto-encoder loss.This work is identified as independent from the paper's approach.