# Wasserstein Task Embedding for Measuring Task Similarities

arXiv: https://arxiv.org/abs/2208.11726 · HTML: https://mlsanity.com/paper/2208.11726/

## Key facts

- Cited by 25 papers in the map (data through 14 Oct 2026).
- Builds on 29 of its 61 references that are in the map.
- First posted on arXiv 24 Aug 2022.

## Builds on (most cited first, 29 of 29)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms](https://mlsanity.com/paper/1708.07747/) (2017)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Computational Optimal Transport](https://mlsanity.com/paper/1803.00567/) (2018)
- [Taskonomy: Disentangling Task Transfer Learning](https://mlsanity.com/paper/1804.08328/) (2018)
- [Deep Learning for Classical Japanese Literature](https://mlsanity.com/paper/1812.01718/) (2018)
- [Joint Distribution Optimal Transportation for Domain Adaptation](https://mlsanity.com/paper/1705.08848/) (2017)
- [Near-linear time approximation algorithms for optimal transport via Sinkhorn iteration](https://mlsanity.com/paper/1705.09634/) (2017)
- [Efficiently Identifying Task Groupings for Multi-Task Learning](https://mlsanity.com/paper/2109.04617/) (2021)
- [Task2Vec: Task Embedding for Meta-Learning](https://mlsanity.com/paper/1902.03545/) (2019)
- [Neural Architecture Search with Bayesian Optimisation and Optimal Transport](https://mlsanity.com/paper/1802.07191/) (2018)
- [LEEP: A New Measure to Evaluate Transferability of Learned Representations](https://mlsanity.com/paper/2002.12462/) (2020)
- [Adaptive Gradient-Based Meta-Learning Methods](https://mlsanity.com/paper/1906.02717/) (2019)
- [Geometric Dataset Distances via Optimal Transport](https://mlsanity.com/paper/2002.02923/) (2020)
- [Transferability and Hardness of Supervised Classification Tasks](https://mlsanity.com/paper/1908.08142/) (2019)
- [Gromov-Wasserstein Learning for Graph Matching and Node Embedding](https://mlsanity.com/paper/1901.06003/) (2019)
- [Gromov-Wasserstein Alignment of Word Embedding Spaces](https://mlsanity.com/paper/1809.00013/) (2018)
- [Representation Similarity Analysis for Efficient Task taxonomy & Transfer Learning](https://mlsanity.com/paper/1904.11740/) (2019)
- [Toward Understanding Catastrophic Forgetting in Continual Learning](https://mlsanity.com/paper/1908.01091/) (2019)
- [Wasserstein Embedding for Graph Learning](https://mlsanity.com/paper/2006.09430/) (2020)
- [Hierarchical Optimal Transport for Document Representation](https://mlsanity.com/paper/1906.10827/) (2019)
- [Learning Wasserstein Embeddings](https://mlsanity.com/paper/1710.07457/) (2017)
- [Domain2Vec: Domain Embedding for Unsupervised Domain Adaptation](https://mlsanity.com/paper/2007.09257/) (2020)
- [An Information-Geometric Distance on the Space of Tasks](https://mlsanity.com/paper/2011.00613/) (2020)
- [Wassmap: Wasserstein Isometric Mapping for Image Manifold Learning](https://mlsanity.com/paper/2204.06645/) (2022)
- [Transferability-Guided Cross-Domain Cross-Task Transfer Learning](https://mlsanity.com/paper/2207.05510/) (2022)
- [P2L: Predicting Transfer Learning for Images and Semantic Relations](https://mlsanity.com/paper/1908.07630/) (2019)
- [Task-similarity Aware Meta-learning through Nonparametric Kernel Regression](https://mlsanity.com/paper/2006.07212/) (2020)
- [Selecting task with optimal transport self-supervised learning for few-shot classification](https://mlsanity.com/paper/2204.00289/) (2022)

## Built on by (most cited first, 25 of 25)

- [InsCL: A Data-efficient Continual Learning Paradigm for Fine-tuning Large Language Models with Instructions](https://mlsanity.com/paper/2403.11435/) (2024)
- [Lifelong Reinforcement Learning with Modulating Masks](https://mlsanity.com/paper/2212.11110/) (2022)
- [Characterizing Out-of-Distribution Error via Optimal Transport](https://mlsanity.com/paper/2305.15640/) (2023)
- [A Survey of Continual Reinforcement Learning](https://mlsanity.com/paper/2506.21872/) (2025)
- [Wasserstein Wormhole: Scalable Optimal Transport Distance with Transformers](https://mlsanity.com/paper/2404.09411/) (2024)
- [Understanding the Transferability of Representations via Task-Relatedness](https://mlsanity.com/paper/2307.00823/) (2023)
- [Sharing Lifelong Reinforcement Learning Knowledge via Modulating Masks](https://mlsanity.com/paper/2305.10997/) (2023)
- [Lightspeed Geometric Dataset Distance via Sliced Optimal Transport](https://mlsanity.com/paper/2501.18901/) (2025)
- [Statistical Context Detection for Deep Lifelong Reinforcement Learning](https://mlsanity.com/paper/2405.19047/) (2024)
- [LCOT: Linear circular optimal transport](https://mlsanity.com/paper/2310.06002/) (2023)
- [Many-Task Federated Fine-Tuning via Unified Task Vectors](https://mlsanity.com/paper/2502.06376/) (2025)
- [Beyond Knowledge Silos: Task Fingerprinting for Democratization of Medical Imaging AI](https://mlsanity.com/paper/2412.08763/) (2024)
- [Task Attribute Distance for Few-Shot Learning: Theoretical Analysis and Applications](https://mlsanity.com/paper/2403.03535/) (2024)
- [Clone-Robust Weights in Metric Spaces: Handling Redundancy Bias for Benchmark Aggregation](https://mlsanity.com/paper/2502.03576/) (2025)
- [Policy Search, Retrieval, and Composition via Task Similarity in Collaborative Agentic Systems](https://mlsanity.com/paper/2506.05577/) (2025)
- [Busemann Functions in the Wasserstein Space: Existence, Closed-Forms, and Applications to Slicing](https://mlsanity.com/paper/2510.04579/) (2025)
- [Learning Embeddings for Sequential Tasks Using Population of Agents](https://mlsanity.com/paper/2306.03311/) (2023)
- [On Wasserstein distances for affine transformations of random vectors](https://mlsanity.com/paper/2310.03945/) (2023)
- [Structured Approximations of Measures](https://mlsanity.com/paper/2310.09149/) (2023)
- [Personalized Federated Learning under Model Dissimilarity Constraints](https://mlsanity.com/paper/2505.07575/) (2025)
- [Beyond Model Base Retrieval: Weaving Knowledge to Master Fine-grained Neural Network Design](https://mlsanity.com/paper/2507.15336/) (2025)
- [RADAR: Relative Angular Divergence Across Representations](https://mlsanity.com/paper/2605.23028/) (2026)
- [LWM-CDE: A Representation Space for Wireless Data Reasoning and Transferability](https://mlsanity.com/paper/2605.24077/) (2026)
- [Making two action heads agree: coordination mechanisms and a runtime collapse certificate for flow-matching policies](https://mlsanity.com/paper/2608.15748/) (2026)
- [How to Find and Reuse Policies for Continuous Adaptation in Lifelong Reinforcement Learning](https://mlsanity.com/paper/2610.03119/) (2026)

## Abstract (arXiv)

Measuring similarities between different tasks is critical in a broad spectrum of machine learning problems, including transfer, multi-task, continual, and meta-learning. Most current approaches to measuring task similarities are architecture-dependent: 1) relying on pre-trained models, or 2) training networks on tasks and using forward transfer as a proxy for task similarity. In this paper, we leverage the optimal transport theory and define a novel task embedding for supervised classification that is model-agnostic, training-free, and capable of handling (partially) disjoint label sets. In short, given a dataset with ground-truth labels, we perform a label embedding through multi-dimensional scaling and concatenate dataset samples with their corresponding label embeddings. Then, we define the distance between two datasets as the 2-Wasserstein distance between their updated samples. Lastly, we leverage the 2-Wasserstein embedding framework to embed tasks into a vector space in which the Euclidean distance between the embedded points approximates the proposed 2-Wasserstein distance between tasks. We show that the proposed embedding leads to a significantly faster comparison of tasks compared to related approaches like the Optimal Transport Dataset Distance (OTDD). Furthermore, we demonstrate the effectiveness of our proposed embedding through various numerical experiments and show statistically significant correlations between our proposed distance and the forward and backward transfer between tasks.