SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Milena Gazdieva*
Skolkovo Institute of Science and Technology
milena.gazdieva@skoltech.ru
Alexander Korotin*
Artificial Intelligence Research Institute
Daniil Selikhanovych
selikhanovychdaniil@gmail.com
In many unpaired image domain translation problems, e.g., style transfer or superresolution, it is important to keep the translated image similar to its respective input image. We propose the extremal transport (ET) which is a mathematical formalization of the theoretically best possible unpaired translation between a pair of domains w.r.t. the given similarity function. Inspired by the recent advances in neural optimal transport (OT), we propose a scalable algorithm to approximate ET maps as a limit of partial OT maps. We test our algorithm on toy examples and on the unpaired image-to-image translation task. The code is publicly available at
https://github.com/milenagazdieva/ExtremalNeuralOptimalTransport
(a) Handbag  $\rightarrow$  shoes  $(128\times 128)$
Figure 1: (Nearly) extremal transport with our Algorithm 1.
(b) Celeba (female)  $\rightarrow$  anime  $(64\times 64)$
Higher  $w$  yields bigger similarity of  $x$  and  $T(x)$  in  $\ell^2$ .
The unpaired translation task [72, Fig. 2] is to find a map  $x \mapsto T(x)$ , usually a neural network, which transports the samples  $x$  from the given source domain to the target domain. The key challenge
here is that the correspondence between available data samples  $x$  from the source and  $y$  from target domains is not given. Thus, the task is ambiguous as there might exist multiple suitable  $T$ .
When solving the task, many methods regularize the translated samples  $T(x)$  to inherit specific attributes of the respective input samples  $x$ . In the popular unpaired translation [72, Fig. 9] and enhancement [69, Equation 3] tasks for images, it is common to use additional unsupervised identity losses, e.g.,  $\| T(x) - x\| _1$ , to make the translated output  $T(x)$  be similar to the input images  $x$ . The same applies, e.g., to audio translation [51]. Therefore, the learning objectives of such methods usually have two components.
Figure 2: Learned transport map  $(w = 2)$  in 'Accept' task.
The first component is the domain loss (main) enforcing the translated sample  $T(x)$  to look like the samples  $y$  from the target domain. The second component is the similarity loss (regularizer, optional) stimulating the translated  $T(x)$  to inherit certain attributes of input  $x$ . A question arises: can one obtain the maximal similarity of  $T(x)$  to  $x$  but still ensure that  $T(x)$  is indeed from the target domain? A straightforward "yes, just increase the weight of the similarity loss" may work but only to a limited extent. We demonstrate this in Appendix C.
Contributions. In this paper, we propose the extremal transport (ET, §3.1) which is a rigorous mathematical task formulation describing the theoretically best possible unpaired domain translation w.r.t. the given similarity function. We explicitly characterize ET maps and plans by establishing an intuitive connection to the nearest neighbors (NN). We show that ET maps can be learned as a limit (§3.3) of specific partial optimal transport (OT) problem which we call incomplete transport (IT, §3.2). For IT, we derive the duality formula yielding an efficient computational algorithm (§3.4). We test our algorithm on toy 2D examples and high-dimensional unpaired image translation (§5).
Notation. We consider compact Polish spaces  $(\mathcal{X},\| \cdot \|_{\mathcal{X}})$ ,  $(\mathcal{V},\| \cdot \|_{\mathcal{V}})$  and use  $\mathcal{P}(\mathcal{X}),\mathcal{P}(\mathcal{V})$  to denote the sets of Radon probability measures on them. We use  $\mathcal{M}_{+}(\mathcal{X})\subset \mathcal{M}(\mathcal{X})$  to denote the sets of finite non-negative and finite signed (Radon) measures on  $\mathcal{X}$ , respectively. They both contain  $\mathcal{P}(\mathcal{X})$  as a subset. For a non-negative  $\mu \in \mathcal{M}_{+}(\mathcal{X})$ , its support is denoted by  $\operatorname {Supp}(\mu)\subset \mathcal{X}$ . It is a closed set consisting of all points  $x\in \mathcal{X}$  for which every open neighbourhood  $A\ni x$  satisfies  $\mu (A) > 0$ . We use  $\mathcal{C}(\mathcal{X})$  to denote the set of continuous functions  $\mathcal{X}\to \mathbb{R}$  equipped with  $\| \cdot \|_{\infty}$  norm. Its dual space is  $\mathcal{M}(\mathcal{X})$  equipped with the  $\| \cdot \| _1$  norm. A sequence  $\mu_1,\mu_2,\dots \in \mathcal{M}(\mathcal{X})$  is said to be weakly-\* converging to  $\mu^{*}\in \mathcal{M}(\mathcal{X})$  if for every  $f\in \mathcal{C}(\mathcal{X})$  it holds that  $\lim_{n\to \infty}\int_{\mathcal{X}}f(x)d\mu_n(x) = \int_{\mathcal{X}}f(x)d\mu^* (x)$ . For a probability measure  $\pi \in \mathcal{P}(\mathcal{X}\times \mathcal{Y})$ , we use  $\pi_x\in \mathcal{P}(\mathcal{X})$  and  $\pi_y\in \mathcal{P}(\mathcal{Y})$  to denote its projections onto  $\mathcal{X},\mathcal{Y}$ , respectively. Disintegration of  $\pi$  yields  $d\pi (x,y) = d\pi_{x}(x)d\pi (y|x)$ , where  $\pi (y|x)$  denotes the conditional distribution of  $y\in \mathcal{V}$  for a given  $x\in \mathcal{X}$ . For  $\mu ,\nu \in \mathcal{M}(\mathcal{V})$  we write  $\mu \leq \nu$  if for all measurable  $A\subset \mathcal{V}$  it holds that  $\mu (A)\leq \nu (A)$ . For a measurable map  $T:\mathcal{X}\rightarrow \mathcal{V}$ , we use  $T\nexists$  to denote the associated pushforward operator  $\mathcal{P}(\mathcal{X})\to \mathcal{P}(\mathcal{V})$ .
2 Background on Optimal Transport
In this section, we give an overview of the OT theory concepts related to our paper. For details on OT, we refer to [62, 65, 56], partial OT - [22, 10].
Standard OT formulation. Let  $c: \mathcal{X} \times \mathcal{Y} \to \mathbb{R}$  be a continuous cost function. For  $\mathbb{P} \in \mathcal{P}(\mathcal{X})$ ,  $\mathbb{Q} \in \mathcal{P}(\mathcal{Y})$ , the OT cost between them is given by
$$
\operatorname {C o s t} (\mathbb {P}, \mathbb {Q}) \stackrel {\text {d e f}} {=} \inf  _ {T \sharp \mathbb {P} = \mathbb {Q}} \int_ {\mathcal {X}} c (x, T (x)) d \mathbb {P} (x), \tag {1}
$$
where inf is taken over measurable  $T:\mathcal{X}\to \mathcal{Y}$  pushing  $\mathbb{P}$  to  $\mathbb{Q}$  (transport maps), see Fig. 3a. Problem (1) is called the Monge's OT problem, and its minimizer  $T^{*}$  is called an OT map.
In some cases, there may be no minimizer  $T^{*}$  of (1). Therefore, it is common to consider Kantorovich's relaxation:
$$
\operatorname {C o s t} (\mathbb {P}, \mathbb {Q}) \stackrel {\text {d e f}} {=} \inf  _ {\pi \in \Pi (\mathbb {P}, \mathbb {Q})} \int_ {\mathcal {X} \times \mathcal {Y}} c (x, y) d \pi (x, y), \tag {2}
$$
where inf is taken over  $\pi \in \mathcal{P}(\mathcal{X}\times \mathcal{Y})$  satisfying  $\pi_{x} = \mathbb{P}$  and  $\pi_y = \mathbb{Q}$ , respectively. A minimizer  $\pi^{*}\in \Pi (\mathbb{P},\mathbb{Q})$  in (2) always exists and is called an OT plan. A widely used example of OT cost for  $\mathcal{X} = \mathcal{Y} = \mathbb{R}^d$  is the Wasserstein-1 distance  $(\mathbb{W}_1)$ , i.e., OT cost (2) for  $c(x,y) = \| x - y\|$ .
Figure 3: Classic optimal transport (OT) formulations.
To provide an intuition behind (2), we disintegrate  $d\pi(x,y) = d\pi_x(x)d\pi(y|x)$ :
$$
\inf  _ {\pi \in \Pi (\mathbb {P}, \mathbb {Q})} \int_ {\mathcal {X}} \left\{\int_ {\mathcal {Y}} c (x, y) d \pi (y | x) \right\} \underbrace {d \mathbb {P} (x)} _ {= d \pi_ {x} (x)}, \tag {3}
$$
i.e., (2) can be viewed as an extension of (1) allowing to split the mass of input points  $x \sim \mathbb{P}$  (Fig. 3b). With mild assumptions on  $\mathbb{P}, \mathbb{Q}$ , the OT cost value (2) coincides with (1), see [62, Theorem 1.33].
Partial OT formulation. Let  $w_0, w_1 \geq m \geq 0$ . We consider
$$
\inf  _ { \begin{array}{l} m \pi_ {x} \leq w _ {0} \mathbb {P} \\ m \pi_ {y} \leq w _ {1} \mathbb {Q} \end{array} } \int_ {\mathcal {X} \times \mathcal {Y}} c (x, y) d [ m \pi ] (x, y), \tag {4}
$$
where inf is taken over  $\pi \in \mathcal{P}(\mathcal{X}\times \mathcal{Y})$  satisfying the inequality constraints  $m\pi_{x}\leq w_{0}\mathbb{P}$  and  $m\pi_y\leq w_1\mathbb{Q}$ . Minimizers  $\pi^*$  of (4) are called partial OT plans (Fig. 4).
Here the inputs are two measures  $w_0\mathbb{P}$  and  $w_1\mathbb{Q}$  with masses  $w_0$  and  $w_1$ . Intuitively, we need to match a  $\frac{m}{w_0}$ -th fraction  $m\pi_x$  of the first measure  $w_0\mathbb{P}$  with a  $\frac{m}{w_1}$ -th fraction  $m\pi_y$  of the second measure  $w_1\mathbb{Q}$  (Fig. 4); choosing  $\pi_x, \pi_y$  is also a part of this problem. The key difference from problem (4) is that the constraints are inequalities. In the particular case  $m = w_0 = w_1$ , problem (4) reduces to (2) as the inequality constraints can be replaced by equalities.
Figure 4: Partial optimal transport formulation.
First, we formulate the extremal transport (ET) problem (§3.1). Next, we prove that ET maps can be recovered as a limit of incomplete transport (IT) maps (§3.2, 3.3). Then we propose an algorithm to solve the IT problem (§3.4). We provide the proofs for all the theorems in Appendix F.
3.1 Extremal Transport Problem
Popular unpaired translation methods, e.g., [72, §3.1] and [33, §3], de-facto assume that available samples  $x, y$  from the input and output domains come from the data distributions  $\mathbb{P}, \mathbb{Q} \in \mathcal{P}(\mathcal{X}), \mathcal{P}(\mathcal{Y})$ . As a result, in their optimization objectives, the domain loss compares the translated  $T(x) \sim T\sharp \mathbb{P}$  and target samples  $y \sim \mathbb{Q}$  by using a metric for comparing probability measures, e.g., GAN loss [27]. Thus, the target domain is identified with the probability measure  $\mathbb{Q}$ .
We pick a different approach to define what the domain is. We still assume that the available data comes from data distributions, i.e.,  $x \sim \mathbb{P}, y \sim \mathbb{Q}$ . However, we say that the target domain is the part of  $\mathcal{V}$  where the probability mass of  $\mathbb{Q}$  lives.2 Namely, it is  $\operatorname{Supp}(\mathbb{Q}) \subset \mathcal{V}$ . We say that a map  $T$  translates the domains if  $\operatorname{Supp}(T \sharp \mathbb{P}) \subset \operatorname{Supp}(\mathbb{Q})$ . This requirement is weaker than the usual  $T \sharp \mathbb{P} = \mathbb{Q}$ . Assume that  $c(x, y)$  is a function estimating the dissimilarity between  $x, y$ . We would like to pick  $T(x) \in \operatorname{Supp}(\mathbb{Q})$  which is maximally similar to  $x$  in terms of  $c(x, y)$ . This preference of  $T$  can be formalized as follows:
$$
\operatorname {C o s t} _ {\infty} (\mathbb {P}, \mathbb {Q}) \stackrel {{d e f}} {{=}} \inf  _ {\operatorname {S u p p} (T \sharp \mathbb {P}) \subset \operatorname {S u p p} (\mathbb {Q})} \int_ {\mathcal {X}} c (x, T (x)) d \mathbb {P} (x), \tag {5}
$$
where the inf is taken over measurable  $T: \mathcal{X} \to \mathcal{Y}$  which map the probability mass of  $\mathbb{P}$  to  $\operatorname{Supp}(\mathbb{Q})$ . We say that (5) is the (Monge's) extremal transport (ET) problem.
$^2$ Following the standard manifold hypothesis [21], real data distribution  $\mathbb{Q}$  is usually supported on a small-dimensional manifold  $M = \operatorname{Supp}(\mathbb{Q}) \subset [-1, 1]^D$  occupying a tiny part of the ambient space  $[-1, 1]^D$ .
(a) Monge's ET formulation.
(b) Kantorovich's ET formulation.
Figure 5: Extremal transport (ET) formulations.
Problem (5) is atypical for the common OT framework. For example, the usual measure-preserving constraint  $T \sharp \mathbb{P} = \mathbb{Q}$  in (1) is replaced with  $\operatorname{Supp}(T \sharp \mathbb{P}) \subset \operatorname{Supp}(\mathbb{Q})$  which is more tricky. Importantly, measure  $\mathbb{Q}$  can be replaced with any other  $\mathbb{Q}' \in \mathcal{P}(\mathcal{V})$  with the same support yielding the same inf. Below we analyse the minimizers  $T^*$  of (5). We define  $c^*(x) \stackrel{\text{def}}{=} \min_{y \in \operatorname{Supp}(\mathbb{Q})} c(x, y)$ . Here the min is indeed attained (for all  $x \in \mathcal{X}$ ) because  $c(x, y)$  is continuous and  $\operatorname{Supp}(\mathbb{Q}) \subset \mathcal{V}$  is a compact set. The value  $c^*(x)$  can be understood as the lowest possible transport cost when mapping the mass of point  $x$  to the support of  $\mathbb{Q}$ . For any admissible  $T$  in (5), it holds (P-almost surely):
$$
c ^ {*} (x) = \min  _ {y \in \operatorname {S u p p} (\mathbb {Q})} c (x, y) \leq c (x, T (x)). \tag {6}
$$
Proposition 1 (Continuity of  $c^*$ ). It holds that  $c^* \in \mathcal{C}(\mathcal{X})$ .
As a consequence of Proposition 1, we see that  $c^*$  is measurable. We integrate (6) w.r.t.  $x \sim \mathbb{P}$  and take inf over all feasible  $T$ . This yields a lower bound on  $\mathrm{Cost}_{\infty}(\mathbb{P}, \mathbb{Q})$ :
$$
\int_ {\mathcal {X}} c ^ {*} (x) d \mathbb {P} (x) \leq \overbrace {\inf  _ {\substack {\operatorname {S u p p} (T \sharp \mathbb {P}) \\ \subset \operatorname {S u p p} (\mathbb {Q})}} \int_ {\mathcal {X}} c (x , T (x)) d \mathbb {P} (x)} ^ {\operatorname {C o s t} _ {\infty} (\mathbb {P}, \mathbb {Q})}. \tag{7}
$$
There exists admissible  $T$  making (7) the equality. Indeed, let  $\mathrm{NN}(x)\stackrel {\mathrm{def}}{=}\{y\in \operatorname {Supp}(\mathbb{Q})\text{s.t.} c(x,y) = c^{*}(x)\}$  be the set of points  $y$  which attain min in the definition of  $c^*$ . These points are the closest to  $x$  points in  $\mathbb{Q}$  w.r.t. the cost  $c(x,y)$ . We call them the nearest neighbors of  $x$ . From this perspective, we see that (7) turns to equality if and only if  $T(x)\in \mathrm{NN}(x)$  holds for  $\mathbb{P}$ -almost all  $x\in \mathcal{X}$ , i.e.,  $T$  maps points  $x\sim \mathbb{P}$  to their nearest neighbors in  $\operatorname {Supp}(\mathbb{Q})$ . We need to make sure that such measurable  $T$  exists (Fig. 5a).
Theorem 1 (Existence of ET maps). There exists at least one measurable map  $T^{*}: \mathcal{X} \to \mathcal{Y}$  minimizing (5). For  $\mathbb{P}$ -almost all  $x \in \mathcal{X}$  it holds that  $T^{*}(x) \in \mathrm{NN}(x)$ . Besides,
$$
\operatorname {C o s t} _ {\infty} (\mathbb {P}, \mathbb {Q}) = \int_ {\mathcal {X}} c ^ {*} (x) d \mathbb {P} (x).
$$
We say that  $\mathrm{Cost}_{\infty}(\mathbb{P},\mathbb{Q})$  is the extremal cost because one can not obtain smaller cost when moving the mass of  $\mathbb{P}$  to  $\mathrm{Supp}(\mathbb{Q})$ . In turn, we say that minimizers  $T^{*}$  are  $ET$  maps.
One may extend the ET problem (8) in the Kantorovich's manner by allowing the mass splitting and stochastic plans:
$$
\operatorname {C o s t} _ {\infty} (\mathbb {P}, \mathbb {Q}) \stackrel {{d e f}} {{=}} \inf  _ {\pi \in \Pi^ {\infty} (\mathbb {P}, \mathbb {Q})} \int_ {\mathcal {X} \times \mathcal {Y}} c (x, y) d \pi (x, y), \tag {8}
$$
where  $\Pi^{\infty}(\mathbb{P},\mathbb{Q})$  are probability measures  $\pi \in \mathcal{P}(\mathcal{X}\times \mathcal{Y})$  s.t.  $\pi_{x} = \mathbb{P}$  and  $\operatorname {Supp}(\pi_y)\subset \operatorname {Supp}(\mathbb{Q})$  To understand the structure of minimizers in (8), it is more convenient to disintegrate  $d\pi (x,y) = d\pi (y|x)d\pi_x(x)$
$$
\operatorname {C o s t} _ {\infty} (\mathbb {P}, \mathbb {Q}) = \inf  _ {\pi \in \Pi^ {\infty} (\mathbb {P}, \mathbb {Q})} \int_ {\mathcal {X}} \int_ {\mathcal {Y}} c (x, y) d \pi (y | x) \underbrace {d \mathbb {P} (x)} _ {= d \pi_ {x} (x)}. \tag {9}
$$
Thus, computing (8) boils down to computing a family of conditional measures  $\pi (\cdot |x)$  minimizing (8). As in (7), for any  $\pi \in \Pi^{\infty}(\mathbb{P},\mathbb{Q})$  , it holds (for  $\mathbb{P}$  -almost all  $x\in \mathcal{X}$  ) that
$$
c ^ {*} (x) = \min  _ {y \in \operatorname {S u p p} (\mathbb {Q})} c (x, y) \leq \int_ {\mathcal {Y}} c (x, y) d \pi (y | x) \tag {10}
$$
because  $\pi$  redistributes the mass of  $\mathbb{P}$  to  $\mathrm{Supp}(\mathbb{Q})$ . By integrating (10) w.r.t.  $x \sim \mathbb{P} = \pi_x$  and taking inf over all admissible plans  $\pi$ , we derive that  $\int_{\mathcal{X}} c^*(x) d\mathbb{P}(x)$  is a lower bound for (8). In particular,
(a) Monge's IT formulation.
(b) Kantorovich's IT formulation.
Figure 6: Incomplete transport (IT) formulations.
the bound is tight for  $\pi^{*}(y|x) = \delta_{T^{*}(x)}$ , where  $T^{*}$  is the ET map from our Theorem 1. Therefore, the value (8) is the same as (5) but possibly admits more minimizers. We call the minimizers  $\pi^{*}$  of (8) the ET plans (Fig. 5b).
From (10) and the definition of  $\mathrm{NN}(x)$ , we see that minimizers  $\pi^{*}$  are the plans for which  $\pi^{*}(y|x)$  redistributes the mass of  $x$  among the nearest neighbors  $y\in \mathrm{NN}(x)$  of  $x$  (for  $\mathbb{P}$ -almost all  $x\in \mathcal{X}$ ). As a result,  $\pi^{*}(y|x)$  can be viewed as a stochastic nearest neighbor assignment between the probability mass of  $\mathbb{P}$  and the support of  $\mathbb{Q}$ .
3.2 Incomplete Transport Problem
In practice, solving extremal problem (5) is challenging because it is hard to enforce  $\operatorname{Supp}(T\sharp \mathbb{P}) \subset \operatorname{Supp}(\mathbb{Q})$ . To avoid enforcing this constraint, we replace it and consider the following problem with finite parameter  $w \geq 1$ :
$$
\operatorname {C o s t} _ {w} (\mathbb {P}, \mathbb {Q}) \stackrel {{d e f}} {{=}} \inf  _ {T \sharp \mathbb {P} \leq w \mathbb {Q}} \int_ {\mathcal {X}} c (x, T (x)) d \mathbb {P} (x). \tag {11}
$$
We call (11) Monge's incomplete transport (IT) problem (Fig. 6a). With the increase of  $w$ , admissible maps  $T$  obtain more ways to redistribute the mass of  $\mathbb{P}$  among  $\operatorname{Supp}(\mathbb{Q})$ . Informally, when  $w \to \infty$ , the constraint  $T \sharp \mathbb{P} \leq w\mathbb{Q}$  in (11) tends to the constraint  $\operatorname{Supp}(T \sharp \mathbb{P}) \subset \operatorname{Supp}(\mathbb{Q})$  in (5), i.e., (11) itself tends to ET problem (5). We will formalize this statement a few paragraphs later (in §3.3).
As in (1), problem (11) may have no minimizer  $T^{*}$  or even may have the empty feasible set. Therefore, it is natural to relax problem (11) in the Kantorovich's manner:
$$
\operatorname {C o s t} _ {w} (\mathbb {P}, \mathbb {Q}) \stackrel {\text {d e f}} {=} \inf  _ {\pi \in \Pi^ {w} (\mathbb {P}, \mathbb {Q})} \int_ {\mathcal {X} \times \mathcal {Y}} c (x, y) d \pi (x, y), \tag {12}
$$
where the inf is taken over the set  $\Pi^w (\mathbb{P},\mathbb{Q})$  of probability measures  $\pi \in \mathcal{P}(\mathcal{X}\times \mathcal{Y})$  whose first marginal is  $\pi_{x} = \mathbb{P}$ , and the second marginal satisfies  $\pi_y\leq w\mathbb{Q}$  (Fig. 6b).
We note that IT problem (12) is a special case of partial OT (4) with  $w_0 = m = 1$  and  $w_1 = w$ . In (12), one may actually replace inf with min, see our proposition below.
Proposition 2 (Existence of IT plans). Problem (12) admits at least one minimizer  $\pi^{*}\in \Pi^{w}(\mathbb{P},\mathbb{Q})$
We say that minimizers of (12) are  $IT$  plans. In the general case, Kantorovich's  $IT$  cost (12) always lower bounds Monge's counterpart (11). Below we show that they coincide in the practically most interesting Euclidean case.
Proposition 3 (Equivalence of Monge's, Kantorovich's IT costs). Let  $\mathcal{X},\mathcal{Y}\subset \mathbb{R}^{D}$  be two compact sets,  $\mathbb{P}\in \mathcal{P}(\mathcal{X})$  be atomless,  $\mathbb{Q}\in \mathcal{P}(\mathcal{Y})$ . Then Monge's (11) and Kantorovich's (12) IT costs coincide.
However, it is not guaranteed that inf in Monge's problem (11) is attained even in the Euclidean case. Still for general Polish spaces  $\mathcal{X},\mathcal{Y}$  it is clear that if there exists a deterministic IT plan in Kantorovich's problem (12) of the form  $\pi^{*} = [\mathrm{id}_{\mathcal{X}},T^{*}]$ , then  $T^{*}$  is an IT map in (11), and the IT Monge's (11) and Kantorovich's (12) costs coincide. Henceforth, for simplicity, we assume that  $\mathcal{X},\mathcal{Y},c,\mathbb{P},\mathbb{Q}$  are such that (11) and (12) coincide, e.g., those from Prop. 3.
IT problem (12) can be viewed as an interpolation between OT (2) and ET problems (8). Indeed, when  $w = 1$ , the constraint  $\pi_y \leq \mathbb{Q}$  is equivalent to  $\pi_y = \mathbb{Q}$  as there is only one probability measure which is  $\leq \mathbb{Q}$ , and it is  $\mathbb{Q}$  itself. Thus, IT (12) with  $w = 1$  coincides with OT (2). In the next section, we show that for  $w \to \infty$  one recovers ET from IT.
3.3 Link between Incomplete and Extremal Transport
Now we connect incomplete (12) and extremal (8) transport tasks.
Theorem 2 (IT costs converge to the ET cost when  $w \to \infty$ ). Function  $w \mapsto \operatorname{Cost}_w(\mathbb{P}, \mathbb{Q})$  is convex, non-increasing in  $w \in [1, +\infty)$  and
$$
\lim  _ {w \to \infty} \operatorname {C o s t} _ {w} (\mathbb {P}, \mathbb {Q}) = \operatorname {C o s t} _ {\infty} (\mathbb {P}, \mathbb {Q}).
$$
A natural subsequent question here is whether IT plans in (12) converge to ET plans (8) when  $w \to \infty$ . Our following result sheds the light on this question.
Theorem 3 (IT plans converge to ET plans when  $w \to \infty$ ). Consider  $w_{1}, w_{2}, w_{3}, \dots \geq 1$  satisfying  $\lim_{n \to \infty} w_{n} = \infty$ . Let  $\pi^{w_{n}} \in \Pi^{w_{n}}(\mathbb{P}, \mathbb{Q})$  be a sequence of IT plans solving (12) with  $w = w_{n}$ , respectively. Then it has a (weakly-*') converging sub-sequence. Every such sub-sequence of IT plans converges to an ET plan  $\pi^{*} \in \Pi^{\infty}(\mathbb{P}, \mathbb{Q})$ .
In general, there may be sub-sequences of IT plans converging to different ET plans  $\pi^{*}\in \Pi^{\infty}(\mathbb{P},\mathbb{Q})$ . However, our following corollary shows that with the increase of weight  $w$ , elements of any subsequence become closer to the set of ET plans.
Corollary 1 (IT plans become closer to the set of ET plans when  $w \to \infty$ ). For all  $\varepsilon > 0 \exists w(\varepsilon) \in [1, \infty)$  such that  $\forall w \geq w(\varepsilon)$  and  $\forall IT$  plan  $\pi^w \in \Pi_w(\mathbb{P}, \mathbb{Q})$  solving Kantorovich's IT problem (12), there exists an ET plan  $\pi^*$  which is  $\varepsilon$ -close to  $\pi^w$  in  $\mathbb{W}_1$ , i.e.,  $\mathbb{W}_1(\pi^*, \pi^w) \leq \varepsilon$ .
Providing a stronger convergence result here is challenging, and we leave this theoretical question open for future studies. Our Theorems 2, 3 and Corollary 1 suggest that to obtain a fine approximation of an ET plan  $(w = \infty)$ , one may use an IT plan for sufficiently large finite  $w$ . In Appendix G.1, we empirically demonstrate this observation through an experiment where the ground-truth deterministic ET plan is analytically known. Below we develop a neural algorithm to compute IT plans.
3.4 Computational Algorithm for Incomplete Transport
To begin with, for IT (12), we derive the dual problem.
Theorem 4 (Dual problem for IT). It holds
$$
\operatorname {C o s t} _ {w} (\mathbb {P}, \mathbb {Q}) = \max  _ {f \leq 0} \int_ {\mathcal {X}} f ^ {c} (x) d \mathbb {P} (x) + w \int_ {\mathcal {Y}} f (y) d \mathbb {Q} (y), \tag {13}
$$
where the max is taken over non-positive  $f \in \mathcal{C}(\mathcal{Y})$  and  $f^c(x) \stackrel{\text{def}}{=} \min_{y \in \mathcal{Y}} \{ c(x, y) - f(y) \}$ .
We call the function  $f$  potential. In the definition of  $f^c$ , min is attained because  $c, f$  are continuous and  $\mathcal{V}$  is compact. The function  $f^c$  is called the  $c$ -transform of  $f$ .
The difference of formula (13) from usual  $c$ -transform-based duality formulas for OT (2), see [62, §1.2], [65, §5], is that  $f$  is required to be non-positive and the second term is multiplied by  $w \geq 1$ . We rewrite the term  $\int_{\mathcal{X}} f^c(x) d\mathbb{P}(x)$  in (13):
$$
\int_ {\mathcal {X}} f ^ {c} (x) d \mathbb {P} (x) = \int_ {\mathcal {X}} \min  _ {y \in \mathcal {Y}} \left\{c (x, y) - f (y) \right\} d \mathbb {P} (x) = \inf  _ {T: \mathcal {X} \rightarrow \mathcal {Y}} \int_ {\mathcal {X}} \left\{c (x, T (x)) - f (y) \right\} d \mathbb {P} (x). \tag {14}
$$