Making transport more robust and interpretable by moving data through a small number of anchor points

Abstract

Optimal transport (OT) is a widely used technique for distribution alignment, with applications throughout the machine learning, graphics, and vision communities. Without any additional structural assumptions on transport, however, OT can be fragile to outliers or noise, especially in high dimensions. Here, we introduce Latent Optimal Transport (LOT), a new approach for OT that simultaneously learns low-dimensional structure in data while leveraging this structure to solve the alignment task. The idea behind our approach is to learn two sets of “anchors” that constrain the flow of transport between a source and target distribution. In both theoretical and empirical studies, we show that LOT regularizes the rank of transport and makes it more robust to outliers and the sampling density. We show that by allowing the source and target to have different anchors, and using LOT to align the latent spaces between anchors, the resulting transport plan has better structural interpretability and highlights connections between both the individual data points and the local geometry of the datasets.

1. Introduction

Optimal transport (OT) (Villani, 2008) is a widely used technique for distribution alignment that learns a transport plan which moves mass from one distribution to match another. With recent advances in tools for regularizing and speeding up OT (Cuturi, 2013), this approach has found applications in many diverse areas of machine learning, including domain adaptation (Courty et al.; Courty et al., 2017), generative modeling (Martin Arjovsky & Bottou, 2017; Tolstikhin et al., 2017), document retrieval (Kusner et al., 2015), computer graphics (Solomon et al., 2014; 2015; Bonneel et al., 2016), and computational neuroscience (Gramfort et al., 2015; Lee et al., 2019).

While the ground metric in OT can be used to impose geometric structure into transport, without any additional assumptions, OT can be fragile to outliers or noise, especially in high dimensions. To overcome this issue, additional structure, either in the data or in the transport plan, can be used to improve alignment or make transport more robust. Examples of methods that incorporate additional structure into OT include approaches that leverage hierarchical structure or cluster consistency (Lee et al., 2019; Yurochkin et al., 2019; Xu et al., 2020), partial class information (Courty et al.,2017; Courty et al.), submodular cost functions (Alvarez-Melis et al., 2018), and low-rank constraints on the transport plan (Forrow et al., 2019; Altschuler et al., 2019). Because of the difficulty of incorporating structure into OT, many of these methods need low-dimensional structure in data to be specified in advance (e.g., estimated clusters or labels).

To simultaneously learn low-dimensional structure and use it to constrain transport, Forrow et al. (2019) recently introduced a statistical approach for OT that builds a factorization of the transport plan to regularize its rank. After factorization, transport from a source to target distribution can be visualized as the flow of mass through a small number of anchors (hubs), which serve as relay stations through which transportation must pass (see Figure 1, a vs. b). Although this idea of moving data through anchors is appealing, in previous work, the anchors used to constrain transport are shared by the source and target. As a result, when the source and target contain different structures or experience domain shift (Courty et al.), shared anchors may not provide an adequate representation for both domains simultaneously.

Figure 1: Comparisons of transport plans obtained for different methods applied to clustered data after domain shift.

Here, we visualize the connection between the source (blue) x and its estimated target (red) $\widehat{y}={\text{argmax}}_{y}p(y\mid x)$. From left to right, we show the standard OT plan (a) and the factored coupling (FC) approach (b). To the right, we show the result of LOT when we use 4 anchors in the target with the same number in the target (c) and 8 anchors in the target (d).

In this work, we propose a new structured transport approach called Latent Optimal Transport (LOT). The main idea behind LOT is to factorize the transport plan into three components, where mass is moved: (i) from individual source points to source anchors, (ii) from the source anchors to target anchors, and (iii) from target anchors to individual target points (Figure 1c–d). The intermediate transport plan captures the high-level structural similarity between the source and target, while the outer transport plans cluster data in their respective spaces. In both theoretical and empirical studies, we show that LOT regularizes the rank of transport and has the effect of denoising the transport plan, making it more robust to outliers and sampling. By allowing the source and target to have different anchors and aligning the latent spaces of the anchors, we show that the mapping between datasets can be more easily interpreted.

Specifically, our contributions are as follows. (i) We introduce LOT, a new form of structured transport, and propose an efficient algorithm that solves our proposed objective (Section 3), (ii) Theoretically, we show that LOT can be interpreted as a relaxation to OT, and from a statistical point-of-view, it overcomes the curse of dimensionality in terms of the sampling rate (Section 5), (iii) We study the robustness of the approach to noise, sampling, and various data perturbations when applied to both synthetic data and domain adaptation problems in computer vision (Section 6).

2. Background

Optimal Transport:

Optimal transport (OT) (Villani, 2008; Santambrogio, 2015; Peyré et al., 2019) is a distribution alignment technique that learns a transport plan that specifies how to move mass from one distribution to match another. Specifically, consider two sets of data points encoded in matrices, the source X = [x₁, …, x_n] and the target Y = [y₁, …, y_m], where $x_i\in \mathcal{X}$, $y_j\in \mathcal{Y}$, ∀i, j. Assume they are endowed with discrete measures $\mu ={\sum }_{i=1}^{N}p\left(x_i\right){\delta }_{x_i}$, $\nu ={\sum }_{j=1}^{M}p\left(y_j\right){\delta }_{y_j}$, respectively. The cost of transporting x_i to y_j is c(x_i, y_j), where c denotes some cost function. OT considers the most cost-efficient transport by solving the following problem:¹

$${\text{OT}}_C(\mu ,\nu )≔\underset{P1=\mu ,P^{T}1=\nu }{\text{min}}\langle C,P\rangle ,$$ (1)

where P ≔ [p(x_i, y_j)]_{i, j} is the source-to-target transport plan matrix (coupling), and C = [c(x_i, y_j)]_{i, j} is the cost matrix. When c(x, y) = d(x, y)^p, where d is a distance function, ${\mathcal{W}}_p≔{\text{OT}}_{\text{C}}^{1/p}$ defines a distance called the p-Wasserstein distance. The objective in (1) is a linear programming problem, where computation speed can be prohibitive if n is large (Pele & Werman, 2009). A common speedup is to replace the objective by an entropy-regularized proxy,

$${\text{OT}}_{C,\epsilon }(\mu ,\nu )≔\underset{P1=\mu ,P^{T}1=\nu }{\text{min}}\langle C,P\rangle -\epsilon H(P)=\underset{P1=\mu ,P^{T}1=\nu }{\text{min}}\epsilon \text{KL}(P\Vert K),$$ (2)

where K is the Gibbs kernel induced by the element-wise exponential of the cost matrix K ≔ exp(−C/ε), H(P) ≔ −∑_ij P_ij log(P_{i, j}) is the Shannon entropy, and ε is a user-specified hyperparameter that controls the amount of entropic regularization that is introduced. We can alternatively write the objective function as a minimization of εKL(P∥K), where KL denotes the Kullback-Leibler divergence. In practice, the entropy-regularized form is often used over the original objective (1) as it admits a fast method called the Sinkhorn algorithm (Cuturi, 2013; Altschuler et al., 2017). Hence, we will use OT to refer to the entropy-regularized form unless specified otherwise in the context.

Optimal Transport via Factored Couplings:

Factored Coupling (FC) is proposed in (Forrow et al., 2019) to reduce the sample complexity of OT in high dimensions. Specifically, it adds an additional constraint to (1) by enforcing the transport plan to be of the following factored form,

$$p\left(x_i,y_j\right)=\sum _{l=1}^{k}p\left(z_l\right)p\left(x_i\mid z_l\right)p\left(y_j\mid z_l\right).$$ (3)

This has a nice interpretation: z_l serves as a common “anchor” that transportation from x_i to y_j must pass through. It turns out that FC is closely related to the Wasserstein barycenter problem (Agueh & Carlier, 2011; Cuturi & Doucet, 2014; Cuturi & Peyré, 2016), ${\text{min}}_{\nu }{\sum }_{i=1}^N{\mathcal{W}}_2^2\left({\mu }_i,\nu \right)$ where ν is the Procrustes mean to distributions μ_i, i = 1, …, N with respect to the squared 2-Wasserstein distance. A crucial insight from (Forrow et al., 2019) is that for N = 2, the barycenter ν could approximate the optimal anchors to a transport plan of the form (3) that minimizes the objective in (1).

3. Latent Optimal Transport

3.1. Motivation

Most datasets have low-dimensional latent structure, but OT does not naturally use it during transport. This motivates the idea that distribution alignment methods should both reveal the latent structure in the data in addition to aligning these latent structures. An illustrative example is provided in Figure 1; here, we show the transport plan for a source (red points) and a target (blue points), both of which exhibit clear cluster structures. Because OT transports points independently, the points can be easily mapped outside of their original cluster (a). In comparison, low-rank OTs (b-d) induce transport plans that are better at preserving clusters. In (b), because factored coupling (FC) transports points via common anchors (black squares), the anchors need to interpolate between both distributions, and it loses the freedom of choosing different structures for the source and target. On the other hand, by specifying different numbers of anchors for the source and target individually (c vs. d), LOT can extract different structures and output different transport plans.

3.2. Problem formulation

Consider data matrices X and Y and their measures μ, ν, as detailed in Section 2. We introduce “anchors” through which points must flow, thus constraining the transportation. The anchors are stacked in data matrices $Z_x≔\left[z_1^x,\dots ,z_{k_x}^x\right]$, $Z_y≔\left[z_1^y,\dots ,z_{k_y}^y\right]$. We denote the measures of the source and target anchors as ${\mu }_z={\sum }_{m=1}^{k_x}p\left(z_m^x\right){\delta }_{z_m^x}$ and $v_z={\sum }_{n=1}^{k_y}p\left(z_n^y\right){\delta }_{z_n^y}$. For any set $\mathcal{A}$, we further denote ${\Delta }_{\mathcal{A}}^k≔\left\{{\sum }_{i=1}^k{\omega }_i{\delta }_{a_i}:{\sum }_{i=1}^k{\omega }_i=1,{\omega }_i\ge 0,a_i\in \mathcal{A},\forall i\right\}$ as the set of probability measures on $\mathcal{A}$ that has discrete support of size up to k. Hence, ${\mu }_z\in {\Delta }_{{\mathcal{Z}}_x}^{k_x}$, ${\nu }_z\in {\Delta }_{{\mathcal{Z}}_y}^{k_y}$ where ${\mathcal{Z}}_x$ (resp. ${\mathcal{Z}}_y$) is the space of source (resp. target) anchors. If we interpret the conditional probability p(a|b) as the strength of transportation from b to a, then, using the chain rule, the concurrence probability p(x_i, y_j) of x_i and y_j can be written as,

$$p\left(x_i,y_j\right)=\sum _{m,n}p\left(x_i\right)p\left(z_m^x\mid x_i\right)p\left(z_n^y\mid z_m^x\right)p\left(y_j\mid z_n^y\right)=\sum _{m,n}p\left(x_i,z_m^x\right)\frac{p\left(z_m^x,z_n^y\right)}{p\left(z_m^x\right)p\left(z_n^y\right)}p\left(z_n^y,y_j\right).$$ (4)

When encoding these probabilities using a transport matrix P ≔ [p(x_i, y_j)]_i,j, the factorized form (4) can be written as,

$$P=P_x\text{diag}\left(u_z^{-1}\right)P_z\text{diag}\left(v_z^{-1}\right)P_y,$$ (5)

where P_x encodes transport from source space to source anchor space (i.e., $p\left(x_i,z_m^x\right)$), P_z encodes transport from source anchor space to target anchor space, P_y encodes transport from target anchor space to target space, and $u_z≔\left[p\left(z_1^x\right),\cdots ,p\left(z_{k_x}^x\right)\right]$, $v_z≔\left[p\left(z_1^y\right),\cdots ,p\left(z_{k_y}^y\right)\right]$ encode the latent distributions of anchors. To learn each of these transport plans, we must first designate the ground metric used to define the cost in each of the three stages. The cost matrices C_x, C_y determine how points will be transported to their respective anchors and thus dictate how the data structure will be extracted. We will elaborate on the choice of costs in Section 3.3.

We now formalize our proposed approach to transport in the following definition.

Definition 1.

Let C_x, C_y denote the cost matrices between the source/target and their representative anchors, and let C_z denote the cost matrix between anchors. We define the latent optimal transport (LOT) problem as,

$${\text{OT}}^{\text{L}}(\mu ,\nu )≔\underset{{\mu }_z\in {\Delta }_{{\mathcal{Z}}_x}^{k_x},{\nu }_z\in {\Delta }_{{\mathcal{Z}}_y}^{k_y}}{\text{inf}}\left\{{\text{OT}}_{C_x}\left(\mu ,{\mu }_{\text{Z}}\right)+{\text{OT}}_{C_{\text{z}}}\left({\mu }_{\text{z}},{\nu }_{\text{z}}\right)+{\text{OT}}_{C_{\text{y}}}\left({\nu }_{\text{z}},\nu \right)\right\},$$

where ${\mathcal{Z}}_x$ and ${\mathcal{Z}}_y$ are the latent spaces of the source and target anchors, respectively.²

The intuition behind Def. 1 is that we use ${\text{OT}}_{C_{\text{x}}}\left(\mu ,{\mu }_{\text{z}}\right)$ and ${\text{OT}}_{C_{\text{y}}}\left({\nu }_{\text{z}},\nu \right)$ to capture group structure in each space, and then ${\text{OT}}_{C_{\text{z}}}\left({\mu }_{\text{z}},{\nu }_{\text{z}}\right)$ to align the source and target by determining the transportation across anchors. Hence, LOT can be interpreted as an optimization of joint clustering and alignment. The flexibility of cost matrices allows LOT to capture different structures and induce different transport plans. In Section 5, we further show that LOT can be regarded as a relaxation of an OT problem.

Remark 1.

In Forrow et al. (2019), the authors introduce the notion of the transport rank for a transport plan P as the minimum number of product probability measures that its corresponding coupling can be composed from, i.e., $p(x,y)={\sum }_{i=1}^r{\lambda }_i\left(p_i(x)\otimes p_i(y)\right)$ λ_i ≥ 0, ∀i. In general, given a transportation plan P, the transport rank rank₊(P) is lower bounded by its usual matrix rank rank(P). In the case of LOT, the transport plan induced by Def. 1 satisfies rank(P) ≤ rank₊(P) ≤ min(k_x, k_y). Thus, by selecting a small number of anchors we naturally induce a low-rank solution for transport.

Next, we show some properties of LOT that highlight its similarity to a metric.

Proposition 1.

Suppose the latent spaces ${\mathcal{Z}}_x={\mathcal{Z}}_y$ are the same as the original data spaces $\mathcal{X}=\mathcal{Y}$, and the cost matrices are defined by C_x[a, b] = C_z[a, b] = C_y[a, b] = d(a, b)^p, where p ≥ 1 and d is some distance function. If we define the latent Wasserstein discrepancy as ${\mathcal{W}}_p^L≔{\left({\text{OT}}^{\text{L}}\right)}^{1/p}$, then there exist κ > 0 such that, for any μ, ν and ζ having latent distributions of support sizes up to k, the discrepancy satisfies,

• $${\mathcal{W}}_p^L(\mu ,\nu )\ge 0$$
• $${\mathcal{W}}_p^L(\mu ,\nu )={\mathcal{W}}_p^L(\nu ,\mu )$$
• $${\mathcal{W}}_p^L(\mu ,\nu )\le \kappa \left({\mathcal{W}}_p^L(\mu ,\zeta )+{\mathcal{W}}_p^L(\zeta ,\nu )\right)$$

The low-rank nature of LOT has a biasing effect that results in ${\mathcal{W}}_p^L(\mu ,\mu )>0$ for a general μ. We can debias it by defining its variant ${\tilde{\mathcal{W}}}_p^L(\mu ,\nu )≔{\left({\left({\mathcal{W}}_p^L(\mu ,\nu )\right)}^p-\underset{z_k\in {\Phi }_x}{\text{min}}{\mathcal{W}}_p^p\left(\mu ,z_k\right)-\underset{z_k^{\prime }\in {\Phi }_y}{\text{min}}{\mathcal{W}}_p^p\left(\nu ,z_k^{\prime }\right)\right)}^{1/p}$, where ${\Phi }_x={\Delta }_{{\mathcal{Z}}_x}^{k_x}$, ${\Phi }_y={\Delta }_{{\mathcal{Z}}_y}^{k_y}$. The following property connects ${\tilde{\mathcal{W}}}_p^L(\mu ,\nu )$ to k-means clustering.

Corollary 1.

Under the assumptions of Proposition 1, if p = 2 and k_x = k_y = k, then ∀μ, ν, we have ${\tilde{\mathcal{W}}}_2^L(\mu ,\nu )\ge 0$. Furthermore, ${\tilde{\mathcal{W}}}_2^L(\mu ,\nu )>0$ if their k-means centroids or sizes of their k-means clusters differ.

3.3. Establishing a ground metric

In what follows, we will focus on the Euclidean space $\mathcal{X}=\mathcal{Y}={ℝ}^d$. Instead of considering every source-to-target distance to build our transportation cost, we can use anchors as proxies for each point. A well-established way of encoding the distance that each point needs to travel to get to its nearest anchor, is to define the cost as:

$$C_x=d_{M_x},C_z=d_{M_z},C_y=d_{M_y},$$ (6)

where d_M denotes the Mahalanobis distance: $d_{\text{M}}^2(x,y)≔{(\text{x}-\text{y})}^{T}M(\text{x}-y)$ and M is some positive semidefinite matrix. The Mahalanobis distance generalizes the squared Euclidean distance and allows us to consider different costs based on correlations between features. The framework of Mahalanobis distance benefits from efficient metric learning techniques (Cuturi & Avis, 2014); recent research also establishes connections between it and robust OT (Paty & Cuturi, 2019; Dhouib et al., 2020). When a simple L2-distance is used (M = I), we will denote this specific variant as LOT-L2.

When LOT moves source points through anchors, the anchors impose a type of bottleneck, and this results in a loss of information that makes it difficult to estimate the corresponding point in the target space. In cases where accurate point-to-point alignment is desired, we propose an alternative strategy for defining the cost matrix C_z. The idea is to represent an anchor as the distribution of points assigned to it. Specifically, we represent z^x, z^y as measures in ${ℝ}^d:{\tilde{z}}^x={\sum }_{i=1}^N{P}_x\left(x_i\mid z^x\right){\delta }_{x_i}$, ${\tilde{z}}^y={\sum }_{j=1}^M{P}_y\left(y_j\mid z^y\right){\delta }_{y_j}$. Then we measure the cost between anchors as the squared Wasserstein distance between their respective distributions,

$$C_z≔{\left[{\mathcal{W}}_2^2\left(P_x\left(\cdot \mid z_m^x\right),P_y\left(\cdot \mid z_n^y\right)\right)\right]}_{m,n}.$$ (7)

Besides the quantity itself, the transport plan returned by calculating C_z is also very important as it provides accurate point-to-point maps. Since the cost matrix is now a function of P_x and P_y, we use an additional alternating scheme to solve the problem: we alternate between updating C_z while keeping P_x and P_y fixed, and then updating P_x, P_y, P_z while keeping C_z fixed. An efficient algorithm is presented in Appendix B.3 to reduce the computation complexity. This variant, LOT-WA, can yield better performance in downstream tasks that require precise alignment at the cost of additional computation.

3.4. Algorithm

In the rest of this section, we will develop our main approach for solving the problem in Def. 1. We provide an outline of the algorithm in Algorithm 1 and an implementation of the algorithm in Python at: http://nerdslab.github.io/latentOT.

(1). Optimizing P_x, P_y and P_z:

To begin, we assume that the anchors and cost matrices C_x, C_z, C_y are already specified. Let K_x, K_z, K_y be the Gibbs kernels induced from the cost matrices C_x, C_z, C_y as in (2). The optimization problem can be written as,

$$\underset{u_z,v_z,P_x,P_z,P_y}{\text{min}}\sum _{i\in \{x,y,z\}}{\epsilon }_i\text{KL}\left(P_i\Vert K_i\right),$$

$$\text{subject to: }{P}_{x}1=\mu ,P_x^{T}1=u_z,P_{z}1=u_z,P_z^{T}1=v_z,P_{y}1=v_z,P_y^{T}1=\nu .$$ (8)

This is a Bregman projection problem with affine constraints. An iterative projection procedure can thus be applied to solve the problem (Benamou et al., 2015). We present the procedure as UpdatePlan in Algorithm 1, where P_x, P_z, P_y are successively projected onto the constrained sets of fixed marginal distributions. We defer the detailed derivation to Appendix B.1.

(2). Optimizing the anchor locations:

Now we consider the case where we are free to select the anchor locations in ${ℝ}^d$. We consider the class of Mahalanobis costs described in Section 3.3. Let M_x, M_z, M_y be the Mahalanobis matrices correspond to C_x, C_z, and C_y, respectively.

Given the transport plans generated after solving (8), we can derive the the first-order stationary condition of OT^L with respect to Z_x and Z_y. Let

$$A=\left[\begin{array}{cc}D\left(u_z\right)\otimes \left(M_x+M_z\right)& P_z\otimes M_z\\ -P_z^T\otimes M_z& D\left(v_z\right)\otimes \left(M_y+M_z\right)\end{array}\right]$$

The update formula is given by

$$\left[\begin{array}{c}\text{vec}\left(Z_x^{*}\right)\\ \text{vec}\left(Z_y^{*}\right)\end{array}\right]=A^{-1}\times \left[\begin{array}{c}\left(P_x^T\otimes M_x\right)\text{vec}(X)\\ \left(P_y\otimes M_y\right)\text{vec}(Y)\end{array}\right],$$ (9)

where vec(·) denotes the operator converting a matrix to a column vector, and D(·) denotes the operator converting a vector to a diagonal matrix. We defer the detailed derivation to Appendix B.2. Pseudo-code for the combined scheme can be found in Algorithm 1.

(3). Robust estimation of data transport:

LOT provides robust transport in the target domain by aligning the data through anchors, which can facilitate regression, and classification in downstream applications. We denote the centroids of the source and target by $Q_x=\text{diag}\left(u_z^{-1}\right)P_x^T{X}^T$, $Q_y=\text{diag}\left(v_z^{-1}\right)P_y{Y}^T$. We propose the estimator $\widehat{X}≔{\sum }_{m,n}p\left(z_m^x,z_n^y\mid x\right)\left(Q_m^y-Q_n^x\right)=\text{diag}\left({\mu }^{-1}\right)P_x\text{diag}\left({\left(P_{z}1\right)}^{-1}\right)P_z\left(Q_y-Q_x\right)$. In contrast to factored coupling (Forrow et al., 2019), where Z_x = Z_y, LOT is robust even when the source and target have different structures (see Table 1 MNIST-DU, Figure 4).

Table 1: Results for concept drift and domain adaptation for handwritten digits.

The classification accuracy and L2-error are computed after transport for MNIST to USPS (left) and coarse dropout (right). Our method is compared with the accuracy before alignment (Original), entropy-regularized OT, k-means plus OT (KOT), and subspace alignment (SA).

	MNIST-USPS	MNIST-DU
	Accuracy	Accuracy	L2 error
Original	79.3	72.6	0.72
OT	76.9	61.5	0.71
KOT	79.4	60.9	0.73
SA	81.3	72.3	-
FC	84.1	67.2	0.59
LOT-WA	86.2	77.7	0.56

Figure 4: Visualization of results on handwritten digits and examples of domain shift.

(a) 2D projections of representations formed in deep neural network before (top) and after different alignment methods (LOT, FC, OT). (b) confusion matrices for LOT (top) and FC (bottom) after alignment. The transport plans are visualized for LOT (c) and FC (d) in the unbalanced case.

(4). Implementation details:

LOT has two primary hyperparameters that must be specified: (i) the number of the source and target anchors k_x, k_y and (ii) the regularization parameter ε. For details on the tuning of these parameters, please refer to Appendix F. In practice, we use centroids from k-means clustering (Arthur & Vassilvitskii, 2006) to initialize the anchors, and for all the experiments we have conducted, LOT typically converges within 20 iterations.

4. Related Work

Interpolation between factored coupling and k-means clustering:

Assume we select the Mahalanobis matrices of the costs defined in Section 3.3 to be M_x = M_y = I, and M_z = λI. If we let λ → ∞ when estimating the transport between source and target anchors, the anchors merge, and our approach reduces to the case of factored coupling (Forrow et al., 2019). At the other end, if we let λ → 0, then LOT becomes separable, and the middle term vanishes. In this case, each remaining term exactly corresponds to a pure clustering task, and LOT reduces to k-means clustering (Arthur & Vassilvitskii, 2006).

Relationship to OT-based clustering methods:

Many methods that combine OT and clustering (Li & Wang, 2008; Ye et al., 2017; Ho et al., 2017; Dessein et al., 2017; Genevay et al., 2019; Alvarez-Melis & Fusi, 2020) focus on using the Wasserstein distance to identify barycenters that serve as the centroids of clusters. When finding barycenters for the source and target separately, this could be seen as LOT with C_z = 0 and C_x, C_y defined using a squared L2 distance. In other related work (Laclau et al., 2017), co-clustering is applied to a transport plan as a post-processing operation, and no additional regularization on the transportation cost in OT is imposed. In contrast, our approach induces explicit regularization by separately defining cost matrices for the transport between the source/target points and their anchors. This yields a transport plan guided by a cluster-level matching.

Relationship to hierarchical OT:

Hierarchical OT (Chen et al., 2018; Lee et al., 2019; Yurochkin et al., 2019; Xu et al., 2020) transports points by moving them within some predetermined subgroup simultaneously based on either their class label or pre-specified structures, and then forms a matching of these subgroups using the Wasserstein distance. The resulting problem solves a multi-layer OT problem which gives rise to its name. With a Wasserstein distance used to build the C_z cost matrix, LOT effectively reduces to hierarchical OT for fixed and hard-class assignment P_x and P_y. However, a crucial difference between LOT and hierarchical OT lies in that the latter imposes the known structure information. In contrast, LOT discovers this structure by simultaneously learning P_x and P_y.

Transportation with anchors:

The notion of moving data points with anchors to match points in heterogeneous spaces has appeared in other work (Sato et al., 2020; Manay et al., 2006). These approaches map each point from one domain into a distribution of the costs, which effectively builds up a common representation for the points from both spaces. In contrast to this work, we use the anchors to encourage clustering of data and to impose rank constraints on the transport plan.

5. Theoretical Analysis

LOT as a relaxation of OT:

We now ask how the optimal value of our original rank-constrained objective in (8) is related to the transportation cost defined in entropy-regularized OT. It turns out their objectives are connected by an inequality described below (see Appendix A for a proof).

Proposition 2.

Let P be a transport plan of the form in (5). Assume that K is some Gibbs kernel that satisfies,

$$K_x{K}_z{K}_y\le K,$$ (10)

where the inequality is over each entry. Then we have,

$$\epsilon KL(P\Vert K)\le \epsilon (KL\left(P_x\Vert K_x\right)+KL\left(P_z\Vert K_z\right)+KL\left(P_y\Vert K_y\right))+\epsilon \left(H\left(u_z\right)+H\left(v_z\right)\right),$$ (11)

where H(a) ≔ −∑_i a_i log a_i denotes the entropy.

The proposition shows that an OT objective, corresponding to a kernel K (resp. C), can be upper bounded by three sub-OT problems defined by subsequent kernels K_x, K_z, K_y (resp. C_x, C_z, C_y) that satisfies (10) (resp. exp(−C_x/ε)exp(−C_z/ε)exp(−C_y/ε) ≤ exp(−C/ε)).

Let us compare the upper bound given by Proposition 2 with Def. 1 and ignore the entropy terms; we recognize that it is precisely the entropy-regularized objective of LOT. In other words, with suitable cost matrices satisfying (10), LOT could be interpreted as a relaxation of an OT problem in a decomposed form. We then ask what C_x, C_z, C_y should be to satisfy (10). In cases where cost C is defined by the L^p-norm to the power p, the following corollary shows that the same form suffices.

Corollary 2.

Let $d(x,y)≔\Vert x-y{\Vert }_p^p$. Consider an optimal transport problem OT_C,ε with cost C[i, j] = d(x_i, y_j), where p ≥ 1. Then for a sufficiently small ε, the latent optimal transport OT^L with cost matrices, $C_x[i,m]=3^{p-1}d\left(x_i-z_m^x\right)$, $C_z[m,n]=3^{p-1}d\left(z_m^x-z_n^y\right)$, $C_y[n,j]=3^{p-1}d\left(z_n^y-y_j\right)$ minimizes an upper bound of the entropy-regularized OT objective in (2).

Corollary (2) provides natural costs for LOT to be posed as a relaxation to a OT problem with L^p norm. More generally, finding the optimal cost functions that obey (10) and minimize the gap in the inequality in Proposition 2 is outside the scope of this work but would be an interesting topic for future investigation.

Sampling complexity: Below we analyze LOT from a statistical point of view. Specifically, we bound the sampling rate of OT^L in Def. 1 when the true distributions μ and ν are estimated by their empirical distributions.

Proposition 3.

Suppose X and Y have distributions μ and ν supported on a compact region Ω in ${ℝ}^d$, the cost functions c_x(·,·) and c_y(·,·) are defined as the squared Euclidean distance, and $\widehat{\mu }$, $\widehat{\nu }$ are empirical distributions of n and m i.i.d. samples from μ and ν, respectively. If the spaces for latent distributions are equal to ${\mathcal{Z}}_x={\mathcal{Z}}_y={ℝ}^d$, and there are k_x and k_y anchors in the source and target, respectively, then with probability at least 1 − δ,

$$\text{Err}\le C\sqrt{\frac{k_{\mathit{\text{max}}}^{3}d\text{log}{k}_{\mathit{\text{max}}}+\text{log}(2/\delta )}{N}},$$ (12)

where $\text{Err }=\left|{\text{OT}}^L(\mu ,\nu )-{\text{OT}}^L(\widehat{\mu },\widehat{\nu })\right|$, k_max = max{k_x, k_y}, N = min{n, m} and C ≥ 0 is some constant not depending on N.

As shown in (Weed et al., 2019), the general sampling rate of a plug-in OT scales with $N^{\frac{1}{d}}$, suffering from the “curse of dimensionality”. On the other hand, as evidence from (Forrow et al., 2019), structural optimal transport paves ways to overcome the issue. In particular, LOT achieves $N^{-\frac{1}{2}}$ scaling by regularizing the transport rank.

Time complexity:

We can bound the time complexity as O(T_i + T_bcd(T_k + T_au + T_pu)), where T_i is the initialization complexity, e.g., if we use k-means, then it equals to O(nk_xdT_x + mk_ydT_y) where T_x and T_y are the iteration numbers of the Floyd algorithm applied to the source and target, respectively, T_bcd is the total number of iterations of block-coordinate descent, T_k = O(nk_x + mk_y) is the computation time for updating the kernels, T_au = O((k_x + k_y)³ + d(nk_x + mk_y)) is the complexity of updating anchors, and T_pu is the complexity for updating plans. Because our updates are based on iterative Bregman projections similar to the Sinkhorn algorithm, it has complexity comparable to OT. Therefore, the overall complexity of LOT is approximately T_bcd times of OT, assuming n, m ≥ d(k_x + k_y). Empirically, T_bcd depends on the structure of data, but we observed that it is usually under 20. Note that the same applies to FC with k_x = k_y = k. In Figure 3, we complement our analysis by simulating a comparison of the time complexity for LOT and FC vs. OT in the setting of a 7-component Gaussian mixture model. We can see the compute time of LOT scales similarly to FC.

Figure 3: Comparisons of the time complexity and loss.

The figure compares the time complexity (dashed) and linear loss (solid) of LOT, FC, and OT in the setting of the 7-component GMM model.

Transport cost:

We also compare the transport loss returned by LOT (blue), FC (orange), and OT (green) as a function of the number of anchors in Figure 3. For a fair comparison, we considered a balanced scenario where 7-component GMMs generate the source and target. The anchors of the source and target are chosen to be equal for LOT. The result shows that the losses are indeed higher for LOT and FC compared to OT but are fairly insensitive with to the chosen number of anchors. Moreover, we find that LOT has a slightly lower loss compared to FC even when we choose the number of source and target anchors to be equal.

6. Experiments

In this section, we conduct empirical investigations. Details of hyperparameter tuning can be found in Appendix F.

E1). Testing robustness to various data perturbations:

To better understand how different types of domain shift impact the transport plans generated by our approach, we considered different transformations between the source and target. To create synthetic data for this task, we generated multiple clusters/components using a k-dimensional Gaussian with random mean and covariance sampled from a Wishart distribution, randomly projected to a 5-dimensional subspace. The source and target are generated independently: we randomly sample a fixed number of points according to the true distribution for each cluster. We compared the performance of the LOT variants proposed in Section 3.3: LOT-L2 (orange curves) and LOT-WA (green curves) with baselines OT (blue curves) and rank regularized factored coupling (FC) (Graf & Luschgy, 2007) (red curves) in terms of their (i) classification rates and (ii) deviation from the original transport plan without perturbations, which we compute as Err(P − P₀) = ∥P − P₀∥_F/∥P₀∥_F, where P₀ is the transport plan obtained before perturbations. The results are averaged over 20 runs, and a 75% confidence interval is used. See Appendix E for further details.

When compared with OT, both our method and FC provide more stable class recovery, even with significant amounts of perturbations (Figure 2). When we examine the error term in the transport plan, we observe that, in most cases, the OT plan deviates rapidly, even for small amounts of perturbations. Both FC and LOT appear to have similar performances across rotations while OT’s performance decreases quickly. In experiment (b), we found that both LOT variants provide substantial improvements on classification subject to outliers, implying the applicability of LOT for noisy data. In experiment (c), we study LOT in the high-dimensional setting; we find that LOT-WA behaves similarly to FC with some degradation in performance after the dimension increases beyond 70. Next, in experiment (d), we fix the number of components in the target to be 10, while varying the number in the source from 4 to 10. In contrast to the outlier experiment in (b), LOT-WA shows more resilience to mismatches between the source and target. At the bottom of plot (d), we show the 2-Wasserstein distance (blue) and latent Wasserstein discrepancy (orange) defined in Proposition 1. This shows that the latent Wasserstein discrepancy does indeed provide an upper bound on the 2-Wasserstein distance. Finally, we look at the effect of transport rank on LOT and FC in (e). The plot shows that the slope for LOT is flatter than FC while maintaining similar performances.

Figure 2: Results on Gaussian mixture models.

In (a), we apply a rotation between the source and target, in (b) we add outliers, in (c) we vary the ambient dimension, in (d) the target is set to have 8 components, and we vary the number of components in the source to simulate source-target mismatch, in (e) we fix the rank to 10 and vary the number of factors (anchors) used in the approximation. Throughout, we simulate data according to a GMM and evaluate performance by measuring the classification accuracy (top) and computing the deviation between the transport plans before and after the perturbations with respect to the Fröbenius norm (bottom).

E2). Domain adaptation application:

In our next experiment, we used LOT to correct for domain shift in a neural network that is trained on one dataset but underperforms on a new but similar dataset (Table 1, Figure 4). MNIST and USPS are two handwritten digits datasets that are semantically similar but that have different pixel-level distributions and thus introduce domain shift (Figure 4a). We train a multi-layer perceptron (MLP) on the training set of the MNIST dataset, freeze the network, and use it for the remaining experiments. The classifier achieves 100% training accuracy and a 98% validation accuracy on MNIST but only achieves 79.3% accuracy on the USPS validation set. We project MNIST’s training samples in the classifier’s output space (logits) and consider the 10D projection to be the target distribution. Similarly, we project images from the USPS dataset in the network’s output space to get our source distribution. We study the performance of LOT in correcting the classifier’s outputs and compare with FC, k-means OT (KOT) (Forrow et al., 2019), and subspace alignment (SA) (Fernando et al., 2013).

In Table 1, we summarize the results of our comparisons on the domain adaptation task (MNIST-USPS). Our results suggest that both FC and LOT perform pretty well on this task, with LOT beating FC by 2% in terms of their final classification accuracy. We also show that LOT does better than naive KOT. In Figure 4a, we use Isomap to project the distribution of USPS images as well as the alignment results for LOT, FC, and OT. For both LOT and FC, we also display the anchors; note that for LOT, we have two different sets of anchors (source, red; target, blue). This example highlights the alignment of the anchors in our approach and contrasts it with that of FC.

Taking inspiration from studies in self-supervised learning (Doersch et al., 2015; He et al., 2020) that use different transformations of an input image (e.g., masking parts of the image) to build invariances into a network’s representations, here we ask how augmentations of the images introduce domain shift and whether our approach can correct/identify it. To test this, we apply coarse dropout on test samples in MNIST and feed them to the classifier to get a new source distribution. We do this in a balanced (all digits in source and target) and an unbalanced setting (2, 4, 8 removed from source, all digits in target). The results of the unbalanced dropout are summarized in Table 1 (MNIST-DU), and the other results are provided in Table S1 in the Appendix. In this case, we have the features of the testing samples pre-transformation, and thus, we can compare the transported features to the ground truth features in terms of their point-to-point error (L2 distance). In the unbalanced case, we observe even more significant gaps between FC and LOT, as the source and target datasets have different structures. To quantify these different class-level errors, we compare the confusion matrices for the classifier’s output after alignment (Figure 4b). By examining the columns corresponding to the removed digits, we see that FC is more likely to misclassify these images. Our results suggest that LOT has comparable performance with FC in a balanced setting and outperforms FC in an unbalanced case.

The decomposition in both LOT and FC allows us to visualize transport between the source, anchors, and the target (Figure 4c–d, S2). This visualization highlights the interpretability of the transport plans learned via our approach, with the middle transport plan P_z providing a concise map of interactions between class manifolds in the unbalanced setting. With LOT (Figure 4c), we find that each source anchor is mapped to the correct target anchor, with some minor interactions with the target anchors corresponding to the removed digits. In comparison, FC (Figure 4d, S2) has more spurious interactions between source, anchors, and target.

E3). Robustness to sampling:

We examined the robustness to sampling for the MNIST to USPS example (Figure 5). In this case, we find that LOT has a stable alignment as we subsample the source dataset, with very little degradation in classification accuracy, even as we reduce the source to only 20 samples. We also observe a significant gap between LOT and other approaches in this experiment, with more than a 10% gap between FC and LOT when very few source samples are provided. Our results demonstrate that LOT is robust to subsampling, providing empirical evidence for Proposition 3.

Figure 5: LOT provides robust alignment, even when given very few samples.

We compare our method with OT and FC on the MNIST-USPS domain adaptation task when different numbers of USPS samples are available. Reported classification rates are averaged over 50 random sets.

7. Discussion

In this paper, we introduced LOT, a new form of structured transport leading to an approach for jointly clustering and aligning data. We provided an efficient optimization method to solve for the transport, and studied its statistical rate via theoretical analysis and robustness to data perturbations with empirical experiments. In the future, we would like to explore the application of LOT to non-Euclidean spaces, and incorporate metric learning into our framework.