PaperlayerView exact revision on arXiv

Source-linked AI summary

Geometric Dataset Distances via Optimal Transport

David Alvarez-Melis, Nicolò Fusi

arXiv:2002.02923v1cs.LGstat.ML

TL;DR

Dataset similarity is difficult to quantify when labels differ, and existing approaches may rely on heuristics, strong label assumptions, or trained probe models. The paper proposes a label-aware optimal-transport distance between feature-label distributions that supports disjoint label sets without training. It reports meaningful dataset comparisons and strong predictive correlation with transferability across experimental settings.

  • Problem

    Dataset similarity lacks an efficient, principled measure that remains useful when label sets differ or methods depend on task-specific trained models.

  • Method

    The paper compares feature-label distributions with a hybrid Euclidean-Wasserstein optimal-transport distance, modeling labels as feature distributions.

  • Results

    The distance provides meaningful, interpretable dataset comparisons and is highly predictive of transferability across domains, tasks, and data modalities.

  • Takeaways & Limitations

    The framework supports scalable and flexible transfer-learning comparisons with minimal assumptions and can compare datasets whose label sets are completely unrelated or disjoint.

  • Takeaways & Limitations

    The framework assumes datasets have feature spaces of the same dimension, while its efficient implementation models same-label point groups as Gaussian distributions.

Abstract

from arXiv · show

The notion of task similarity is at the core of various machine learning paradigms, such as domain adaptation and meta-learning. Current methods to quantify it are often heuristic, make strong assumptions on the label sets across the tasks, and many are architecture-dependent, relying on task-specific optimal parameters (e.g., require training a model on each dataset). In this work we propose an alternative notion of distance between datasets that (i) is model-agnostic, (ii) does not involve training, (iii) can compare datasets even if their label sets are completely disjoint and (iv) has solid theoretical footing. This distance relies on optimal transport, which provides it with rich geometry awareness, interpretable correspondences and well-understood properties. Our results show that this novel distance provides meaningful comparison of datasets, and correlates well with transfer learning hardness across various experimental settings and datasets.

1. Introduction

Dataset similarity is important for transferring knowledge when target labels are scarce, but principled and efficient comparison remains difficult, especially across differing label sets. The paper addresses these challenges with a geometric, label-aware optimal-transport distance and reports predictive transfer-learning evidence.

  • Scarce application-specific labels motivate combining, adapting, and transferring knowledge across datasets and domains.
  • Quantifying dataset distance remains open because labels are difficult to compare when tasks have nonidentical label sets.
  • Existing similarity measures often use heuristic proxies, strong label assumptions, or probe models, while lacking guarantees and requiring task-specific information.
  • The proposed distance compares feature-label distributions with optimal transport, modeling labels as distributions over feature vectors.
  • The framework compares datasets with completely unrelated or disjoint label sets whenever a feature distance is available.
  • The paper introduces scalable computation and reports empirical evidence that its distance predicts transfer-learning success across domains, tasks, and modalities.

2. Related Work

Related approaches measure dataset similarity through discrepancy functions, parameter sensitivity, distributional embeddings, hierarchical optimal transport, or label-guided transport. These methods differ in their dependence on hypothesis classes, trained probe networks, within-dataset comparisons, or domain-specific label information.

  • Discrepancy Distance: Discrepancy distances depend on a loss function and hypothesis class and quantify distributional dissimilarity through a supremum over that class.
  • Dataset Distance via Parameter Sensitivity: Fisher-information task representations can recover taxonomic similarities and predict whether a feature extractor will perform well on a new task.
  • Dataset Distance via Parameter Sensitivity: Fisher-information similarity differs from the paper’s approach because it uses a trained probe network whose characteristics influence the resulting geometry.
  • Optimal Transport-based distributional distances: Distributional-embedding methods represent complex objects as distributions and compare them with optimal transport, including elliptical distributions or discrete measures.
  • Hierarchical OT distances: The proposed distance is hierarchical optimal transport because its ground metric is itself defined through an optimal-transport problem.
  • Optimal Transport for Domain Adaptation: Prior label-guided transport methods target unsupervised domain adaptation, use labels from only one domain, and do not define metrics between labels or proper distances.

3. Background on Optimal Transport

Optimal transport compares probability measures by minimizing a geometry-based ground cost over couplings, with Wasserstein distances as a metric-based special case. Entropy regularization enables more efficient computation, while label-aware dataset comparison can reverse conclusions from label-agnostic transport.

  • Optimal transport compares probability distributions using the geometry of their underlying space.
  • The Kantorovich formulation minimizes ground transport cost over couplings whose marginals equal the compared probability measures.
  • Using a metric-powered ground cost yields the p-Wasserstein distance, with p = 1 also called the Earth Mover’s Distance.
  • Finite samples are represented as discrete measures, reducing the transport problem to a linear program with pairwise costs stored in a matrix.
  • Entropy regularization makes optimal-transport computation more efficient and improves sample complexity relative to the original problem.
  • Figure 1 shows that label-aware and label-agnostic OT can rank the relative closeness of two dataset pairs in opposite ways.

4. Optimal Transport between Datasets

The paper defines datasets as distributions over feature-label pairs and lifts feature- and label-level geometry into an optimal-transport distance. Labels are represented by distributions over their associated features, yielding a proper metric that also admits a tractable Gaussian realization.

  • Datasets are sets of feature-label pairs z=(x,y) in a product space Z=X×Y, represented probabilistically through their joint distributions.
  • Labels are compared through distributions of their associated feature vectors, using Wasserstein distance rather than requiring a predefined metric between label sets.The associated feature set is N_D(y)={x∈X|(x,y)∈D}.
  • Optimal transport lifts the point-wise feature-label metric into a distance between dataset measures, defined as the minimum cost of coupling them.
  • OTDD is a proper distance: d_OT(D_A,D_B) defines a valid metric on measures over feature and label-distribution pairs.
  • The tractable d_OT-N realization models each label-feature distribution as a Gaussian whose parameters are the sample mean and covariance of its feature neighborhood.The Gaussian choice uses the analytic 2-Wasserstein distance between Gaussian distributions.
  • The Gaussian realization lower-bounds the distance computed from the true label distributions and is exact when those distributions are Gaussian or elliptical.

5. Computational Considerations

The paper addresses the computational cost of nested optimal-transport calculations with regularization, precomputation, approximations, and Gaussian label representations. The Gaussian distance is substantially more scalable in the stated worst-case comparison, while some speed-ups rely on restrictive covariance assumptions.

  • Scalability is central because dataset distances may need repeated computation on large samples, while general Wasserstein estimation is computationally expensive.
  • A fast augmented-representation implementation requires the feature metric to match the transport ground metric and all covariance matrices to commute.The paper focuses instead on an exact version without diagonal or commuting-covariance approximations.
  • The implementation precomputes all label-to-label distances and uses a two-pass stable online batch algorithm for class means and covariances.The online algorithm avoids loading every example from a class into memory simultaneously.
  • For d_OT-N, label-to-label distances can be accelerated with Newton-Schulz approximations to matrix square roots instead of exact eigendecompositions.
  • For datasets of sizes n and m, both distances incur O(nm log(max{n,m})τ^-3) cost for solving the global OT problem τ-approximately.
  • O(n^5 log n) versus O(n^2d+d^3): when n=m and the largest class has size O(n), d_OT is prohibitive while d_OT-N is more tractable.This comparison concerns the pairwise-distance computation step, which dominates in most practical applications.

6. Experiments

Experiments evaluate OTDD for dataset selection, augmentation comparison, and transfer learning across image and text datasets. Across these settings, dataset distance is strongly correlated with adaptation outcomes, while also revealing interpretable correspondences and sample-size behavior.

  • Dataset Selection for Transfer Learning: Most optimal correspondences between MNIST and USPS lie along the block diagonal, indicating class-coherent matches across datasets.The corresponding digit-class distances are generally smaller than off-diagonal distances, although USPS class 0 is unusually far from MNIST digits.
  • Dataset Selection for Transfer Learning: As sample sizes grow, OTDD converges toward a fixed value, whereas small USPS samples exhibit greater variability than MNIST samples.The wider USPS variability suggests greater heterogeneity in that dataset.
  • Dataset Selection for Transfer Learning: OTDD pairwise distances identify MNIST–EMNIST as the closest pair, while FASHION-MNIST is comparatively far from the other datasets.MNIST and USPS are not the closest pair despite both containing digits.
  • Dataset Selection for Transfer Learning: OTDD is highly predictive of transferability across *NIST datasets, with EMNIST producing the best adaptation to MNIST.Transferability is measured as the relative decrease in classification error after source pretraining and target fine-tuning, averaged over 10 random-seed repetitions.
  • Distance-Driven Data Augmentation: Across MNIST-to-USPS and Tiny-ImageNet-to-CIFAR-10 transformations, OTDD strongly correlates with adaptation: cropping improves transfer while most rotations degrade it.The augmentation experiments compare transformed source replicas against the target dataset using their distance and transferability.
  • Transfer Learning for Text Classification: For text classification, sentences are embedded with base BERT before OTDD computation, and OTDD remains highly correlated with transferability under 100 target examples per class.In this setting, adaptation can drastically degrade performance, suggesting that off-the-shelf BERT may already provide strong initialization for many tasks.

7. Discussion

The proposed dataset distance is presented as scalable and flexible for realistic transfer-learning scenarios, while retaining theoretical, interpretability, and low-assumption benefits. The discussion identifies extensions to feature-space comparability, class-conditional distribution models, and task-specific information.

  • 7. Discussion: The distance is scalable and flexible enough for realistic transfer-learning scenarios while requiring minimal assumptions on the underlying datasets.The authors also emphasize theoretical properties and interpretable comparisons.
  • 7. Discussion: The framework assumes that the compared datasets use feature spaces of the same dimension.Gromov-Wasserstein distance is suggested for datasets whose features are not directly comparable.
  • 7. Discussion: The efficient implementation models groups of points sharing a label as Gaussian distributions.The authors suggest extending it to broader distributions with analytically or efficiently computable Wasserstein distances.
  • 7. Discussion: The distance excludes the loss function and predictor function class, although future work may incorporate them with minimal training.The paper frames independence from these choices as important while leaving their integration open.

A. Proof of Proposition 4.1

The proof establishes that the feature-label-pair cost is a metric, allowing the associated optimal transport value to define a dataset distance. It also uses the Gelbrich bound, with equality for Gaussian distributions and certain elliptical distributions.

  • A. Proof of Proposition 4.1: The optimal transport value is a distance when its cost function is a metric on the underlying space.The proof therefore focuses on showing that dZ is a distance.
  • A. Proof of Proposition 4.1: dZ is symmetric because both the feature distance dX and Wasserstein distance Wp are symmetric.The proof also invokes the fact that both component functions are distances.
  • A. Proof of Proposition 4.1: Minkowski’s inequality supplies the triangle inequality, making dZ a metric on Z = X × P(X).Consequently, the optimal transport value using dZ as its cost function is itself a distance.
  • A. Proof of Proposition 4.1: The proof applies Gelbrich’s bound to the 2-Wasserstein distance between distributions using their means and covariance matrices.The lemma is applied to feature-label pairs and couplings in the optimal transport problem.
  • A. Proof of Proposition 4.1: Gelbrich’s bound is exact for Gaussian distributions and also for elliptical distributions sharing the same density generator.This yields equality between the corresponding quantities in equation (13).

C. Time Complexity Analysis

The complexity analysis decomposes computation into pairwise feature-label distances and a global optimal transport problem. It provides worst-case costs for the general and Gaussian-approximation variants, alongside the dataset and dimensionality assumptions used in the analysis.

  • C. Time Complexity Analysis: The analysis assumes source and target datasets have n and m labeled examples in R^d and ks and kt classes.The implementation discussion also takes matrix multiplication exponent ω = 3 for simplicity.
  • C. Time Complexity Analysis: The distance computation has two stages: pairwise pointwise distances requiring label-to-label OT subproblems, followed by a global OT problem between the samples.The global problem operates on discrete distributions containing n and m examples.
  • C. Time Complexity Analysis: The global OT problem can be solved exactly with interior-point methods or Orlin’s algorithm, or approximately with Sinkhorn iterations.The stated approximate complexity is O(nm log(max{n, m})τ^-3).
  • C. Time Complexity Analysis: For the Gaussian approximation, precomputing pointwise distances costs O(nmd + ksktd^3 + d^2n(ks + kt)).The terms include pairwise feature-distance computation, class-pair matrix square roots, and per-class mean and covariance computation.

E. Optimization and Training Details

The experiments use distinct neural architectures and optimization schedules for digit, image-adaptation, and text-classification settings. Text inputs are embedded with pretrained BERT, and the implementation relies on POT and GeomLoss.

  • E. Optimization and Training Details: NIST adaptation experiments use LeNet-5 with ReLU, 20 source-training epochs, and 10 target fine-tuning epochs.ADAM uses learning rate 1 × 10^-3 and weight decay 1 × 10^-6 in both stages.
  • E. Optimization and Training Details: Tiny-ImageNet-to-CIFAR-10 adaptation uses ResNet-50, trained for 300 epochs and fine-tuned for 30 target epochs.SGD uses momentum 0.9, weight decay 1 × 10^-4, and learning rates 0.1 for source training and 0.01 for target fine-tuning.
  • E. Optimization and Training Details: Text experiments embed sentences with pretrained bert-base-uncased BERT before source and restricted-target fine-tuning.Source fine-tuning lasts 10 epochs and target fine-tuning lasts 2 epochs using ADAM with learning rate 2 × 10^-5.
  • E. Optimization and Training Details: The OTDD implementation uses the POT and GeomLoss Python packages.The cited package resources are listed alongside the implementation description.

F. Robustness of the Distance

The robustness analysis examines distances computed on subsets of varying size, comparing MNIST and USPS across repeated trials and two regularization settings.

  • Distances are computed on subsets with varying sizes to assess robustness.
  • Rows represent MNIST and columns represent USPS in the analysis.
  • The experiment uses 10 random repetitions and two values of the regularization parameter ε.
Loading 2002.02923v1…