PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning with a Wasserstein Loss

Charlie Frogner, Chiyuan Zhang, Hossein Mobahi, Mauricio Araya-Polo, Tomaso Poggio

arXiv:1506.05439v3cs.LGcs.CVstat.ML

TL;DR

Multi-label prediction can benefit from metrics that capture semantic relationships among output labels, but standard losses often ignore this structure. The paper introduces a Wasserstein loss with efficient regularized optimization and an extension to unnormalized measures. It reports improved performance over a baseline without the output metric and characterizes the loss with a statistical learning bound.

  • Problem

    Many learning tasks have output spaces with known or estimated semantic similarity, while common divergence losses treat output dimensions independently and ignore that structure.

  • Method

    The paper uses Wasserstein distance as a loss over finite non-negative measures, optimizes an entropically regularized approximation, and extends it to unnormalized measures.

  • Results

    The Wasserstein loss improves performance over a baseline without the output metric on a real-data tag prediction problem and produces predictions closer to ground truth under semantically related label noise.

  • Takeaways & Limitations

    The loss can encourage prediction smoothness with respect to a chosen metric on the output space.

  • Takeaways & Limitations

    The exact Wasserstein loss is computationally costly, motivating the paper's use of an efficiently computed regularized approximation.

Abstract

from arXiv · show

Learning to predict multi-label outputs is challenging, but in many problems there is a natural metric on the outputs that can be used to improve predictions. In this paper we develop a loss function for multi-label learning, based on the Wasserstein distance. The Wasserstein distance provides a natural notion of dissimilarity for probability measures. Although optimizing with respect to the exact Wasserstein distance is costly, recent work has described a regularized approximation that is efficiently computed. We describe an efficient learning algorithm based on this regularization, as well as a novel extension of the Wasserstein distance from probability measures to unnormalized measures. We also describe a statistical learning bound for the loss. The Wasserstein loss can encourage smoothness of the predictions with respect to a chosen metric on the output space. We demonstrate this property on a real-data tag prediction problem, using the Yahoo Flickr Creative Commons dataset, outperforming a baseline that doesn't use the metric.

1 Introduction

The paper proposes a Wasserstein loss that uses known semantic structure among output labels, enabling predictions to remain similar to ground truth under semantically related errors. It develops an efficient learning procedure and reports improved performance over a metric-agnostic baseline.

  • Problem setting: The paper formulates multi-label prediction as learning a non-negative measure over a finite output set.This includes score or probability vectors in classification and pixel-support measures in semantic segmentation.
  • Metric structure: Unlike information-divergence losses that treat output dimensions independently, the proposed loss uses semantic relationships among labels.The ground metric can encode relationships such as those represented by the WordNet hierarchy for ImageNet categories.
  • Proposed loss: The Wasserstein loss incorporates a chosen metric on the output space to measure transport cost between predictions and target labels.It defines dissimilarity through the cost of moving predicted mass to match the target mass.
  • Synthetic experiment: The Wasserstein loss yields predictions closer to ground truth across all tested noise levels when labels are randomly switched to nearby semantic classes.The synthetic experiment compares against the standard multiclass logistic loss.
  • Contributions: The authors develop empirical-risk minimization with entropic regularization, extend the loss to unnormalized measures, and provide a statistical learning bound.These contributions address efficient optimization and broader output representations.

2 Related work

Prior work used smoothness in label spaces through regularization and graphical-model post-processing, while optimal transport supplied a natural metric for probability distributions. This paper positions its contribution as the first use of optimal transport as a loss in discriminative learning.

  • Label-space smoothness: Existing approaches exploit label-space smoothness through regularization and post-processing with graphical models.These methods incorporate prior metric structure without using optimal transport directly as the discriminative loss.
  • Optimal transport: Optimal transport provides a natural distance for probability distributions over metric spaces and has been applied to several related estimation problems.The cited applications include barycenter estimation and related distribution-matching tasks.
  • Positioning: The paper claims to be the first to use optimal transport as a loss function in a discriminative learning framework.It distinguishes this setting from prior theoretical work on estimators minimizing empirical optimal transport cost.

3 Learning with a Wasserstein loss

The learning framework predicts non-negative measures over a metrized finite output space and evaluates them with optimal transport. The resulting Wasserstein loss penalizes transporting mass according to output distance, while exact optimization motivates efficient approximations.

  • Learning setup: The framework maps inputs X into non-negative measures over a finite output set K equipped with a ground metric d_K.The metric represents semantic similarity between output dimensions, and predictors are parameterized within a hypothesis space.
  • Learning setup: Empirical risk minimization uses a loss function as a tractable surrogate for expected prediction risk.Training examples are sampled independently from an unknown joint distribution over inputs and outputs.
  • Motivation: Information-divergence losses evaluate output components independently and therefore ignore metric structure on the output set.The Wasserstein formulation is introduced to account for relationships among output dimensions.
  • Optimal transport: The Wasserstein distance measures the cheapest transport plan that moves predicted mass to match target mass under a cost derived from the ground metric.For discrete probability measures, the outputs are histograms in the simplex.
  • Exact Wasserstein loss: The exact Wasserstein loss penalizes longer transport more heavily according to the ground metric.Its optimal transport plan moves the predicted mass distribution h(x) to the target distribution y.

4 Efficient optimization via entropic regularization

The paper replaces costly exact Wasserstein optimization with entropic regularization and extends the approach to unnormalized measures through relaxed marginal constraints. The resulting procedures are efficiently computable and recover smoothed or unregularized transport behavior under suitable parameter limits.

  • Computing an exact Wasserstein subgradient requires solving a linear program with O(K^2) constraints, which can be prohibitive for gradient-based learning.
  • 4.1 Entropic regularization of optimal transport: Entropic regularization makes the transport objective strictly convex and enables efficient approximation of the transport matrix and loss subgradient.
  • 4.1 Entropic regularization of optimal transport: The regularized transport matrix is obtained by diagonal scaling of K = e^−λM−1, reducing optimization to a matrix-balancing problem solvable with Sinkhorn-Knopp iterations.
  • 4.1 Entropic regularization of optimal transport: For simplex-valued predictions and targets, the regularized objective approximates exact Wasserstein distance for sufficiently large λ, with gradients derived from the optimal scaling vector.
  • 4.2 Extending smoothed transport to the learning setting: The paper extends smoothed transport to unnormalized measures by replacing equality constraints on transport marginals with soft generalized-KL penalties.
  • 4.3 Relaxed transport: The relaxed optimal transport matrix remains a diagonal scaling of K and is computed through a Sinkhorn-like fixed-point iteration, with an explicit gradient for unconstrained predictions.
  • 4.3 Relaxed transport: For normalized measures, the relaxed objective converges to smoothed transport and then to unregularized Wasserstein distance as λ, γa, and γb increase.
  • 4.3 Relaxed transport: Empirically, iteration convergence is nearly independent of output dimension K, but the paper does not prove a convergence guarantee.

5 Statistical Properties of the Wasserstein loss

The paper establishes statistical guarantees for empirical risk minimization with the Wasserstein loss, including convergence toward the best achievable loss in the hypothesis space and bounds on semantic prediction error.

  • The framework uses predictors formed by composing a softmax layer with a base hypothesis space, producing outputs in the probability simplex.The learning setup assumes i.i.d. samples and empirical risk minimization over this composed hypothesis space.
  • The Rademacher complexity term decays with training-set size for common neural-network and kernel-machine hypothesis spaces.This makes the guarantee approach the best achievable loss for H as the sample size increases.
  • In multiclass classification, the guarantee concerns expected semantic distance between predictions and ground-truth classes rather than classification error.The predictor is obtained by taking the argmax of the empirical-risk-minimizing model's output.

6 Empirical study

The empirical study examines metric-sensitive prediction on MNIST and Yahoo/Flickr tag data. Wasserstein-based objectives produce semantically smoother predictions and outperform the KL-divergence baseline on the tag-prediction task, with a tunable trade-off against maximum likelihood.

  • 6.1 Impact of the ground metric: As the MNIST metric exponent p increases, predicted probability mass spreads more evenly across numerically neighboring digits.At p → 0, the true digit receives probability near 1; as p → ∞, predictions converge toward a uniform distribution.
  • 6.2 Flickr tag prediction: The Flickr experiment combines Wasserstein loss W_p with KL loss, using KL alone as the baseline for top-K cost and AUC evaluation.The combination is trained on 1000 tags and two random sets of 10,000 images, with tag distances derived from word2vec embeddings.
  • 6.2 Flickr tag prediction: The combined Wasserstein loss continues to outperform the baseline after redundant labels are removed, although both algorithms perform worse on the harder dataset.This tests a setting that simulates a single user's selection among semantically similar applicable tags.
  • 6.2 Flickr tag prediction: The top-K-cost optimum weights Wasserstein more heavily than the AUC optimum, reflecting a trade-off between semantic smoothness and maximum likelihood.The paper attributes this to Wasserstein favoring semantically similar mispredictions, sometimes at the cost of lower AUC.
  • 6.2 Flickr tag prediction: Figure 7 compares ground-truth tags with tags proposed by the Wasserstein-based algorithm and the baseline on selected Flickr images.The examples include cases where both algorithms produce semantically relevant predictions despite little overlap with the ground truth.

7 Conclusions and future work

The paper concludes that Wasserstein loss provides an efficient, statistically justified way to learn non-negative measures while encouraging metric-based smoothness. It identifies connections to Markov random fields as a direction for future work.

  • The proposed loss learns non-negative measures over finite sets using the Wasserstein distance and supports smoothness relative to a chosen output-space metric.The paper demonstrates improved performance over a baseline that does not incorporate the metric on a real-data tag-prediction problem.
  • Entropic regularization makes Wasserstein optimization efficiently computable, and the paper extends the regularized loss to unnormalized measures while preserving efficiency.The paper also provides a statistical learning bound for the loss.
  • Future work may investigate connections between Wasserstein loss and Markov random fields, which can encourage prediction smoothness through inference at prediction time.

B Statistical Learning Bounds

The paper derives statistical generalization control for Wasserstein-loss classifiers by combining uniform concentration, Lipschitz properties, and Rademacher-complexity bounds. A softmax-based hypothesis construction makes the loss compatible with these tools and with standard classifiers.

  • Generalization analysis: Uniform concentration bounds control the deviation between empirical Wasserstein risk and expected risk with high probability.The analysis uses Rademacher complexity and McDiarmid’s inequality to obtain this control.
  • Lipschitz reduction: The Wasserstein loss is Lipschitz continuous, enabling its function class to be bounded through generalized Talagrand’s lemma.The resulting complexity bound involves the underlying classifier class and an identity-function class.
  • Probability-output construction: A softmax layer maps arbitrary model outputs to valid probability distributions, resolving the fact that Wasserstein loss is originally defined only on probability distributions.This construction is consistent with logistic regression and neural networks.
  • Special case: For one-hot ground-truth labels, the transport plan is unique, so the Wasserstein loss can be computed directly.The paper obtains the corresponding generalization result by applying the main theorem with p = 1.
  • Optimization: The empirical studies use linear logistic regression, stochastic gradient descent, mini-batches of 100, and 100,000 training iterations.The Wasserstein loss and subgradient are computed through the Sinkhorn-based routine described by the algorithms.

E.1 Noisy label example

The noisy-label experiment models semantically related classes as nearby lattice points and compares Wasserstein loss with KL-divergence loss under increasing label noise. Performance is evaluated by geometric distance between predicted and true classes.

  • Synthetic setup: Labels are placed on a D × D lattice, with Euclidean distance representing semantic similarity between classes.Observations are Gaussian around class vertices, and neighboring labels are treated as semantically similar.
  • Synthetic setup: Training labels are flipped with probability t to uniformly selected neighboring categories, producing controlled semantic label noise.Figure 8 illustrates a 3×3 lattice at noise levels t = 0.1 and t = 0.5.
  • Evaluation: The experiment repeats training across t = 0.1, 0.2, . . . , 0.9 and D = 3, 4, . . . , 7 using KL-divergence or Wasserstein loss.Performance is measured by mean Euclidean distance between predicted and true classes on the test set.

E.2 Full figure for the MNIST example

The appendix provides the complete version of the MNIST probability-curve figure from Section 6.1. The expanded figure is presented as Figure 9.

  • Figure reference: Figure 9 contains the full version of Figure 4 from the MNIST experiment.The passage identifies Figure 9 as an expanded presentation of the earlier figure.
  • Figure reference: The figure belongs to the paper’s MNIST experiment on Wasserstein-loss behavior.
  • Figure reference: Figure 9 is located in the appendix as a supplementary version of the Section 6.1 result.

E.3 Details of the Flickr tag prediction experiment

The Flickr experiment constructs a filtered 1,000-tag dataset, derives semantic distances from word2vec embeddings, and trains a Wasserstein–KL logistic model. It also evaluates a reduced-redundancy variant with randomly selected representative tags.

  • Metric and features: Tag distances are computed as Euclidean distances between unit word2vec embeddings.This metric represents semantic relationships among the selected tags.
  • Training setup: The model combines Wasserstein loss with KL divergence in a multiclass linear logistic regression model.The Wasserstein term uses 10 Sinkhorn–Knopp iterations, p = 13, and λ = 50.
  • Reduced-redundancy variant: A reduced-redundancy dataset clusters semantically similar ground-truth tags using threshold 1.3 and randomly retains one tag per cluster.This simulates images tagged by a single user rather than with multiple equivalent tags.
  • Qualitative examples: Figures 10 and the associated examples show ground-truth tags alongside predictions from the proposed algorithm and the baseline.
Loading 1506.05439v3…