# Wasserstein Task Embedding for Measuring Task Similarities

Xinran Liu1, Yikun Bai1, Yuzhe Lu1, Andrea Soltoggio2, Soheil Kolouri1

###### Abstract

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 on a wide variety of image recognition datasets.

## Introduction

Learning from a broad spectrum of tasks and transferring knowledge between
them is a cornerstone of intelligence, and primates perfectly exemplify this
characteristic. Modern Machine Learning (ML) is rapidly moving toward multi-
task learning, and there is great interest in methods that can integrate,
rapidly adapt, and seamlessly transfer knowledge between tasks. When learning
from multiple possibly heterogeneous tasks, it is essential to understand the
relationships between the tasks and their fundamental properties. It is,
therefore, highly desirable to define (dis)similarity measures between tasks
that will allow one to cluster tasks, have better control over the forward and
backward transfer, and ultimately require less supervision for learning tasks.

There has been an increasing interest in assessing task similarities and their
relationship with forward and backward knowledge transfer among tasks. For
instance, various recent works look into the selection of good source
tasks/models for a given target task to maximize the forward transfer to the
target task (Achille et al. 2019; Zamir et al. 2018; Bao et al. 2019;
Bhattacharjee et al. 2020; Fifty et al. 2021). Others have demonstrated the
relationship between negative backward transfer (i.e., catastrophic
forgetting) and task similarities (Nguyen et al. 2019).

Many existing methods for measuring task similarities depend on the choice of
model(s), architecture(s), and the training process (Leite and Brazdil 2005;
Zamir et al. 2018; Achille et al. 2019; Khodak, Balcan, and Talwalkar 2019;
Nguyen et al. 2020; Venkitaraman, Hansson, and Wahlberg 2020; Gao and
Chaudhari 2021). For example, Zamir et al. (2018); Venkitaraman, Hansson, and
Wahlberg (2020); Gao and Chaudhari (2021) use pre-trained task specified
models to measure a notion of forward transfer and define it as task
similarity. Achille et al. (2019) embed tasks into a vector space that relies
on a partially-trained network. Khodak, Balcan, and Talwalkar (2019) use the
optimal parameters as a proxy for each task and Leite and Brazdil (2005) use
the learning curves of a pre-specified model to measure task similarities.
Besides being model-dependent, these approaches are often computationally
expensive as they involve training deep models (or require pre-trained
models).

Model-agnostic task similarity measures provide a fundamentally different
approach to quantifying task relationships (Ben-David et al. 2006; Alvarez-
Melis and Fusi 2020; Tran, Nguyen, and Hassner 2019; Tan, Li, and Huang 2021).
These methods often measure the similarity between tasks as a function of the
similarity between the joint or conditional input/output distributions,
sometimes also taking the loss function into account. The classic theoretical
results for such similarity measures (Ben-David et al. 2006; Batu et al. 2000)
focus on information theoretic divergences between the source and target
distributions. More recently, Optimal Transport (OT) based approaches
(Alvarez-Melis and Fusi 2020; Tan, Li, and Huang 2021; Xu et al. 2022) have
shown promise in modeling task similarities. Notably, Alvarez-Melis and Fusi
(2020) approach measuring task similarities through the lens of a hierarchical
OT (Yurochkin et al. 2019) where they solve an inner OT problem to calculate
the label distance between the class-conditional distributions of two
supervised learning tasks. The label distance is then incorporated into the
transportation cost of an outer OT problem, resulting in a distance between
two datasets that integrates both sample and label discrepancies. Tan, Li, and
Huang (2021) treats the optimal transport plan between the input distributions
of two tasks as a joint probability distribution and use conditional entropy
to measure the difference between the two tasks. One major shortcoming of
these OT-based approaches is their computational complexity. These methods
require the pairwise calculation of OT (or entropy regularized OT) between
different tasks, which can be prohibitively expensive in applications
requiring frequent evaluations of task similarities, e.g., in continual
learning.

We propose a novel OT-based task embedding for supervised learning problems
that is model-agnostic and computationally efficient. On the one hand, our
proposed approach is similar to (Achille et al. 2019) and (Peng, Li, and
Saenko 2020), which embed datasets into a vector space in which one can easily
measure the difference between tasks, e.g., via the Euclidean distance between
embedded vectors. On the other hand, our approach is inspired by the Optimal
Transport Dataset Distance (OTDD) (Alvarez-Melis and Fusi 2020) framework, and
it essentially provides a Euclidean embedding for a hierarchical OT-based
distance between tasks. To calculate such a task embedding, we use the
Wasserstein embedding framework (Wang et al. 2013; Kolouri et al. 2020).
Importantly, our approach alleviates the need for pairwise calculation of OT
problems between tasks, turning it into a more desirable solution than
previously proposed methods.

Contributions. We propose a computationally efficient and model-agnostic task
embedding, denoted as Wasserstein Task Embedding (WTE), in which the Euclidean
distance between embedded vectors approximates a hierarchical OT distance
between the tasks. We provide extensive numerical experiments and demonstrate
that: 1) our calculated distances between embedded tasks are highly correlated
with the OTDD distance (Alvarez-Melis and Fusi 2020), 2) our proposed
embedding and similarity calculation is significantly faster than the OTDD
distance, and 3) our proposed similarity measure provides strong and
statistically significant correlation with both forward and backward transfer.

## Related work

Model-based task similarity. Most existing approaches to measuring task
similarity are model-dependent and use forward transferability as a proxy for
similarity. For example, Zamir et al. (2018) use pre-trained models on source
tasks and measure their performance on a target task to obtain an asymmetric
notion of similarity between source and target tasks. Following Zamir et al.
(2018)’s work, Dwivedi and Roig (2019) measure the transferability in a more
efficient manner by applying the Representation Similarity Analysis (RSA)
between the trained models (e.g. DNNs) from different tasks. Similarly, Nguyen
et al. (2020) assume the source and target tasks share the same set of inputs
but have different sets of labels, and estimate the transferability by the
empirical conditional distribution of target labels given the inputs computed
by a pre-trained model on the source task.

Another class of approaches embed the tasks into a vector space and then
define the (dis)similarity on the embedded vector representations. Achille et
al. (2019) discuss processing data (images) through a partially trained “probe
network” and obtain vector embedding by computing the Fisher information
matrix (FIM). The (dis)similarity of two tasks is then computed from the
difference between the the FIMs. Similarly, Peng, Li, and Saenko (2020)
propose a domain (labeled dataset) to vector technique. In particular, given a
domain, they feed the data to a pre-trained CNN to compute the Gram matrices
of the activations of the hidden convolutional layers, and apply feature
disentanglement to extract the domain-specified features. Concatenation of the
diagonal entries of Gram matrices and the domain-specified features gives the
final domain embedding. These methods, however, highly rely on the pre-trained
models and training process, and lack theoretical guarantees. On the opposite
side of the spectrum is directly measuring the discrepancy between domains.

Discrepancy measures of domains. Over the years, numerous notions of
discrepancy to measure the (dis)similarity of datasets (domains) were
proposed, including $L_{1}$-norm (Batu et al. 2000), generalized Kolmogorov-
Smirnov distance (Devroye, Györfi, and Lugosi 1996; Kifer, Ben-David, and
Gehrke 2004), and loss-oriented discrepancy distance (Mansour, Mohri, and
Rostamizadeh 2009). In the context of domain adaptation, generalized
Kolmogorov-Smirnov distance (later known as the $\mathcal{A}$-distance) is a
principled notation of discrepancy, which is a relaxation of total variation.
Ben-David et al. (2006) show that the target performance (generalization
error) is controlled by the empirical estimate of the source domain error and
the $\mathcal{A}$-distance between source and target domains. Another widely
used distance is the Maximum Mean Discrepancy (MMD) (Gretton et al. 2006),
which captures the (dis)similarity of the embedding of distribution measures
in a reproducing kernel Hilbert space. Pan et al. (2010) propose to learn
transfer components across domains in reproducing kernel Hilbert space using
MMD, and show that the subspace spanned by these transfer components preserves
data properties. Such domain discrepancy methods, however, can not take labels
into account, and thus may not be enough to reflect the similarity of tasks.

Optimal transport based task similarity. In recent years, metrics rooted in
the optimal transport problem, e.g., the “Wasserstein distance” (Villani 2009,
2021) (or the “earth mover’s distance” (Rubner, Tomasi, and Guibas 2000;
Solomon et al. 2014)), have attracted growing interest in the machine learning
community. Wasserstein distance is a rigorous metric of probability measures
endowed with desired statistical convergence behavior, in contrast to other
classical discrepancies (e.g. KL-divergence, total variation, JS-divergence,
Hellinger distance, Maximum mean discrepancy, etc). OT based metrics are
widely used in generative modeling (Arjovsky, Chintala, and Bottou 2017; Liu,
Gu, and Samaras 2019), domain adaptation (Courty et al. 2017; Courty, Flamary,
and Tuia 2014; Alvarez-Melis and Jaakkola 2018), graph embedding (Kolouri et
al. 2020; Xu et al. 2019), and neural architecture search (Kandasamy et al.
2018).

Alvarez-Melis and Fusi (2020) propose a notion of distance between two
datasets in a supervised learning setting. They introduce Optimal Transport
Dataset Distance (OTDD) based on the OT theory, which can be thought as a
hierarchical OT distance where the transportation cost measures the distance
between samples as well as labels. With the assumption that the label-induced
distributions can be approximated by Gaussians, the distance between labels is
defined as the Bures-Wasserstein distance.

Tan, Li, and Huang (2021) introduce another OT-based method to measure the
transferability, named OTCE (Optimal Transport Conditional Entropy) score. In
particular, they first use the entropic optimal transport to estimate domain
differences and then use the optimal coupling between the source and target
distributions to compute the conditional entropy of the target task given
source task. The OTCE is defined by the linear combination of the OT distance
and the conditional entropy. Both OTDD and OTCE were shown to be effectively
aligned with forward transfer, however, the computation of the pairwise
Wasserstein distances among increasing number of datasets remains expensive.
This hinders the application of these methods to problems where one needs to
perform nearest dataset retrieval frequently (e.g., memory replay approaches
in continual learning).

Computation Cost of OT Distance. Calculating the Wasserstein distance involves
solving an $n^{2}$ dimension linear programming and the computational cost is
$\mathcal{O}(n^{3}\log(n))$ for a pair of $n$-size empirical distributions
(Pele and Werman 2009). To facilitate the computation, one common method is
adding entropic regularization (Cuturi 2013; Peyré, Cuturi et al. 2017), by
which the original linear programming problem is converted into a strictly
convex problem. By applying the Sinkhorn-Knopp algorithm (Peyré, Cuturi et al.
2017; Chizat et al. 2018) to find an $\epsilon$\- accurate solution, the
computational complexity reduces to $\mathcal{O}(n^{2}\log(n)/\epsilon^{3})$
(Altschuler, Niles-Weed, and Rigollet 2017). However, this technique suffers a
stability-accuracy trade-off. When the regularity coefficient is high, the
objective is biased toward the entropy term; when it is small, the Sinkhorn
algorithm will not be numerically stable.

Figure 1: Wasserstein Task Embedding framework. Given tasks $\tau_{1}$ and
$\tau_{2}$ with input space $\mathbb{R}^{d}$, WTE first map them into
$\mathbb{R}^{d+L}$ as probability distributions $\mu_{1}$ and $\mu_{2}$ by
MDS, then apply WE to get vector $v_{1}$ and $v_{2}$ with respect to a fixed
reference measure $\mu_{0}$. Here $M$ is the size of reference set.

## Preliminaries

### Multidimensional scaling (MDS)

Multidimensional scaling (MDS) (Cox and Cox 2008) is a non-linear
dimensionality reduction approach that embeds $N$ samples into an
$l$-dimensional Euclidean space while preserving their pairwise distances.
Given a set of high-dimensional data $\mathcal{X}=\\{x_{n}\\}_{n=1}^{N}$ and
the proximity matrix $D\in\mathbb{R}^{N\times N}$, where
$D_{i,j}=d(x_{i},x_{j})$, and $d(\cdot,\cdot)$ denotes the metric in
$\mathcal{X}$, the goal of MDS is to construct a distance-preserving map from
$\mathcal{X}$ to a lower-dimensional Euclidean space $\mathbb{R}^{l}$.
Depending on the objective and inputs, MDS can be classified into metric MDS
and non-metric MDS. Specifically, metric MDS aims to find a map
$\psi:\mathcal{X}\rightarrow\mathbb{R}^{l}$ such that

$\displaystyle\min_{\psi}\sqrt{\frac{\sum_{i,j}\left(d(x_{i},x_{j})-\|\psi(x_{i}),\psi(x_{j})\|\right)^{2}}{\sum_{i,j}d(x_{i},x_{j})^{2}}},$
(1)

which can be solved by Algo. 1.

Algorithm 1 Multidimensional Scaling

procedure
MDS($\mathcal{X}=\\{x_{n}\\}_{n=1}^{N},~{}D=[d(x_{i},x_{j})]_{i,j},~{}l$)

$B=-\frac{1}{2}(id-\frac{1}{N}\mathds{1}_{N})D(id-\frac{1}{N}\mathds{1}_{N})$

Eigen-decomposition $B=V\Lambda V^{T}$

Rearrange $\Lambda$ into $\hat{\Lambda}$ with descending order of variances

Rearrange $V$ into $\hat{V}$ in correspondence with $\hat{\Lambda}$

$\hat{\Lambda}_{(l)}=\hat{\Lambda}[:l,:l]$; $\hat{V}_{(l)}=\hat{V}[:l]$ return
$\psi(\mathcal{X})=\hat{\Lambda}_{(l)}^{\frac{1}{2}}\hat{V}_{(l)}$

end procedure

Note that MDS not only works for Euclidean distances, but also for other
dissimilarities such as Wasserstein distances (Wang et al. 2011; Hamm,
Henscheid, and Kang 2022).

### Wasserstein Distances

Let $\mu,\nu$ be Borel probability measures on
$\mathcal{X}\subseteq\mathbb{R}^{d}$ with finite $p^{\text{th}}$ moment, and
the corresponding probability density functions are $p_{\mu}$ and $p_{\nu}$,
i.e. $d\mu=p_{\mu}dx$, $d\nu=p_{\nu}dx$. The 2-Wasserstein distance between
$\mu$ and $\nu$ is defined as (Villani 2009):

$\mathcal{W}_{2}(\mu,\nu)=\left(\inf_{\gamma\in\Gamma(\mu,\nu)}\int_{\mathcal{X}\times\mathcal{X}}\|x-x^{\prime}\|^{2}d\gamma(x,x^{\prime})\right)^{\frac{1}{2}},$
(2)

where $\Gamma(\mu,\nu)$ is the set of all transport plans between $\mu$ and
$\nu$, i.e. probability measures on $\mathcal{X}\times\mathcal{X}$ with
marginals $\mu$ and $\nu$. We also note that by Brenier theorem (Brenier
1991), given two absolutely continuous probability measures $\mu,\nu$ on
$\mathbb{R}^{d}$ with densities $p_{\mu},p_{\nu}$, there exists a convex
function $\phi$ such that $T=\nabla\phi$ is a transport map sending $\mu$ to
$\nu$. Moreover, it is the optimal map in the Monge-Kantorovitch optimal
transport problem with quadratic cost:

$\mathcal{W}_{2}(\mu,\nu)=\left(\int_{\mathcal{X}}\|x-T(x)\|^{2}p_{\mu}dx\right)^{\frac{1}{2}},$
(3)

where $T=\nabla\phi$ pushes $\mu$ to $\nu$, denoted by $T_{\\#}\mu=\nu$.

### Wasserstein Embedding (WE)

Wasserstein Embedding (Wang et al. 2013; Kolouri et al. 2016; Courty, Flamary,
and Ducoffe 2017; Kolouri et al. 2021) provides a Hilbertian embedding for
probability distributions such that the Euclidean distance between the
embedded vectors approximates the 2-Wasserstein distance between the two
distributions. Let $\\{\mu_{i}\\}_{i=0}^{I}$ be a set of a probability
distributions over $\mathcal{X}\subseteq\mathbb{R}^{d}$ with densities
$\\{p_{i}\\}_{i=0}^{I}$. We fix $\mu_{0}$ as the reference measure. Assume
$T_{i}$ is the optimal transport map that pushes $\mu_{0}$ to $\mu_{i}$, the
Wasserstein embedding of $\mu_{i}$ is through a function $\Phi$ defined as

$\Phi(\mu_{i})=(T_{i}-id)\sqrt{p_{0}}$ (4)

where the $id$ is the identity function, i.e., $id(x)=x$. $\Phi$ admits nice
properties including but not limited to (Kolouri et al. 2021):

  1. 1.

$d(\mu_{i},\mu_{j})\triangleq\|\Phi(\mu_{i})-\Phi(\mu_{j})\|_{2}$ is a true
metric between $\mu_{i}$ and $\mu_{j}$, moreover, it approximates the
2-Wasserstein distance:
$d(\mu_{i},\mu_{j})\approx\mathcal{W}_{2}(\mu_{i},\mu_{j})$.

  2. 2.

In particular,
$\|\Phi(\mu_{i})\|_{2}=\|\Phi(\mu_{i})-\Phi(\mu_{0})\|_{2}\approx\mathcal{W}_{2}(\mu_{i},\mu_{0})$.
Here we leveraged the fact $\Phi(\mu_{0})=0$.

Although these hold true for both continuous and discrete measures
$\\{\mu_{i}\\}_{i=0}^{I}$, we focus on the (uniformly distributed) discrete
setting in this paper and provide the following numerical computation details.
Let $p_{i}=\frac{1}{N_{i}}\sum_{n=1}^{N_{i}}\delta_{x_{n}^{i}}$, where
$\delta_{x}$ is the Dirac delta function centered at $x\in\mathcal{X}$ and
$X_{i}=\\{x_{n}^{i}\\}_{n=1}^{N_{i}}$ is the set of locations of non-negative
mass for $\mu_{i}$. Then the Kantorovich problem with quadratic cost between
$\mu_{i}$ and $\mu_{0}$ can be formulated as

$\min_{\pi\in\Pi_{i}}\sum_{n=1}^{N_{0}}\sum_{k=1}^{N_{i}}\pi_{nk}\|x_{n}^{i}-x_{k}^{0}\|_{2}^{2}$
(5)

where the feasible set is

$\Pi_{i}\triangleq\\{\pi\in\mathbb{R}^{N_{0}\times
N_{i}}|N_{0}\sum_{n=1}^{N_{i}}\pi_{nk}=N_{i}\sum_{k=1}^{N_{0}}\pi_{nk}=1,~{}\forall
n,k\\}.$ (6)

The optimal transport plan $\pi_{i}^{*}$ is the minimizer of the above
optimization problem, which is solved by linear program at cost
$\mathcal{O}(N^{3}\log(N))$, $N$ being the number of input samples. To avoid
mass splitting, the barycentric projection (Wang et al. 2013) assigns each
$x_{j}^{0}$ in the reference distribution to the center of mass it is sent to
and thus outputs an approximated Monge map $T_{i}$. Then the Wasserstein
Embedding for input $X_{i}$ is calculated by

$\Phi(X_{i})=(T_{i}-X_{0})/\sqrt{N_{0}}\in\mathbb{R}^{N_{0}\times d}.$ (7)

One of the motivations behind Wasserstein embedding is to reduce the need for
computing pairwise Wasserstein distances. Given $M$ datasets, computation of
$\frac{M(M-1)}{2}$ Wasserstein distances across all distinct pairs is
impractically expensive especially when $M$ is large, while leveraging
Wasserstein embedding, it suffices to calculate only $M$ Wasserstein distances
and the pairwise Euclidean distances between the embedded distributions.

## Method

In this section, we specify the problem setting, review the OTDD framework,
and then propose our Wasserstein task embedding (WTE).

### Problem Setting

In supervised classification problems, tasks are represented by input-label
pairs and can be denoted as
$\tau=\mathcal{X}\times\mathcal{Y}=\\{(x_{n},y_{n})\\}_{n=1}^{N}$, where
$\mathcal{X}\subseteq\mathbb{R}^{d}$ is the data/inputs and $\mathcal{Y}$ is
the labels. We aim to define a similarity/dissimilarity measure for tasks that
enables task clustering and allows for better control over the forward and
backward transfer.

Figure 2: Label-to-label Bures-Wasserstein distance (left) and label MDS
embedding Euclidean distances (middle) between MNIST and USPS datasets,
squared error is provided on the right.

### Optimal Transport Dataset Distance (OTDD)

Let $\mathcal{X}=\\{x_{n}\in\mathbb{R}^{d}\\}_{n=1}^{N}$ be the input set with
labels (classes) $\mathcal{Y}=\\{y_{j}\\}_{j=1}^{J}$. For each $j$, let
$C_{y_{j}}\subseteq\mathcal{X}$ denote the class with label $y_{j}$. Following
the OTDD framework, let
$\tau=\\{(x_{i},y_{j})\in\mathcal{X}\times\mathcal{Y}~{}|~{}x_{i}\in
C_{y_{j}}\\}_{i,j}$ denote the set of data-label pairs. OTDD encodes the label
$y_{j}$ as distribution $\nu_{y_{j}}$, where
$\nu_{y_{j}}=\frac{1}{|C_{y_{j}}|}\sum_{x_{i}\in C_{y_{j}}}\delta_{x_{i}}$.
The ground distance in $\tau$ is then defined by combining the Euclidean
distance between the data points and the 2-Wasserstein distance between label
distributions:

$\displaystyle
d_{\tau}\left((x,y),(x^{\prime},y^{\prime})\right):=\left(\|x-x^{\prime}\|^{2}+W_{2}^{2}(\nu_{y},\nu_{y^{\prime}})\right)^{\frac{1}{2}}.$
(8)

Based on this metric, the OT distance between two distributions $\mu_{i}$ and
$\mu_{j}$ on $\tau$ is

$\displaystyle
d_{OT}(\mu_{i},\mu_{j})=\inf_{\pi\in\Pi(\mu_{i},\mu_{j})}\int_{\tau\times\tau}d_{\tau}(z,z^{\prime})^{2}d\pi(z,z^{\prime}),$
(9)

where $\Pi(\mu_{i},\mu_{j})$ denotes the set of transport plans between
$\mu_{i}$ and $\mu_{j}$. Note that Eq. 9 is a hierarchical transport problem,
as the transportation cost itself depends on calculation of the Wasserstein
distance. To avoid the computational cost of a hierarchical optimal transport
problem, Alvarez-Melis and Fusi (2020) replace the Wasserstein distance in Eq.
8 with the Bures-Wasserstein distance (Malago, Montrucchio, and Pistone 2018;
Bhatia, Jain, and Lim 2019), which assumes that $\nu_{y}$s are Gaussian
distributions. Throughout the paper, we consider only the exact-OTDD, as
opposed to the entropy-regularized and other variants.

### Wasserstein Task Embedding

We define a task-2-vec framework (Fig. 1) using Wasserstein embedding (WE)
such that the (squared) Euclidean distance between two vectors approximates
the OTDD between the original tasks, and denote this embedding by WTE. We
later show in the experiment section that the Euclidean distance between the
embedded task vectors is not only highly predictive of forward
transferability, but also significantly correlates with the backward
transferability (catastrophic forgetting).

Label Embedding via MDS. The combination of optimal transport metric with MDS
technique was first introduced as an approach to characterize and contrast the
distribution of nuclear structure in different tissue classes (normal, benign,
cancer, etc.) (Wang et al. 2011), and further studied in image manifold
learning (Hamm, Henscheid, and Kang 2022). In short, it seeks to isometrically
map probability distributions to vectors in relatively low-dimensional space.
We leverage the prior work and define an approximated isometry on the label
distributions by 1) calculating the pairwise Wasserstein distances and 2)
applying MDS to obtain embedded vectors. We adopt the same simplification as
in OTDD, that is, assuming the label distributions are Gaussians to replace
Wasserstein distances with the closed form Bures-Wasserstein distance:

$\mathcal{W}_{2}^{2}(\nu_{y},\nu_{y^{\prime}})=\|u_{y}-u_{y^{\prime}}\|_{2}^{2}+\text{Tr}(\Sigma_{y}+\Sigma_{y^{\prime}}-2(\Sigma_{y}^{\frac{1}{2}}\Sigma_{y^{\prime}}\Sigma_{y}^{\frac{1}{2}})^{\frac{1}{2}})$
(10)

where $u$ and $\Sigma$ denote the mean and covariance matrix of Gaussian
distributions. In consistence with previous notations, let us denote the label
(MDS) embedding operator by $\psi$, then

$\mathcal{W}_{2}^{2}(\nu_{y},\nu_{y^{\prime}})\approx\|\psi(\nu_{y})-\psi(\nu_{y^{\prime}})\|_{2}^{2},$
(11)

where $\psi(\nu_{y}),\psi(\nu_{y^{\prime}})\in\mathbb{R}^{l}$ are vectors
whose dimension $l$ is selected to balance the trade-off between accuracy and
computation cost (Tenenbaum, Silva, and Langford 2000). Having both inputs and
labels represented as vectors, we concatenate these two components and map the
data-label pairs $\tau$ to $\tau^{\prime}\subseteq\mathbb{R}^{d+l}$ such that

$\displaystyle d_{\tau}\left((x,y),(x^{\prime},y^{\prime})\right)$
$\displaystyle\approx(\|x-x^{\prime}\|_{2}^{2}+\|\psi(\nu_{y})-\psi(\nu_{y^{\prime}})\|_{2}^{2})^{\frac{1}{2}}$
$\displaystyle=(\|[x,\psi(\nu_{y})]-[x^{\prime},\psi(\nu_{y^{\prime}})]\|_{2}^{2})^{\frac{1}{2}}$
$\displaystyle=\|[x,\psi(\nu_{y})]-[x^{\prime},\psi(\nu_{y^{\prime}})]\|_{2}$
(12)

where $[\cdot,\cdot]$ denotes the concatenation operator and the domain
$\tau^{\prime}=\\{[x,\psi(\nu_{y})]\\}_{x\in\mathcal{X}}$ is equipped with
$l_{2}$ norm. Fig. 2 shows this approximation performance among labels in
MNIST (LeCun and Cortes 2005) and USPS (Hull 1994) datasets. MDS embeddings
can capture the pairwise relationships with a maximum of $7.26\%$ error by
$10$-dimensional vectors.

Figure 3: (Top row) forward transfer error drop and (bottom row) catastrophic
forgetting against WTE distance on *NIST (left), Split-CIFAR100 (middle) and
Split-Tiny ImageNet (right) over five runs. Pearson’s $r$ and the
corresponding $p$-value are reported on top of each experiment setting.

Wasserstein Embedding. By Eq. Wasserstein Task Embedding and Eq. 9, OTDD can
be approximated by the squared 2-Wasserstein distance between the
distributions over input-(label MDS embedding) pairs, $\tau^{\prime}$. Then we
leverage the Wasserstein embedding framework to embed the updated task
distributions into a Hilbert space, with the goal of reducing the cost of
computing pairwise Wasserstein distances. Again, we emphasize that this can
bring down the cost from quadratic to linear with the number of task
distributions.

The WTE algorithm is summarized in Algo. 2. The outputs are the vector
representations of input tasks with respect to a pre-determined MDS dimension
and WE reference distribution.

Algorithm 2 Wasserstein Task Embedding

procedure
WTE($\\{X_{i}=\\{(x^{i}_{n},y^{i}_{n})\\}_{n=1}^{N_{i}}\\}_{i=1}^{I}$)

Calculate label-to-label distance matrix W (Eq. 10)

Calculate $\psi(\nu_{y})$ for all distinct labels $y$ (Algo. 1)

Stack each input with its label vector: $x\rightarrow[x,\psi(\nu_{y})]$

Calculate the WE
$v_{i}=\Phi(\\{[x^{i}_{n},\psi(\nu_{y^{i}_{n}})]\\}_{n=1}^{N_{i}})$ (Eq. 7)
return $\\{v_{i}\\}_{i=1}^{I}$

end procedure

## Experiments

To assess the effectiveness of our WTE framework, we empirically validate the
correlation between WTE distance and forward/backward transferability on
several datasets. Moreover, we provide both qualitative and quantitative
comparison results with OTDD, and show WTE distance is well aligned with OTDD,
and meanwhile is notably faster to compute. We use the MDS toolkit in scikit-
learn and the exact linear programming solver in Python Optimal Transport
(POT) (Flamary et al. 2021) library for implementing WE. We carry out the
distance calculations on CPU and all the model training experiments on a 24GB
NVIDIA RTX A5000 GPU.

Figure 4: Pairwise OTDD (left) and WTE distances (middle) on the *NIST task
group, and their correlation diagram (right). Notice that OTDD (Eq. 9) is the
squared $\mathcal{W}_{2}$, we report the squared WTE distances and adjust to
the same scale according to the cost function. Adjusted WTE distance is
strongly correlated with OTDD, with correlation coefficient $r=0.91$.

### Datasets

We conduct experiments on the following three task groups:

*NIST task group consists of the handwritten digits dataset MNIST (LeCun and Cortes 2005) and its extensions EMNIST (Cohen et al. 2017), FashionMNIST (Xiao, Rasul, and Vollgraf 2017), KMNIST (Clanuwat et al. 2018) along with USPS (Hull 1994). We choose the mnist split for EMNIST dataset and thus all tasks contain 10 classes of gray-scale images. All datasets have a training set of 60,000 samples and a test set of 10,000 samples, except USPS, with a total of 9,298 samples. We resize the images from USPS into $28\times 28$ pixel level to match with the others.

Split-CIFAR100 task group is generated by randomly splitting the CIFAR-100
(Krizhevsky 2009) dataset with 100 image categories into 10 smaller tasks,
each of which is a classification with 10 classes. There are 600 $32\times 32$
color images in the training set and 100 in the test set per class.

Split-Tiny ImageNet task group follows the same splitting scheme as in Split-
CIFAR100. We randomly divide the Tiny ImageNet (Le and Yang 2015) into 10
disjoint tasks with 20 classes. Each class contains 500 training images, 50
validating images and 50 test images. For better model performance, we first
rescale each sample to $256\times 256$ and then perform a center crop to get
$224\times 224$ pixel images.

### Results

To study the transfer behaviors against the WTE distances, we fix a model
architecture for each task group. Specifically, we use ResNet18 (He et al.
2016) on both *NIST and Split-Tiny ImageNet, and ResNet34 (He et al. 2016) on
Split-CIFAR100. In the forward transfer setting, for each source-target task
pair, we first train the head (i.e., the classifier) of a randomly initialized
backbone on the target task, and use the test performance as the baseline.
Next, we adapt from a model pre-trained on the source task and finetune the
head on the target task. We define the forward transferability of the source-
target pair as the performance gain, i.e. error drop when adapting from the
source task. To analyse backward transfer, all source tasks are trained
jointly during the first phase to avoid task bias, then in the second phase
the model learns only the target task and suffers from “forgetting” the
previous tasks. We use the catastrophic forgetting, i.e., negative backward
transfer as a measure of backward (in)transferability. In implementations of
WE, the reference distribution is fixed for each task group, and is randomly
generated by upsampling random images at a lower spatial resolution to entail
some smooth structure.

Figure 5: Wall-clock computation time comparison on the *NIST (left) and
Split-CIFAR100 (right) task groups.

Fig. 3 summarizes the correlation diagrams between our proposed WTE distance
and the forward/backward transferability on the aforementioned three task
groups. WTE distance is negatively correlated with forward transferability,
and positively correlated with catastrophic forgetting. In all scenarios, the
correlation is strong and statistically significant, which confirms the
efficacy of WTE distance as a measure of task similarities. We also visualize
the comparison between WTE distance and OTDD on the *NIST task group in Fig.
4, showing strong correlation between the two distances.

### Computation Complexity

As we mentioned before, OTDD suffers from a prohibitive computational cost as
the number of tasks grows large. The pairwise OTDD calculation for a set of
$M$ tasks requires $\mathcal{O}(M^{2}N^{3}\log(N))$ time in the worst case,
where $N$ is the largest number of samples among the tasks. WTE distance
requires solving only $M$ optimal transport problem, leading to
$\mathcal{O}(MN^{3}\log(N))$ complexity. To better demonstrate the efficiency
of WTE distance, we report the wall-clock time comparison on the *NIST and
Split-CIFAR100 in Fig. 5.

## Conclusion

In this paper, we propose Wasserstein task embedding (WTE), a model-agnostic
task embedding framework for measuring task (dis)similarities in supervised
classification problems. We perform a label embedding through multi-
dimensional scaling and leverage the 2-Wasserstein embedding framework to
embed tasks into a vector space, in which the Euclidean distance between the
embedded points approximates the 2-Wasserstein distance between tasks. We
demonstrate that our proposed task embedding distance is correlated with
forward and backward transfer on *NIST, Split-CIFAR100 and Split-Tiny ImageNet
task groups while being significantly faster than existing methods. In
particular, we show statistically significant negative correlation between the
WTE distances and the forward transfer, and positive correlation with the
catastrophic forgetting (i.e. negative backward transfer). Lastly, we show the
alignment of WTE distance with OTDD, but with a significant computational
advantage as the number of tasks grows.

## References

  * Achille et al. (2019) Achille, A.; Lam, M.; Tewari, R.; Ravichandran, A.; Maji, S.; Fowlkes, C. C.; Soatto, S.; and Perona, P. 2019.  Task2vec: Task embedding for meta-learning.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , 6430–6439.
  * Altschuler, Niles-Weed, and Rigollet (2017) Altschuler, J.; Niles-Weed, J.; and Rigollet, P. 2017.  Near-linear time approximation algorithms for optimal transport via Sinkhorn iteration.  _Advances in neural information processing systems_ , 30.
  * Alvarez-Melis and Fusi (2020) Alvarez-Melis, D.; and Fusi, N. 2020.  Geometric dataset distances via optimal transport.  _Advances in Neural Information Processing Systems_ , 33: 21428–21439.
  * Alvarez-Melis and Jaakkola (2018) Alvarez-Melis, D.; and Jaakkola, T. S. 2018.  Gromov-Wasserstein alignment of word embedding spaces.  _arXiv preprint arXiv:1809.00013_.
  * Arjovsky, Chintala, and Bottou (2017) Arjovsky, M.; Chintala, S.; and Bottou, L. 2017.  Wasserstein generative adversarial networks.  In _International conference on machine learning_ , 214–223. PMLR.
  * Bao et al. (2019) Bao, Y.; Li, Y.; Huang, S.-L.; Zhang, L.; Zheng, L.; Zamir, A.; and Guibas, L. 2019\.  An information-theoretic approach to transferability in task transfer learning.  In _2019 IEEE International Conference on Image Processing (ICIP)_ , 2309–2313. IEEE.
  * Batu et al. (2000) Batu, T.; Fortnow, L.; Rubinfeld, R.; Smith, W. D.; and White, P. 2000.  Testing that distributions are close.  In _Proceedings 41st Annual Symposium on Foundations of Computer Science_ , 259–269. IEEE.
  * Ben-David et al. (2006) Ben-David, S.; Blitzer, J.; Crammer, K.; and Pereira, F. 2006.  Analysis of representations for domain adaptation.  _Advances in neural information processing systems_ , 19.
  * Bhatia, Jain, and Lim (2019) Bhatia, R.; Jain, T.; and Lim, Y. 2019.  On the Bures–Wasserstein distance between positive definite matrices.  _Expositiones Mathematicae_ , 37(2): 165–191.
  * Bhattacharjee et al. (2020) Bhattacharjee, B.; Kender, J. R.; Hill, M.; Dube, P.; Huo, S.; Glass, M. R.; Belgodere, B.; Pankanti, S.; Codella, N.; and Watson, P. 2020.  P2L: Predicting transfer learning for images and semantic relations.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops_ , 760–761.
  * Brenier (1991) Brenier, Y. 1991.  Polar Factorization and Monotone Rearrangement of Vector-Valued Functions.  _Communications on Pure and Applied Mathematics_ , 44: 375–417.
  * Chizat et al. (2018) Chizat, L.; Peyré, G.; Schmitzer, B.; and Vialard, F.-X. 2018.  Scaling algorithms for unbalanced optimal transport problems.  _Mathematics of Computation_ , 87(314): 2563–2609.
  * Clanuwat et al. (2018) Clanuwat, T.; Bober-Irizar, M.; Kitamoto, A.; Lamb, A.; Yamamoto, K.; and Ha, D. 2018.  Deep Learning for Classical Japanese Literature.  _ArXiv_ , abs/1812.01718.
  * Cohen et al. (2017) Cohen, G.; Afshar, S.; Tapson, J.; and Van Schaik, A. 2017.  EMNIST: Extending MNIST to handwritten letters.  In _2017 international joint conference on neural networks (IJCNN)_ , 2921–2926. IEEE.
  * Courty, Flamary, and Ducoffe (2017) Courty, N.; Flamary, R.; and Ducoffe, M. 2017.  Learning wasserstein embeddings.  _arXiv preprint arXiv:1710.07457_.
  * Courty et al. (2017) Courty, N.; Flamary, R.; Habrard, A.; and Rakotomamonjy, A. 2017.  Joint distribution optimal transportation for domain adaptation.  _Advances in Neural Information Processing Systems_ , 30.
  * Courty, Flamary, and Tuia (2014) Courty, N.; Flamary, R.; and Tuia, D. 2014.  Domain adaptation with regularized optimal transport.  In _Joint European Conference on Machine Learning and Knowledge Discovery in Databases_ , 274–289. Springer.
  * Cox and Cox (2008) Cox, M. A.; and Cox, T. F. 2008.  Multidimensional scaling.  In _Handbook of data visualization_ , 315–347. Springer.
  * Cuturi (2013) Cuturi, M. 2013.  Sinkhorn distances: Lightspeed computation of optimal transport.  _Advances in neural information processing systems_ , 26.
  * Devroye, Györfi, and Lugosi (1996) Devroye, L.; Györfi, L.; and Lugosi, G. 1996.  Parametric Classification.  In _A Probabilistic Theory of Pattern Recognition_ , 263–278. Springer.
  * Dwivedi and Roig (2019) Dwivedi, K.; and Roig, G. 2019.  Representation similarity analysis for efficient task taxonomy & transfer learning.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 12387–12396.
  * Fifty et al. (2021) Fifty, C.; Amid, E.; Zhao, Z.; Yu, T.; Anil, R.; and Finn, C. 2021.  Efficiently identifying task groupings for multi-task learning.  _Advances in Neural Information Processing Systems_ , 34: 27503–27516.
  * Flamary et al. (2021) Flamary, R.; Courty, N.; Gramfort, A.; Alaya, M. Z.; Boisbunon, A.; Chambon, S.; Chapel, L.; Corenflos, A.; Fatras, K.; Fournier, N.; Gautheron, L.; Gayraud, N. T.; Janati, H.; Rakotomamonjy, A.; Redko, I.; Rolet, A.; Schutz, A.; Seguy, V.; Sutherland, D. J.; Tavenard, R.; Tong, A.; and Vayer, T. 2021.  POT: Python Optimal Transport.  _Journal of Machine Learning Research_ , 22(78): 1–8.
  * Gao and Chaudhari (2021) Gao, Y.; and Chaudhari, P. 2021.  An information-geometric distance on the space of tasks.  In _International Conference on Machine Learning_ , 3553–3563. PMLR.
  * Gretton et al. (2006) Gretton, A.; Borgwardt, K.; Rasch, M.; Schölkopf, B.; and Smola, A. 2006.  A kernel method for the two-sample-problem.  _Advances in neural information processing systems_ , 19.
  * Hamm, Henscheid, and Kang (2022) Hamm, K.; Henscheid, N.; and Kang, S. 2022.  Wassmap: Wasserstein Isometric Mapping for Image Manifold Learning.  _arXiv preprint arXiv:2204.06645_.
  * He et al. (2016) He, K.; Zhang, X.; Ren, S.; and Sun, J. 2016.  Deep residual learning for image recognition.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 770–778.
  * Hull (1994) Hull, J. 1994.  A database for handwritten text recognition research.  _IEEE Transactions on Pattern Analysis and Machine Intelligence_ , 16(5): 550–554.
  * Kandasamy et al. (2018) Kandasamy, K.; Neiswanger, W.; Schneider, J.; Poczos, B.; and Xing, E. P. 2018.  Neural architecture search with bayesian optimisation and optimal transport.  _Advances in neural information processing systems_ , 31.
  * Khodak, Balcan, and Talwalkar (2019) Khodak, M.; Balcan, M.-F. F.; and Talwalkar, A. S. 2019.  Adaptive gradient-based meta-learning methods.  _Advances in Neural Information Processing Systems_ , 32.
  * Kifer, Ben-David, and Gehrke (2004) Kifer, D.; Ben-David, S.; and Gehrke, J. 2004.  Detecting change in data streams.  In _VLDB_ , volume 4, 180–191. Toronto, Canada.
  * Kolouri et al. (2020) Kolouri, S.; Naderializadeh, N.; Rohde, G. K.; and Hoffmann, H. 2020.  Wasserstein embedding for graph learning.  _arXiv preprint arXiv:2006.09430_.
  * Kolouri et al. (2021) Kolouri, S.; Naderializadeh, N.; Rohde, G. K.; and Hoffmann, H. 2021.  Wasserstein Embedding for Graph Learning.  In _International Conference on Learning Representations_.
  * Kolouri et al. (2016) Kolouri, S.; Tosun, A. B.; Ozolek, J. A.; and Rohde, G. K. 2016.  A continuous linear optimal transport approach for pattern analysis in image datasets.  _Pattern recognition_ , 51: 453–462.
  * Krizhevsky (2009) Krizhevsky, A. 2009.  Learning Multiple Layers of Features from Tiny Images.
  * Le and Yang (2015) Le, Y.; and Yang, X. S. 2015.  Tiny ImageNet Visual Recognition Challenge.
  * LeCun and Cortes (2005) LeCun, Y.; and Cortes, C. 2005.  The mnist database of handwritten digits.
  * Leite and Brazdil (2005) Leite, R.; and Brazdil, P. 2005.  Predicting relative performance of classifiers from samples.  In _Proceedings of the 22nd international conference on machine learning_ , 497–503.
  * Liu, Gu, and Samaras (2019) Liu, H.; Gu, X.; and Samaras, D. 2019.  Wasserstein gan with quadratic transport cost.  In _Proceedings of the IEEE/CVF international conference on computer vision_ , 4832–4841.
  * Malago, Montrucchio, and Pistone (2018) Malago, L.; Montrucchio, L.; and Pistone, G. 2018.  Wasserstein riemannian geometry of positive definite matrices.  _arXiv preprint arXiv:1801.09269_.
  * Mansour, Mohri, and Rostamizadeh (2009) Mansour, Y.; Mohri, M.; and Rostamizadeh, A. 2009.  Domain adaptation: Learning bounds and algorithms.  _arXiv preprint arXiv:0902.3430_.
  * Nguyen et al. (2020) Nguyen, C.; Hassner, T.; Seeger, M.; and Archambeau, C. 2020.  Leep: A new measure to evaluate transferability of learned representations.  In _International Conference on Machine Learning_ , 7294–7305. PMLR.
  * Nguyen et al. (2019) Nguyen, C. V.; Achille, A.; Lam, M.; Hassner, T.; Mahadevan, V.; and Soatto, S. 2019\.  Toward understanding catastrophic forgetting in continual learning.  _arXiv preprint arXiv:1908.01091_.
  * Pan et al. (2010) Pan, S. J.; Tsang, I. W.; Kwok, J. T.; and Yang, Q. 2010.  Domain adaptation via transfer component analysis.  _IEEE transactions on neural networks_ , 22(2): 199–210.
  * Pele and Werman (2009) Pele, O.; and Werman, M. 2009.  Fast and robust earth mover’s distances.  In _2009 IEEE 12th international conference on computer vision_ , 460–467. IEEE.
  * Peng, Li, and Saenko (2020) Peng, X.; Li, Y.; and Saenko, K. 2020.  Domain2vec: Domain embedding for unsupervised domain adaptation.  In _European conference on computer vision_ , 756–774. Springer.
  * Peyré, Cuturi et al. (2017) Peyré, G.; Cuturi, M.; et al. 2017.  Computational optimal transport.  _Center for Research in Economics and Statistics Working Papers_ , (2017-86).
  * Rubner, Tomasi, and Guibas (2000) Rubner, Y.; Tomasi, C.; and Guibas, L. J. 2000.  The earth mover’s distance as a metric for image retrieval.  _International journal of computer vision_ , 40(2): 99–121.
  * Solomon et al. (2014) Solomon, J.; Rustamov, R.; Guibas, L.; and Butscher, A. 2014.  Earth mover’s distances on discrete surfaces.  _ACM Transactions on Graphics (ToG)_ , 33(4): 1–12.
  * Tan, Li, and Huang (2021) Tan, Y.; Li, Y.; and Huang, S.-L. 2021.  qq.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 15779–15788.
  * Tenenbaum, Silva, and Langford (2000) Tenenbaum, J. B.; Silva, V. d.; and Langford, J. C. 2000.  A global geometric framework for nonlinear dimensionality reduction.  _science_ , 290(5500): 2319–2323.
  * Tran, Nguyen, and Hassner (2019) Tran, A. T.; Nguyen, C. V.; and Hassner, T. 2019.  Transferability and hardness of supervised classification tasks.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , 1395–1405.
  * Venkitaraman, Hansson, and Wahlberg (2020) Venkitaraman, A.; Hansson, A.; and Wahlberg, B. 2020.  Task-similarity aware meta-learning through nonparametric kernel regression.  _arXiv preprint arXiv:2006.07212_.
  * Villani (2009) Villani, C. 2009.  _Optimal transport: old and new_ , volume 338.  Springer.
  * Villani (2021) Villani, C. 2021.  _Topics in optimal transportation_ , volume 58.  American Mathematical Soc.
  * Wang et al. (2011) Wang, W.; Ozolek, J. A.; Slepčev, D.; Lee, A. B.; Chen, C.; and Rohde, G. K. 2011\.  An Optimal Transportation Approach for Nuclear Structure-Based Pathology.  _IEEE Transactions on Medical Imaging_ , 30(3): 621–631.
  * Wang et al. (2013) Wang, W.; Slepčev, D.; Basu, S.; Ozolek, J. A.; and Rohde, G. K. 2013.  A linear optimal transportation framework for quantifying and visualizing variations in sets of images.  _International journal of computer vision_ , 101(2): 254–269.
  * Xiao, Rasul, and Vollgraf (2017) Xiao, H.; Rasul, K.; and Vollgraf, R. 2017.  Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms.  _arXiv preprint arXiv:1708.07747_.
  * Xu et al. (2019) Xu, H.; Luo, D.; Zha, H.; and Duke, L. C. 2019.  Gromov-wasserstein learning for graph matching and node embedding.  In _International conference on machine learning_ , 6932–6941. PMLR.
  * Xu et al. (2022) Xu, R.; Yang, X.; Liu, B.; Zhang, K.; and Liu, W. 2022.  Selecting task with optimal transport self-supervised learning for few-shot classification.  _arXiv preprint arXiv:2204.00289_.
  * Yurochkin et al. (2019) Yurochkin, M.; Claici, S.; Chien, E.; Mirzazadeh, F.; and Solomon, J. M. 2019.  Hierarchical optimal transport for document representation.  _Advances in Neural Information Processing Systems_ , 32.
  * Zamir et al. (2018) Zamir, A. R.; Sax, A.; Shen, W.; Guibas, L. J.; Malik, J.; and Savarese, S. 2018\.  Taskonomy: Disentangling task transfer learning.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 3712–3722.