$\displaystyle=\underbrace{\widehat{\mathsf{W}}_{\alpha,k}(\widehat{\mu}_{k})-\widehat{\mathsf{W}}_{\alpha,k}(\mu^{\star})}_{(a):~{}\mbox{\scriptsize
first
stage}}+\underbrace{\widehat{\mathsf{W}}_{\alpha,k}(\mu^{\star})-\mathsf{W}_{\alpha}(\mu^{\star})}_{(b):~{}\mbox{\scriptsize
second stage}}$
$\displaystyle\quad\quad+\underbrace{\mathbb{E}_{\widehat{P}_{k}}[\widehat{\tau}_{k}(Z)(\ell(\theta(X);Y)-\widehat{\mu}_{k}(Z))\mid
I_{k}^{c}]}_{(c):~{}\mbox{\scriptsize debiasing term}}$ (B.4)

Term $(b)$ can be bounded using Proposition 8 as before. Without assuming
$\mu^{\star}\in\mathcal{H}$, recall that any bounded random variable taking
values in $[0,B]$ is sub-Gaussian with parameter $B^{2}/4$, so Proposition 8
yields

$\displaystyle|(b)|=\left|\widehat{\mathsf{W}}_{\alpha,k}(\mu^{\star})-\mathsf{W}_{\alpha}(\mu^{\star})\right|\leq
C\frac{B}{\alpha}\sqrt{\frac{\log(2/\delta)}{|I_{k}|}}~{}~{}\mbox{with
probability at least}~{}1-\delta.$

We first notice the following bound on term $(a)$.

$\displaystyle(a)=\left|\widehat{\mathsf{W}}_{\alpha,k}(\widehat{\mu}_{k})-\widehat{\mathsf{W}}_{\alpha,k}(\mu^{\star})\right|$
$\displaystyle\leq\frac{1}{\alpha}\sup_{\eta}\left|\mathbb{E}_{\widehat{P}_{k}}\left[\left({\widehat{\mu}_{k}(Z_{i})-\eta}\right)_{+}-\left({\mu^{\star}(Z_{i})-\eta}\right)_{+}\mid
I_{k}^{c}\right]\right|$
$\displaystyle\leq\frac{1}{\alpha}\mathbb{E}_{\widehat{P}_{k}}\left[\left|\widehat{\mu}_{k}(Z)-\mu^{\star}(Z)\right|\mid
I_{k}^{c}\right]\equiv\frac{1}{\alpha}\left\|{\widehat{\mu}_{k}(Z)-\mu^{\star}(Z)}\right\|_{L^{1}(\widehat{P}_{k}|I_{k}^{c})}.$
(B.5)

Next we will bound the debiasing term $(c)$ with the same quantity. We start
by observing Hoeffding’s inequality implies

$\displaystyle\left|\mathbb{E}_{\widehat{P}_{k}}[\widehat{\tau}_{k}(Z)(\ell(\theta(X);Y)-\widehat{\mu}_{k}(Z))\mid
I_{k}^{c}]-\mathbb{E}[\widehat{\tau}_{k}(Z)(\ell(\theta(X);Y)-\widehat{\mu}_{k}(Z))\mid
I_{k}^{c}]\right|\leq\frac{B}{\alpha}\sqrt{\frac{\log(2/\delta)}{|I_{k}|}}$

with probability at least $1-\delta$ because
$\widehat{\tau}_{k}\in\\{0,1/\alpha\\}$ and $\ell,\widehat{\mu}_{k}\in[0,B]$
almost surely. Then notice by definition of
$\mu^{\star}\equiv\mathbb{E}[\ell(\theta(X);Y)\mid Z]$ we know
$\mathbb{E}[\widehat{\tau}_{k}(Z)(\ell(\theta(X);Y)-\widehat{\mu}_{k}(Z))\mid
I_{k}^{c}]=\mathbb{E}[\widehat{\tau}_{k}(Z)(\mu^{\star}(Z)-\widehat{\mu}_{k}(Z))\mid
I_{k}^{c}]$ by conditioning on $Z$. Then we again invoke Hoeffding’s
inequality to argue with probability at least $1-\delta$,

$\displaystyle\left|\mathbb{E}[\widehat{\tau}_{k}(Z)(\mu^{\star}(Z)-\widehat{\mu}_{k}(Z))\mid
I_{k}^{c}]-\mathbb{E}_{\widehat{P}_{k}}[\widehat{\tau}_{k}(Z)(\mu^{\star}(Z)-\widehat{\mu}_{k}(Z))\mid
I_{k}^{c}]\right|\leq\frac{B}{\alpha}\sqrt{\frac{\log(2/\delta)}{|I_{k}|}}.$

Lastly we notice

$\displaystyle\left|\mathbb{E}_{\widehat{P}_{k}}[\widehat{\tau}_{k}(Z)(\mu^{\star}(Z)-\widehat{\mu}_{k}(Z))\mid
I_{k}^{c}]\right|$
$\displaystyle\leq\frac{1}{\alpha}\mathbb{E}_{\widehat{P}_{k}}[|\mu^{\star}(Z)-\widehat{\mu}_{k}(Z)|\mid
I_{k}^{c}]$
$\displaystyle\equiv\frac{1}{\alpha}\left\|{\mu^{\star}(Z)-\widehat{\mu}_{k}(Z)}\right\|_{L^{1}(\widehat{P}_{k}|I_{k}^{c})}$

because $\widehat{\tau}_{k}\in\\{0,1/\alpha\\}$. Hence, we conclude with
probability at least $1-2\delta$,

$|(c)|\leq\frac{1}{\alpha}\left\|{\mu^{\star}(Z)-\widehat{\mu}_{k}(Z)}\right\|_{L^{1}(\widehat{P}_{k}|I_{k}^{c})}+\frac{2B}{\alpha}\sqrt{\frac{\log(2/\delta)}{|I_{k}|}}.$

Thus we have shown with probability at least $1-3\delta$,

$|\widehat{\omega}_{\alpha,k}-\mathsf{W}_{\alpha}^{\star}|\leq\frac{2}{\alpha}\left\|{\mu^{\star}(Z)-\widehat{\mu}_{k}(Z)}\right\|_{L^{1}(\widehat{P}_{k}|I_{k}^{c})}+\frac{CB}{\alpha}\sqrt{\frac{\log(2/\delta)}{|I_{k}|}}.$

We now present two approaches to bounding the empirical $L^{1}$-norm of of the
error $\mu^{\star}(Z)-\widehat{\mu}_{k}(Z)$ for whether the model class
$\mathcal{H}$ is convex. Before we move on, notice the following identity that
is useful for both cases, which can be interpreted geometrically as the cosine
theorem in the $L^{2}(\widehat{P}_{k})$ space. For any two functions
$\mu,\tilde{\mu}:\mathcal{Z}\to[0,B]$,

$\left\|{\tilde{\mu}(Z)-\mu(Z)}\right\|_{L^{2}(\widehat{P}_{k})}=\Delta_{I_{k}}(\tilde{\mu})-\Delta_{I_{k}}(\mu)+2\mathbb{E}_{\widehat{P}_{k}}\left[(\ell(\theta(X);Y)-\mu(Z))\left(\tilde{\mu}(Z)-\mu(Z)\right)\right].$
(B.6)

####  B.3.1 Continuing proof of Theorem 3 with a non-convex model class

First note Hölder’s inequality implies

$\left\|{\mu^{\star}(Z)-\widehat{\mu}_{k}(Z)}\right\|_{L^{1}(\widehat{P}_{k}|I_{k}^{c})}\leq\left\|{\mu^{\star}(Z)-\widehat{\mu}_{k}(Z)}\right\|_{L^{2}(\widehat{P}_{k}|I_{k}^{c})},$

so it suffices to bound the $L^{2}$-norm. In order to use the identity B.6
with $\mu=\mu^{\star}$ and $\tilde{\mu}=\widehat{\mu}_{k}$, we notice by
definition of $\mu^{\star}\equiv\mathbb{E}[\ell(\theta(X);Y)\mid Z]$ that
$\mathbb{E}[(\ell(\theta(X);Y)-\mu^{\star}(Z))(\widehat{\mu}_{k}(Z)-\mu(Z))]=0$
for all $\tilde{\mu}$. Since
$(\ell(\theta(X);Y)-\mu^{\star}(Z))(\widehat{\mu}_{k}(Z)-\mu^{\star}(Z))$ is
almost surely bounded in $[-B^{2},B^{2}]$ and i.i.d., Hoeffding inequality
[112, Ch. 2] yields

$\displaystyle\left|\mathbb{E}_{\widehat{P}_{k}}[(\ell(\theta(X);Y)-\mu^{\star}(Z))(\widehat{\mu}_{k}(Z)-\mu^{\star}(Z))]\right|\leq
B^{2}\sqrt{\frac{2\log(2/\delta)}{|I_{k}|}}~{}~{}\mbox{with probability at
least}~{}1-\delta.$

Hence, the identity B.6 implies with probability at least $1-\delta$,

$\left\|{\mu^{\star}(Z)-\widehat{\mu}_{k}(Z)}\right\|_{L^{2}(\widehat{P}_{k}|I_{k}^{c})}^{2}\leq\left[\Delta_{I_{k}}(\widehat{\mu}_{k})-\Delta_{I_{k}}(\bar{\mu})\right]+\left[\Delta_{I_{k}}(\bar{\mu})-\Delta_{I_{k}}(\mu^{\star})\right]+B^{2}\sqrt{\frac{2\log(2/\delta)}{|I_{k}|}}$

Similarly, Hoeffding inequality implies with probability at least $1-\delta$,

$\displaystyle\Delta_{I_{k}}(\bar{\mu})-\Delta_{I_{k}}(\mu^{\star})$
$\displaystyle\leq\mathbb{E}(\ell(\theta(X);Y)-\bar{\mu}(Z))^{2}-\mathbb{E}(\ell(\theta(X);Y)-\mu^{\star}(Z))^{2}+B^{2}\sqrt{\frac{2\log(1/\delta)}{|I_{k}|}}$
$\displaystyle=\left\|{\bar{\mu}-\mu^{\star}}\right\|_{L^{2}}^{2}+B^{2}\sqrt{\frac{2\log(1/\delta)}{|I_{k}|}},$

where the equality follows by the definition of the conditional risk
$\mu^{\star}(Z)\equiv\mathbb{E}[\ell(\theta(X);Y)\mid Z]$. Hence, with
probability at least $1-2\delta$,

$\displaystyle\left\|{\mu^{\star}(Z)-\widehat{\mu}_{k}(Z)}\right\|_{L^{2}(\widehat{P}_{k}|I_{k}^{c})}$
$\displaystyle\leq\sqrt{\Delta_{I_{k}}(\widehat{\mu}_{k})-\Delta_{I_{k}}(\bar{\mu})+\left\|{\bar{\mu}-\mu^{\star}}\right\|_{L^{2}}^{2}+2B^{2}\sqrt{\frac{2\log(2/\delta)}{|I_{k}|}}}$
$\displaystyle\leq\sqrt{[\Delta_{I_{k}}(\widehat{\mu}_{k})-\Delta_{I_{k}}(\bar{\mu})]_{+}}+\left\|{\bar{\mu}-\mu^{\star}}\right\|_{L^{2}}+\sqrt{2}B\left(\frac{2\log(2/\delta)}{|I_{k}|}\right)^{1/4}.$

Therefore, we conclude that with probability at least $1-5\delta$,

$\left|\mathsf{W}_{\alpha}^{\star}-\widehat{\omega}_{\alpha,k}\right|\leq\frac{2}{\alpha}\left(\sqrt{[\Delta_{I_{k}}(\widehat{\mu}_{k})-\Delta_{I_{k}}(\bar{\mu})]_{+}}+\left\|{\bar{\mu}-\mu^{\star}}\right\|_{L^{2}}+CB\left(\frac{2\log(2/\delta)}{|I_{k}|}\right)^{1/4}\right).$

####  B.3.2 Continuing proof of Theorem 3 with a convex model class

First notice with probability at least $1-\delta$,

$\displaystyle\left\|{\mu^{\star}(Z)-\widehat{\mu}_{k}(Z)}\right\|_{L^{1}(\widehat{P}_{k}|I_{k}^{c})}$
$\displaystyle\leq\left\|{\bar{\mu}(Z)-\widehat{\mu}_{k}(Z)}\right\|_{L^{1}(\widehat{P}_{k}|I_{k}^{c})}+\left\|{\mu^{\star}(Z)-\bar{\mu}(Z)}\right\|_{L^{1}(\widehat{P}_{k})}$
$\displaystyle\leq\left\|{\bar{\mu}(Z)-\widehat{\mu}_{k}(Z)}\right\|_{L^{2}(\widehat{P}_{k}|I_{k}^{c})}+\left\|{\mu^{\star}(Z)-\bar{\mu}(Z)}\right\|_{L^{1}(\widehat{P}_{k})}$
$\displaystyle\leq\left\|{\bar{\mu}(Z)-\widehat{\mu}_{k}(Z)}\right\|_{L^{2}(\widehat{P}_{k}|I_{k}^{c})}+\left\|{\mu^{\star}(Z)-\bar{\mu}(Z)}\right\|_{L^{1}(P)}+\frac{B}{2}\sqrt{\frac{\log(1/\delta)}{|I_{k}|}},$

where the first inequality follows by the triangle inequality, the second by
Hölder’s inequality, and the last by Hoeffding inequality because
$\mu^{\star},\widehat{\mu}_{k}\in[0,B]$. The identity B.6 implies

$\left\|{\bar{\mu}(Z)-\widehat{\mu}_{k}(Z)}\right\|_{L^{2}(\widehat{P}_{k}|I_{k}^{c})}^{2}=\Delta_{I_{k}}(\widehat{\mu}_{k})-\Delta_{I_{k}}(\bar{\mu})+2\mathbb{E}_{\widehat{P}_{k}}[(\ell(\theta(X_{i});Y_{i})-\bar{\mu}(Z_{i}))(\widehat{\mu}_{k}(Z_{i})-\bar{\mu}(Z_{i}))].$

Since we assume the model class $\mathcal{H}$ is convex and
$\widehat{\mu}_{k}\in\mathcal{H}$, the first-order condition of
$\bar{\mu}\in\arg\min_{\mu\in\mathcal{H}}\mathbb{E}(\ell(\theta(X);Y)-\mu(Z))^{2}$
gives

$\mathbb{E}[(\ell(\theta(X);Y)-\bar{\mu}(Z))(\widehat{\mu}_{k}(Z)-\bar{\mu}(Z))\mid
Z,I_{k}^{\complement}]\leq 0,$

so Hoeffding inequality implies with probability at least $1-\delta$,

$\displaystyle\mathbb{E}_{\widehat{P}_{k}}[(\ell(\theta(X);Y)-\bar{\mu}(Z))(\widehat{\mu}_{k}(Z)-\bar{\mu}(Z))]\leq
B^{2}\sqrt{\frac{2\log(1/\delta)}{|I_{k}|}}.$ (B.7)

Hence, with probability at least $1-2\delta$,

$\displaystyle\left\|{\mu^{\star}(Z)-\widehat{\mu}_{k}(Z)}\right\|_{L^{1}(\widehat{P}_{k}|I_{k}^{c})}$
$\displaystyle\leq\left\|{\bar{\mu}-\mu^{\star}}\right\|_{L^{1}}+\sqrt{\Delta_{I_{k}}(\widehat{\mu}_{k})-\Delta_{I_{k}}(\bar{\mu})+2B^{2}\sqrt{\frac{2\log(1/\delta)}{|I_{k}|}}}+\frac{B}{2}\sqrt{\frac{\log(1/\delta)}{|I_{k}|}}$
$\displaystyle\leq\sqrt{[\Delta_{I_{k}}(\widehat{\mu}_{k})-\Delta_{I_{k}}(\bar{\mu})]_{+}}+\left\|{\bar{\mu}-\mu^{\star}}\right\|_{L^{1}}+CB\left(\frac{2\log(2/\delta)}{|I_{k}|}\right)^{1/4}.$

Therefore, we conclude that with probability at least $1-5\delta$,

$\left|\mathsf{W}_{\alpha}^{\star}-\widehat{\omega}_{\alpha,k}\right|\leq\frac{2}{\alpha}\left(\sqrt{[\Delta_{I_{k}}(\widehat{\mu}_{k})-\Delta_{I_{k}}(\bar{\mu})]_{+}}+\left\|{\bar{\mu}-\mu^{\star}}\right\|_{L^{1}}+CB\left(\frac{2\log(2/\delta)}{|I_{k}|}\right)^{1/4}\right).$

##  Appendix C Additional experiment details

In this section, we present additional experiments for the Functional Map of
the World (FMoW) dataset. Due to the ever-changing nature of aerial images and
the uneven availability of data from different regions, it is imperative that
ML models maintain good performance under temporal (learn from the past and
generalize to future) and spatial distribution shifts (learn from one region
and generalize to another). Without having access to the out-of-distribution
samples, our diagnostic raises awareness on brittleness of model performance
against subpopulation shifts.

###  C.1 Dataset Description

The original Functional Map of the World (FMoW) dataset by [33] consists of
over 1 million images from over 200 countries. We use a variant, FMoW-WILDS,
proposed by Koh et al. [68], which temporally groups observations to simulate
distribution shift across time. Each data point includes an RGB satellite
image $x$, and a corresponding label $y$ on the land / building use of the
image (there are 62 different classes). FMoW-WILDS splits data into non-
overlapping time periods: we train and validate models $\theta(\cdot)$ on data
collected from years 2002-2013, and simulate distribution shift by looking at
data collected during 2013-2018. Data collected during 2002-2013 (“in-
distribution”) is split into training ($n=$76,863), validation ($n=$19,915),
and test ($n=$11,327). Data collected during 2013-2018 (“out-of-distribution”)
is split into two sets: one consisting of observations from years 2013-2016
($n=$19,915), and another consisting of observations from years 2016-2018
($n=$22,108). All data splits contain images from a diverse array of
geographic regions. We evaluate the worst-case subpopulation performance on
in-distribution validation data, and study model performance under
distribution shift on data after 2016.

###  C.2 Models Evaluated

We consider _DenseNet_ models as reported by Koh et al. [68], including the
vanilla empirical risk minimization (ERM) model and models trained with
robustness interventions (IRM [6] method; Koh et al. [68] notes that ERM’s
performance closely match or outperform “robust” counterparts even under
distribution shift. We also evaluate _ImageNet_ pre-trained _DPN_ -68 model
from Miller et al. [79]. As separate experiments, we also consider _ResNet_
-18 and _VGG_ -11 from Miller et al. [79], and the results are reported in
C.6.

CLIP (Contrastive Language-Image Pre-training) is a newly proposed model pre-
trained on 400M image-text pairs, and has been shown to exhibit strong zero-
shot performance on out-of-distribution samples [86]. Although not
specifically designed for classification tasks, CLIP can be used for
classification by predicting the class whose encoded text is the closest to
the encoded image. We consider the weight-space ensembled _CLIP WiSE_ models
proposed in [118] as it is observed that these models exhibit robust behavior
on FMoW. _CLIP WiSE_ models are constructed by linearly combining the model
weights of _CLIP ViT-B16 Zeroshot model_ and _CLIP ViT-B16 FMoW end-to-end
finetuned_ model.

To illustrate the usage of our method, we choose the _CLIP WiSE_ model that
has similar ID validation accuracy as the _DenseNet_ Models. This turns out to
be putting $60\%$ weight on _CLIP ViT-B16 Zeroshot model_ and $40\%$ weight on
_CLIP ViT-B16 FMoW end-to-end finetuned_. _DenseNet_ Models have average ID
validation loss $2.4-2.8$, but _CLIP WiSE_ has average ID validation loss
$1.6$. To ensure fair comparison, we calibrate the temperature parameter such
that the average loss of _CLIP WiSE_ matches the worst average loss of the
models considered. We deliberately make _CLIP WiSE_ no better than any
_DenseNet_ Models, in the hope that out metric will recover its robustness
property.

# | Text Prompt
---|---
1 | “CLASSNAME”
2 | “a picture of a CLASSNAME.”
3 | “a photo of a CLASSNAME.”
4 | “an image of an CLASSNAME”
5 | “an image of a CLASSNAME in asia.”
6 | “an image of a CLASSNAME in africa.”
7 | “an image of a CLASSNAME in the americas.”
8 | “an image of a CLASSNAME in europe.”
9 | “an image of a CLASSNAME in oceania.”
10 | “satellite photo of a CLASSNAME”
11 | “satellite photo of an CLASSNAME”
12 | “satellite photo of a CLASSNAME in asia.”
13 | “satellite photo of a CLASSNAME in africa.”
14 | “satellite photo of a CLASSNAME in the americas.”
15 | “satellite photo of a CLASSNAME in europe.”
16 | “satellite photo of a CLASSNAME in oceania.”
17 | “an image of a CLASSNAME”
Table 2: Text prompts for CLIP text encoders

###  C.3 Flexibility of our metric

We implement Algorithm 1 by partitioning the ID validation data into two; we
estimate $\mu^{\star}(Z)$ using XGBoost on one sample, and estimate
$\mathsf{W}_{\alpha}(\cdot)$ at varying subpopulation size $\alpha$ on the
other. By switching the role of each split, our final estimator averages two
versions of $\widehat{\mathsf{W}}_{\alpha,k}(\widehat{\mu})$.

####  C.3.1 A less conservative Z

In Section 5, we report results when $Z$ is defined over all metadata
consisting of (longitude, latitude, cloud cover, region, year), as well as the
label $Y$. Defining subpopulations over such a wide range of variables may be
overly conservative in some scenarios, and to illustrate the flexibility of
our approach, we now showcase a more tailored definition of subpopulations.
Since FMoW-WILDS is specifically designed for spatiotemporal shifts, a natural
choice of $Z$ is to condition on (region, year). Motivated by our observation
that some classes are harder to predict than others (Figure 4(b)), we also
consider $Z$= (region, year, label $Y$). We plot our findings in Figure 6. If
we simply define $Z$= (year, region), the corresponding worst-case
subpopulation performance is less pessimistic. However, when we add labels to
$Z$, we again see a drastic decrease in the worst-case subpopulation
performance, and that CLIP WiSE-FT outperforms all other models by a
significant amount. This is consistent with our motivation in defining
subpopulations over labels; our procedure automatically takes into account the
interplay between class labels and spatiotemporal information.

Figure 6: In the left panel $Z$ = (year, region); in the right panel $Z$ =
(year, region, label $Y$). Here we take $Z$ to contain only spatial and
temporal information, a less conservative counterpart to the experiment
reported in the main text. We again see that introduction of labels in $Z$
drastically increase our metric, showing varying difficulty in learning
different labels.

####  C.3.2 Using semantics of the labels

Alternatively, we may wish to define subpopulations over rich natural language
descriptions on the input $X$. To illustrate the flexibility of our procedure
in such scenarios, we consider subpopulations defined over the semantic
meaning of the class names: CLIP-encoded class names using the 17 prompts
reported in Table 2. For comparison, we report the (estimated) worst-case
subpopulation performance (1.2) when we take $Z$ = (all metadata, encoded
labels) and (all metadata, label Y, encoded labels) in Figure 7. We observe
that in this case, the semantics of the class names do not contribute to
further deterioration in robustness, and the relative ordering across models
remains unchanged.

Figure 7: In the left panel, $Z$ = (all meta, encoded labels); in the right
panel, $Z$ = (all meta, label Y, encoded labels). We see that in this case no
significant difference is introduced when semantics of the class names are
included.

###  C.4 Analysis of spatiotemporal distribution shift

The significant performance drop in the Africa region on data collected from
2016-2018 was also observed in [68, 118]. In Figures 8-9, we plot the number
of samples collected from Africa over data splits. In particular, we observe a
large number of single-unit and multi-unit residential instances emerge in the
OOD data. Data collection systems are often biased against the African
continent—often as a result of remnants of colonialism—and addressing such
bias is an important topic of future research.

Figure 8: Instances by class, ID 2002-2013, Africa Figure 9: Instances by
class, test 2016-2018, Africa

###  C.5 Estimator of model loss

One potential limitation of our approach is $\hat{h}$ does not always estimate
the tail losses accurately, and this is important because our approach
precisely is designed to counter ML models that perform poorly on tail
subpopulations. Figure 10 plots a histogram of model losses and the estimated
conditional risk $\hat{h}$ for _DenseNet_ ERM and _CLIP WiSE_ , where the
y-axis is plotted on a log-scale. It is clear that _DenseNet_ ERM has more
extreme losses compared to the _CLIP WiSE_ model, suggesting that at least
part of the reason why _DenseNet_ ERM suffers poor loss on subpopulations: it
is overly confident when it’s incorrect. While a direct comparison is not
appropriate since the conditional risk $\mu(Z)$ represent _smoothed_ losses,
we observe that naive estimators of $\mu(\cdot)$ may consistently
underestimate. In this particular instance, since the extent of
underestimation is more severe for _ImageNet_ pre-trained models, our
experiments are fortuitously providing an even more conservative comparison
between the two model classes, instilling confidence in the relative
robustness of the _CLIP WiSE_ model.

Figure 10: Histograms of model losses and fitted losses $\hat{h}$. Y-axis
count is plotted in log-scale. For _DenseNet_ ERM model, fitted $\hat{h}$
underestimates the extreme losses (right-tail).

Alternatively, we can directly define the worst-case subpopulation performance
(1.2) using the 0-1 loss. The discrete nature of the 0-1 loss pose some
challenges in estimating $\mu(\cdot)$. While we chose to focus on the cross
entropy loss that aligns with model training, we leave a thorough study of 0-1
loss to future work.

###  C.6 Additional comparisons

We use the ensembled _CLIP WiSE_ model constructed by averaging the network
weights of _CLIP zero-shot_ and _CLIP finetuned_ models. So far, we used
proportion $\lambda=0.4$ to match the ID validation accuracy of _CLIP WiSE_ to
that of _DenseNet_ models and _DPN_ -68. In this subsection, we provide
alternative choices:

  1. 1.

$\lambda=0.24$ to match ID accuracy of _ResNet_ -18 of $47\%$

  2. 2.

$\lambda=0.27$ to match ID accuracy of _VGG_ -11 of $51\%$.

Similar to _DPN_ -68, _ResNet_ -18 and _VGG_ -11 are _ImageNet_ pretrained
models fine-tuned on FMoW as evaluated by Miller et al. [79]. We refer to the
two _CLIP WiSE_ models as _CLIP WiSE_ -24 and _CLIP WiSE_ -27 respectively,
and report all model performances below. Again, we observe that our approach
successfully picks out the more robust _CLIP WiSE_ models, in contrast to the
non-robust models chosen by ID accuracy or ID loss.

| ID, 2002-2013 | OOD, 2016-2018
---|---|---
Model | Accuracy | Loss | $\mathsf{W}_{0.10}$ | Accuracy | Loss
_CLIP WiSE_ -24 | 0.47 | 2.84 | 3.47 | 0.45 | 2.85
_ResNet_ -18 | 0.48 | 2.84 | 5.05 | 0.40 | 3.36
_CLIP WiSE_ -27 | 0.51 | 3.07 | 3.54 | 0.48 | 3.08
_VGG_ -11 | 0.51 | 3.06 | 6.07 | 0.45 | 3.68
Table 3: Additional experiments showcasing our approach successfully
identifies more robust models.

##  Appendix D Proof of equivalence results

In this section, we discuss in detail how our notion of generalized worst-case
subpopulation performance is closely connected to distributional robustness
and coherent risk measures. The reader is recommended to refer to Section 6 of
the lectures notes by Shapiro et al. [103] for more detail.

###  D.1 Proof of Theorem 5

To see the first claim, fix a nonempty class $\Lambda$ of probability measures
on $(0,1]$. Notice $\mathsf{W}_{\Lambda}(\cdot)$ is proper because $\Lambda$
is nonempty and $\mathsf{W}_{\Lambda}(0)=0<\infty$. It is lower semi-
continuous because it is a pointwise supremum of
$\mu\mapsto\int_{(0,1]}\mathsf{W}_{\alpha}(\mu)d\lambda(\alpha)$, which are
lower semi-continuous. Coherence can be shown by verifying the definition.
Clearly, $\mathsf{W}_{\Lambda}$ is law-invariant because it is defined using
only $\mathsf{W}_{\alpha}(\cdot)$ which is law-invariant.

For the converse relation, recall the discussion preceding Lemma 3. The
biconjugacy relation (6.4) gives the variational representation

$\rho(\mu)=\sup_{L\in\mathop{\rm dom}\rho^{\star}}\mathbb{E}_{P}[L\mu].$

Since $\rho$ is law-invariant, we have the tautological reformulation

$\rho(\mu)=\sup_{\mu^{\prime}\in\mathcal{L}^{k}}\left\\{\rho(\mu^{\prime}):\mu^{\prime}\stackrel{{\scriptstyle
d}}{{=}}\mu\right\\}=\sup_{L\in\mathop{\rm
dom}\rho^{\star},\mu^{\prime}\in\mathcal{L}^{k}}\left\\{\mathbb{E}_{P}[L\mu^{\prime}]:\mu^{\prime}\stackrel{{\scriptstyle
d}}{{=}}\mu\right\\}.$ (D.1)

Next, we use a generalization of the Hardy–Littlewood inequality. Let
$P_{1-\alpha}^{-1}(\mu)$ denote the $(1-\alpha)$-quantile of the random
variable $\mu\in\mathcal{L}^{k}$.

######  Lemma 11 (Shapiro et al. [102, Lemma 6.25]).

Suppose $P$ is nonatomic. For $\mu\in\mathcal{L}^{k}$ and
$L\in\mathcal{L}^{k_{*}}$,

$\sup_{\mu^{\prime}\in\mathcal{L}^{k}}\left\\{\mathbb{E}_{P}[L\mu^{\prime}]:\mu^{\prime}\stackrel{{\scriptstyle
d}}{{=}}\mu\right\\}=\int_{0}^{1}P_{1-t}^{-1}(L)P_{1-t}^{-1}(\mu)dt.$

In particular, $\mathbb{E}_{P}[\mu]=\int_{0}^{1}P_{1-t}^{-1}(\mu)dt$.

We use the following elementary identity to rewrite $P_{1-t}^{-1}(L)$ in the
preceding display. We defer its proof to the end of the subsection.

######  Lemma 12.

For any random variable $L\in\mathcal{L}^{k_{*}}(P)$,

$P_{1-t}^{-1}(L)=\int_{t}^{1}\alpha^{-1}d\lambda_{L}(\alpha)~{}~{}\mbox{for}~{}~{}t\in(0,1),$

where $\lambda_{L}(\cdot)$ is a probability distribution on $(0,1]$, if and
only if $L\geq 0$ $P$-a.s., $\mathbb{E}_{P}[L]=1$,
$\lambda_{L}(\alpha)=\mathbb{E}_{P}\left({L-P_{1-\alpha}^{-1}(L)}\right)_{+}$
for $\alpha\in(0,1)$, and $\lambda_{L}(1)=1$.

For any fixed $L\in\mathcal{L}^{k_{*}}$, conclude

$\sup_{\mu^{\prime}\in\mathcal{L}^{k}}\left\\{\mathbb{E}_{P}[L\mu^{\prime}]:\mu^{\prime}\stackrel{{\scriptstyle
d}}{{=}}\mu\right\\}=\int_{0}^{1}P_{1-t}^{-1}(L)P_{1-t}^{-1}(\mu)dt=\int_{0}^{1}\int_{(t,1]}\alpha^{-1}d\lambda_{L}(\alpha)~{}P_{1-t}^{-1}(\mu)dt.$

Applying Fubini-Tonelli to the RHS of the preceding display,

$\displaystyle\sup_{\mu^{\prime}\in\mathcal{L}^{k}}\left\\{\mathbb{E}_{P}[L\mu^{\prime}]:\mu^{\prime}\stackrel{{\scriptstyle
d}}{{=}}\mu\right\\}=\int_{(0,1]}\frac{1}{\alpha}\int_{0}^{\alpha}P_{1-t}^{-1}(\mu)dt~{}d\lambda_{L}(\alpha)=\int_{(0,1]}\mathsf{W}_{\alpha}(\mu)d\lambda_{L}(\alpha),$

where the final equality follows from the change-of-variables reformulation
(2.2). To obtain the representation (D.1), we take the supremum over
$L\in\mathop{\rm dom}\rho^{\star}$

$\displaystyle\rho(\mu)=\sup_{L\in\mathop{\rm
dom}\rho^{\star}}\int_{(0,1]}\mathsf{W}_{\alpha}(\mu)d\lambda_{L}(\alpha).$

####  D.1.1 Proof of Lemma 12

Before proving the equivalence, we first show the following identity for any
$L$ and $\lambda_{L}$.

$\int_{0}^{\alpha}[h(t)-h(\alpha)]\,\mathrm{d}t=\mathbb{E}_{P}(L-P_{1-\alpha}^{-1}(L))_{+}-\lambda_{L}(\alpha),$
(D.2)

where we define
$h(t):=P_{1-t}^{-1}(L)-\int_{(t,1]}a^{-1}\,\mathrm{d}\lambda_{L}(a)$ for
convenience. We separately consider the two differences in the integrand

$\displaystyle\int_{0}^{\alpha}[h(t)-h(\alpha)]\,\mathrm{d}t=\int_{0}^{\alpha}P_{1-t}^{-1}(L)-P_{1-\alpha}^{-1}(L)-\left(\int_{t}^{1}a^{-1}\,\mathrm{d}\lambda_{L}(a)-\int_{\alpha}^{1}a^{-1}\,\mathrm{d}\lambda_{L}(a)\right)\,\mathrm{d}t.$

For $t\in(0,\alpha)$

$\begin{split}\int_{0}^{\alpha}\left[P_{1-t}^{-1}(L)-P_{1-\alpha}^{-1}(L)\right]\,\mathrm{d}t&=\int_{0}^{1}\left[P_{1-t}^{-1}(L)-P_{1-\alpha}^{-1}(L)\right]_{+}\,\mathrm{d}t\\\
&=\int_{0}^{1}P_{1-t}^{-1}\left(\left[L-P_{1-\alpha}^{-1}(L)\right]_{+}\right)\,\mathrm{d}t=\mathbb{E}_{P}\left[L-P_{1-\alpha}^{-1}(L)\right]_{+},\end{split}$

and

$\int_{0}^{\alpha}\left[\int_{(t,1]}a^{-1}\,\mathrm{d}\lambda_{L}(a)-\int_{(\alpha,1]}a^{-1}\,\mathrm{d}\lambda_{L}(a)\right]\,\mathrm{d}t=\int_{(0,\alpha]}\left(\int_{0}^{a}\,\mathrm{d}t\right)a^{-1}\,\mathrm{d}\lambda_{L}(a)=\lambda_{L}(\alpha).$

Now we show the “if” part. It is clear that $\lambda_{L}$ is nondecreasing, as
$P_{1-\alpha}^{-1}($) is nonincreasing in $\alpha$, and
$\lambda_{L}(\alpha)\leq\mathbb{E}_{P}[L]=1=\lambda_{L}(1)$ because $L\geq 0$
$P$-a.s. We know $\lambda_{L}$ is right-continuous because
$P_{1-\alpha}^{-1}(\cdot)$ is left continuous in $\alpha$. Then we conclude
$\lambda_{L}$ is a probability distribution on $(0,1]$ by noticing
$\lim_{\alpha\downarrow 0}\lambda_{L}(\alpha)=0$ because

$\displaystyle\lambda_{L}(\alpha)=\left\|{\left({L-P_{1-\alpha}^{-1}(L)}\right)_{+}}\right\|_{1}$
$\displaystyle\leq\left\|{\left({L-P_{1-\alpha}^{-1}(L)}\right)_{+}}\right\|_{k_{*}}$
$\displaystyle\leq\left\|{L}\right\|_{k_{*}}\mathbb{P}\\{L>P_{1-\alpha}^{-1}(L)\\}^{1/k_{*}}$
$\displaystyle\leq\left\|{L}\right\|_{k_{*}}\alpha^{1/k_{*}},$

where the first inequality follows by Hölder’s inequality, the second because
$\left({L-P_{1-\alpha}^{-1}(L)}\right)_{+}=(L-P_{1-\alpha}^{-1}(L))\textbf{1}\\{L>P_{1-\alpha}^{-1}(L)\\}$
and $L\geq 0$ $P$-a.s., and the last by definition of $P_{1-\alpha}^{-1}(L)$.

By the definition of $\lambda_{L}$, the RHS of Identity (D.2) is zero, so
$\alpha h(\alpha)=\int_{0}^{\alpha}h(t)\,\mathrm{d}t$. In particular,
$h(\cdot)$ is differentiable in $(0,1)$ since
$\alpha\mapsto\int_{0}^{\alpha}h(t)\,\mathrm{d}t$ is differentiable. Taking
derivatives on both sides of the preceding display, we have $\alpha
h^{\prime}(\alpha)+h(\alpha)=h(\alpha)$ and in particular,
$h^{\prime}(\alpha)=0$ for $\alpha\in(0,1]$. To show that $h(t)\equiv h$ is
uniformly zero, notice

$h=\int_{0}^{1}h(t)\,\mathrm{d}t=\int_{0}^{1}\left[P_{1-t}^{-1}(L)-\int_{t}^{1}a^{-1}\,\mathrm{d}\lambda_{L}(a)\right]\,\mathrm{d}t=\mathbb{E}_{P}L-1=0,$

where we used Fubini-Tonelli and $\lambda_{L}(1)=1$ in the third equality.

Next, we show the “only if” part. Clearly the LHS of Identity (D.2) is zero,
so
$\lambda_{L}(\alpha)=\mathbb{E}_{P}\left({L-P_{1-\alpha}^{-1}(L)}\right)_{+}$.
We know $L\geq 0$ $P$-a.s. because $\lambda_{L}\geq 0$. Lastly,
$\mathbb{E}_{P}[L]=\int_{0}^{1}P_{1-t}^{-1}(L)dt=\int_{0}^{1}\int_{t}^{1}\alpha^{-1}d\lambda_{L}(\alpha)dt=1$
by Fubini-Tonelli.

###  D.2 Proof of Proposition 6

First, to see the duality result, notice the minimax theorem and Hölder’s
inequality imply

$\begin{split}\rho_{k}(\mu)&=\inf_{\eta\in\mathbb{R}}\\{\alpha^{-1}||(\mu-\eta)_{+}||_{k}+\eta\\}\\\
&=\inf_{\eta\in\mathbb{R}}\sup\\{\mathbb{E}_{P}[L(\mu-\eta)]+\eta:L\geq
0,\left\|{L}\right\|_{k_{*}}\leq\alpha^{-1}\\}\\\
&=\sup\left\\{\inf_{\eta\in\mathbb{R}}\\{\mathbb{E}_{P}[L(\mu-\eta)]+\eta\\}:L\geq
0,\left\|{L}\right\|_{k_{*}}\leq\alpha^{-1}\right\\}\\\
&=\sup_{L}\left\\{\mathbb{E}_{P}[L\mu]:L\geq
0,\mathbb{E}_{P}[L]=1,\left\|{L}\right\|_{k_{*}}\leq\alpha^{-1}\right\\}.\end{split}$

This means $\mathop{\rm dom}\rho^{\star}_{k}=\\{L\in\mathcal{L}^{k}:L\geq
0,\mathbb{E}_{P}[L]=1,\left\|{L}\right\|_{k_{*}}\leq\alpha^{-1}\\}$. Theorem 5
implies

$\rho_{k}(\mu)=\sup\left\\{\int_{(0,1]}\mathsf{W}_{\alpha}(\mu)d\lambda_{L}(\alpha):\lambda_{L}(\alpha)=\mathbb{E}_{P}\left({L-P_{1-\alpha}^{-1}(L)}\right)_{+},L\geq
0,\mathbb{E}_{P}[L]=1,\left\|{L}\right\|_{k_{*}}\leq\alpha^{-1}\right\\}.$

Lemma 12 further implies
$\\{L,\lambda_{L}:\lambda_{L}(\alpha)=\mathbb{E}_{P}\left({L-P_{1-\alpha}^{-1}(L)}\right)_{+},L\geq
0,\mathbb{E}_{P}[L]=1\\}=\\{L,\lambda_{L}\in\Delta((0,1]):P_{1-t}^{-1}(L)=\int_{t}^{1}\alpha^{-1}d\lambda_{L}(\alpha)\\}$.
Notice $P_{1-t}^{-1}(L)=\int_{t}^{1}\alpha^{-1}d\lambda_{L}(\alpha)$ implies

$\left\|{L}\right\|_{k_{*}}^{k_{*}}=\int_{\Omega}L^{k_{*}}d\mathbb{P}=\int_{0}^{1}(P_{1-t}^{-1}(L))^{k_{*}}dt=\int_{0}^{1}\left(\int_{t}^{1}\alpha^{-1}d\lambda_{L}(\alpha)\right)^{k_{*}}dt,$

so

$\rho_{k}(\mu)=\sup\left\\{\int_{(0,1]}\mathsf{W}_{\alpha}(\mu)d\lambda_{L}(\alpha):\lambda_{L}\in\Lambda_{k},P_{1-t}^{-1}(L)=\int_{t}^{1}\alpha^{-1}d\lambda_{L}(\alpha)\right\\}.$

Lastly, since $L$ can always be defined based on any
$\lambda_{L}\in\Delta((0,1])$, we drop the last equality and conclude

$\rho_{k}(\mu)=\sup\left\\{\int_{(0,1]}\mathsf{W}_{\alpha}(\mu)d\lambda_{L}(\alpha):\lambda_{L}\in\Lambda_{k}\right\\}=\mathsf{W}_{\Lambda_{k}}(\mu).$

##  Appendix E Certificate of robustness

Instead of estimating the worst-case subpopulation performance for a fixed
subpopulation size $\alpha$, it may be natural to posit a level of acceptable
performance (upper bound $\bar{\ell}$ on the loss) and study $\alpha^{\star}$,
the smallest subpopulation size (2.1) over which the model $\theta(\cdot)$ can
guarantee acceptable performance. Our plug-in estimator $\widehat{\alpha}$
given in Eq. (3.4) enjoys similar concentration guarantees as given above. The
following theorem—whose proof we give in Appendix E.1—states that the true
$\alpha^{\star}$ is either close to our estimator $\widehat{\alpha}$ or it is
sufficiently small, certifying the robustness of the model against
subpopulation shifts.

######  Theorem 9.

Let Assumption D hold, let $U(\delta)>0$ be such that for any fixed
$\alpha\in(0,1]$,
$|\widehat{\mathsf{W}}_{\alpha,k}(\widehat{\mu})-\mathsf{W}_{\alpha}(\mu^{\star})|\leq
U(\delta)/\alpha$ with probability at least $1-\delta$. Then given any
$\underline{\alpha}\in(0,1]$, either $\alpha^{\star}<\underline{\alpha}$, or

$\small\left|\frac{\alpha^{\star}}{\widehat{\alpha}}-1\right|\leq\frac{U(\delta)}{\widehat{\mathbb{E}}\left[\widehat{\mu}(Z)-\widehat{P}_{1-\underline{\alpha}\wedge\widehat{\alpha}}^{-1}\left(\widehat{\mu}(Z)\right)\right]_{+}}$

with probability at least $1-\delta$, where $\widehat{\mathbb{E}}$ and
$\widehat{P}^{-1}_{1-\alpha}$ denote the expectation and the
$(1-\alpha)$-quantile under the empirical probability measure induced by
$I_{k}$.

Our approach simultaneously provides localized Rademacher complexity bounds
and dimension-free guarantees. Our bound becomes large as
$\underline{\alpha}\to 0$ and we conjecture this to be a fundamental
difficulty as the worst-case subpopulation performance (1.2) focuses on
$\alpha$-faction of the data.

###  E.1 Proof of Theorem 9

For ease of notation, we suppress any dependence on the prediction model
$\theta(X)$ under evaluation. Consider any $\alpha_{1},\alpha_{2}\in(0,1]$
with $\alpha_{1}<\alpha_{2}$. Denote by $\widehat{\mathbb{P}}$ the empirical
probability measure induced by $(Z_{i}:i\in I_{k})$. For convenience denote
$\xi_{1}:=\widehat{P}_{1-\alpha_{1}}^{-1}(\widehat{\mu}(Z))$ and
$\xi_{2}:=\widehat{P}_{1-\alpha_{2}}^{-1}(\widehat{\mu}(Z))$, so
$\xi_{1}\geq\xi_{2}$ and
$\widehat{\mathsf{W}}_{\alpha_{1}}(\widehat{\mu})\geq\widehat{\mathsf{W}}_{\alpha_{2}}(\widehat{h})$.
Notice that

$\begin{split}\widehat{\mathbb{E}}[\widehat{\mu}(Z)-\xi_{2}]_{+}-\widehat{\mathbb{E}}[\widehat{\mu}(Z)-\xi_{1}]_{+}&=\widehat{\mathbb{E}}[\widehat{\mu}(Z)-\xi_{2};\widehat{\mu}(Z)>\xi_{2}]-\widehat{\mathbb{E}}[\widehat{\mu}(Z)-\xi_{1};\widehat{\mu}(Z)\geq\xi_{1}]\\\
&=\underbrace{\widehat{\mathbb{E}}[\widehat{\mu}(Z)-\xi_{1};\xi_{2}<\widehat{\mu}(Z)<\xi_{1}]}_{\leq
0}+(\xi_{1}-\xi_{2})\underbrace{\widehat{\mathbb{P}}(\widehat{\mu}(Z)>\xi_{2})}_{\leq\alpha_{2}}\\\
&\leq(\xi_{1}-\xi_{2})\alpha_{2}.\end{split}$

Hence, by Lemma 9,

$\widehat{\mathsf{W}}_{\alpha_{1}}(\widehat{\mu})-\widehat{\mathsf{W}}_{\alpha_{2}}(\widehat{\mu})=\left(\frac{\widehat{\mathbb{E}}[\widehat{\mu}(Z)-\xi_{1}]_{+}}{\alpha_{1}}+\xi_{1}\right)-\left(\frac{\widehat{\mathbb{E}}[\widehat{\mu}(Z)-\xi_{2}]_{+}}{\alpha_{2}}+\xi_{2}\right)\geq\frac{\widehat{\mathbb{E}}[\widehat{\mu}(Z)-\xi_{1}]_{+}}{\alpha_{1}\alpha_{2}}(\alpha_{2}-\alpha_{1}),$

meaning

$|\alpha_{1}-\alpha_{2}|\leq\frac{\alpha_{1}\alpha_{2}|\widehat{\mathsf{W}}_{\alpha_{1}}(\widehat{\mu})-\widehat{\mathsf{W}}_{\alpha_{2}}(\widehat{\mu})|}{\widehat{\mathbb{E}}[\widehat{\mu}(Z)-\widehat{P}^{-1}_{1-\alpha_{1}}(\widehat{\mu}(Z))]_{+}}.$

Now suppose $\alpha^{\star}\geq\underline{\alpha}$. Notice
$\mathsf{W}_{\alpha}$ and $\widehat{\mathsf{W}}_{\alpha,k}$ are continuous and
nonincreasing in $\alpha$, so the definitions (2.1) and (3.4) imply
$\mathsf{W}_{\alpha^{\star}}(\mu^{\star})=\bar{\ell}=\widehat{\mathsf{W}}_{\widehat{\alpha}}(\widehat{\mu})$.
Plugging $\widehat{\alpha}$ and $\alpha^{\star}$ into the inequality above, we
know with probability at least $1-\delta$,

$|\alpha^{\star}-\widehat{\alpha}|\leq\frac{\widehat{\alpha}\alpha^{\star}\left|\widehat{\mathsf{W}}_{\alpha^{\star}}(\widehat{\mu})-{\mathsf{W}}_{\alpha^{\star}}(\mu^{\star})\right|}{\widehat{\mathbb{E}}\left[\widehat{\mu}(Z)-\widehat{P}_{1-\alpha^{\star}\wedge\widehat{\alpha}}^{-1}\left(\widehat{\mu}(Z)\right)\right]_{+}}\leq\frac{\widehat{\alpha}U(\delta)}{\widehat{\mathbb{E}}\left[\widehat{\mu}(Z)-\widehat{P}_{1-\underline{\alpha}\wedge\widehat{\alpha}}^{-1}\left(\widehat{\mu}(Z)\right)\right]_{+}}.$


# Algorithms for Weak Optimal Transport with an Application to Economics

François-Pierre Paty
CREST-ENSAE
Institut Polytechique de Paris
<EMAIL_ADDRESS>
&Philippe Choné
CREST-ENSAE
Institut Polytechique de Paris
<EMAIL_ADDRESS>
&Francis Kramarz
CREST-ENSAE
Institut Polytechique de Paris
<EMAIL_ADDRESS>
###### Abstract

The theory of weak optimal transport (WOT), introduced by [23], generalizes
the classic Monge-Kantorovich framework by allowing the transport cost between
one point and the points it is matched with to be nonlinear. In the so-called
barycentric version of WOT, the cost for transporting a point $x$ only depends
on $x$ and on the barycenter of the points it is matched with. This
aggregation property of WOT is appealing in machine learning, economics and
finance. Yet algorithms to compute WOT have only been developed for the
special case of quadratic barycentric WOT, or depend on neural networks with
no guarantee on the computed value and matching. The main difficulty lies in
the transportation constraints which are costly to project onto. In this
paper, we propose to use mirror descent algorithms to solve the primal and
dual versions of the WOT problem. We also apply our algorithms to the variant
of WOT introduced by [13] where mass is distributed from one space to another
through unnormalized kernels (WOTUK). We empirically compare the solutions of
WOT and WOTUK with classical OT. We illustrate our numerical methods to the
economic framework of [12], namely the matching between workers and firms on
labor markets.

##  1 Introduction

Over the past few years, optimal transport (OT) has gained importance in the
machine learning community as a useful tool to analyze data, with applications
to various domains such as graphics [37, 8], imaging [34, 15], generative
models [5, 35], biology [25, 36], NLP [24, 3], finance [7, 22, 1] or economics
[21, 20, 30].

The key in making the optimal transport approach work in these applications
lies in the different forms of regularization added to the classical optimal
transport problem. Although many different types of regularizations have been
proposed in the literature, among which the celebrated entropic regularization
[14], most of these correspond to penalized versions of the Monge-Kantorovich
problem, possibly with relaxed [18, 11] or tightened [7, 32] constraints.

These variants of OT are unable to capture the aggregation property that plays
a crucial role in some applications [6, 12]. In these applications, the cost
for matching points $x$ and $y$ does not only depends on $x$ and $y$ through a
cost function, but also on all the other points $y^{\prime}$ matched with $x$.
This idea was formalized by [23, 4] and gave rise to the notion of weak
optimal transport, which generalizes OT. Even more recently, and guided by the
economic application of [12], [13] proposed to relax the WOT problem and
defined the notion of weak optimal transport with unnormalized kernel,
henceforth WOTUK.

In the discrete setting, when the cost function is convex, weak optimal
transport (WOT) defines a convex optimization problem over the transportation
polytope [33]. Nevertheless, algorithms to compute WOT when the measures are
discrete have only been proposed in the special case of quadratic barycentric
WOT [9] (see subsection 2.2 for a precise definition). In their recent
preprint, [29] propose to use neural networks to approximate the WOT problem,
but do not not provide guarantees for their optimization procedure.

#### Contributions

The goal of this paper is twofold: first, we present recent optimal transport
problems of interest in economics, the so-called WOT and WOTUK problems, to
the machine learning community; second, we introduce algorithms to solve these
problems numerically, when the considered measures are discrete. We compare
these variants of OT with classical OT and apply our algorithms to the
economics framework of [12].

#### Structure of the paper

We begin in section 2 by reminders on optimal transport and weak optimal
transport problems, which we introduce in the context of the economic
application of [12]. We continue with the definition and the key properties of
weak optimal transport with unnormalized kernel in section 3. We present our
algorithms in section 4, and conclude the paper in section 5 with numerical
experiments and applications in economics.

##  2 From Optimal Transport to Weak Optimal Transport

###  2.1 Optimal Transport and Matching

A classic problem in labor economics [28, 26, 17] is to understand the
matching between workers and firms, i.e. to explain why workers work in their
employing firms, and conversely, why firms hire some employees and not others.
Optimal transport has been used in the economics literature [30, 20] to model
workers-to-firms matching.

Firms differ in technologies and workers differ in skills. Let
$\mathcal{X}\subset\mathbb{R}^{p}$, $p\geq 1$, denote the set of firms’ types
(or technologies) and likewise, let $\mathcal{Y}\subset\mathbb{R}^{q}$, $q\geq
1$, denote the set of workers’ types in the economy. Given a probability
distribution of firm types $\mu\in\mathscr{P}(\mathcal{X})$ and a probability
distribution of worker types $\nu\in\mathscr{P}(\mathcal{Y})$, a coupling
$\pi\in\mathscr{P}(\mathcal{X}\times\mathcal{Y})$ of $\mu$ and $\nu$
represents the matching between firms and workers, in the sense that
$\pi(A\times B)$, for $A\subset\mathcal{X},B\subset\mathcal{Y}$ Borel sets, is
the proportion of firms whose type is in $A$ that employ a worker whose type
is in $B$. The primal problem is to maximize the total output in the economy:

$\displaystyle\operatorname{OT}(\mu,\nu)\operatorname{\overset{\operatorname{def}}{=}}\sup_{\pi\in\Pi(\mu,\nu)}\iint_{\mathcal{X}\times\mathcal{Y}}F(x,y)\,d\pi(x,y)$
(1)

where $\Pi(\mu,\nu)$ is the set of all couplings between $\mu$ and $\nu$, i.e.

$\Pi(\mu,\nu)=\Big{\\{}\pi\in\mathscr{P}(\mathcal{X}\times\mathcal{Y})\,,\,\forall
A\subset\mathcal{X},B\subset\mathcal{Y}\text{
Borel},\pi(A\times\mathcal{Y})=\mu(A),\pi(\mathcal{X}\times
B)=\nu(B)\Big{\\}},$

and $F:\mathcal{X}\times\mathcal{Y}\to\mathbb{R}$ is the production function,
i.e. $F(x,y)$ is the output (in $\$$) produced by a worker of type
$y\in\mathcal{Y}$ working in a firm of type $x\in\mathcal{X}$.

Problem (1) corresponds to the definition of the Kantorovich [27] problem in
optimal transport with cost function $-F$. It admits the following dual
formulation:

$\displaystyle\operatorname{OT}(\mu,\nu)=\inf_{\begin{subarray}{c}\chi\in\mathcal{C}(\mathcal{X}),\varphi\in\mathcal{C}(\mathcal{Y})\\\
\chi\oplus\varphi\geq F\end{subarray}}\int\chi\,d\mu+\int\varphi\,d\nu$ (2)

where for $\mathcal{C}(\mathcal{X})$ (resp. $\mathcal{C}(\mathcal{Y})$) is the
set of real continuous functions over $\mathcal{X}$ (resp. over
$\mathcal{Y}$), and
$\chi\oplus\varphi\in\mathcal{C}(\mathcal{X}\times\mathcal{Y})$ is the
function $\chi\oplus\varphi:(x,y)\mapsto\chi(x)+\varphi(y)$. In the labor
market context, $\chi(x)$ and $\varphi(y)$ represent respectively the profit
of firms with type $x$ and the wage of workers with type $y$.

###  2.2 Weak Optimal Transport

The Kantorovich problem (1) can be rewritten as:

$\displaystyle\sup_{\pi\in\Pi(\mu,\nu)}\int_{\mathcal{X}}\left[\int_{\mathcal{Y}}F(x,y)\,d\pi_{x}(y)\right]d\mu(x)$
(3)

where $\left(\pi_{x}\right)_{x\in\mathcal{X}}\subset\mathscr{P}(\mathcal{Y})$
is the ($\mu$-almost surely unique) probability kernel that allows to
disintegrate $\pi$ with respect to $\mu$ as $d\pi(x,y)=d\mu(x)d\pi_{x}(y)$. In
other words, $\pi_{x}$ is the law of $Y|X=x$ when $(X,Y)\sim\pi$.

Gozlan et al. [23] introduce the weak optimal transport problem as

$\displaystyle\operatorname{WOT}(\mu,\nu)\operatorname{\overset{\operatorname{def}}{=}}\sup_{\pi\in\Pi(\mu,\nu)}\int_{\mathcal{X}}\mathcal{F}(x,\pi_{x})\,d\mu(x)$
(4)

where $\mathcal{F}:\mathcal{X}\times\mathscr{P}(\mathcal{Y})\to\mathbb{R}$,
i.e. $\mathcal{F}(x,p)$ now denotes the production (in $\$$) of a firm type
$x\in\mathcal{X}$ hiring employees with distribution
$p\in\mathscr{P}(\mathcal{Y})$. The classic Kantorovich problem is the special
case of WOT where $\mathcal{F}(x,p)=\int_{\mathcal{Y}}F(x,y)\,dp(y)$.

In our economic context, the probability kernel $\pi_{x}$ represents the
distribution of types among the workers hired by firms with types $x$. The
economic interpretation of the reworded problem (3) is that the output
produced by a firm of type $x\in\mathcal{X}$ is the sum of the output produced
by its employees, $\int F(x,y)\,d\pi_{x}(y)$. In particular, the production of
a firm type $x\in\mathcal{X}$ depends linearly on the distribution of its
employees’ types, $\pi_{x}\in\mathscr{P}(\mathcal{Y})$. Choné and Kramarz [12]
relax this restriction and allow firms to aggregate the skills of their
employees in more general way. The production of a firm type $x\in\mathcal{X}$
depends non-linearly on the distribution $\pi_{x}\in\mathscr{P}(\mathcal{Y})$
of its employees’ types.111WOT is connected to classic problems in other
fields, such as entropic transport and martingale transport, see Appendix A.

#### Barycentric WOT problem

We will say that the WOT problem (4) is _barycentric_ when $\mathcal{F}(x,p)$
only depends on the barycenter of $p$, that is when
$\mathcal{F}(x,p)=F\left(x,\int y\,dp(y)\right)$ for some function
$F:\mathcal{X}\times\operatorname{conv}(\mathcal{Y})\to\mathbb{R}$, where
$\operatorname{conv}(\mathcal{Y})$ is the convex hull of $\mathcal{Y}$.222The
particular case where $F(x,y)=\|x-y\|^{2}$ is called the quadratic barycentric
WOT problem. In the economic context, the barycentric specification is valid
if the production of a firm depends on the distribution of its employees’
types, $p\in\mathscr{P}(\mathcal{Y})$, only through their aggregate skills,
$\int y\,dp(y)$.

###  2.3 Duality

Just like the Kantorovich problem (1) admits the dual formulation (2), the WOT
problem (4) also admits a dual formulation under some assumptions on
$\mathcal{F}$. For the WOT problem to be convex, and hence hope for strong
duality to hold, we require that $p\mapsto\mathcal{F}(x,p)$ is convex for all
$x\in\mathcal{X}$. We refer to [23, Section 9] for the technical assumptions
and details.

Under these assumptions, the WOT problem (4) admits the following dual
formulation [23, Theorem 9.5]:

$\displaystyle\operatorname{WOT}(\mu,\nu)=\inf_{\varphi\in\mathcal{C}(\mathcal{Y})}\int_{\mathcal{X}}R_{\mathcal{F}}(\varphi)\,d\mu+\int_{\mathcal{Y}}\varphi\,d\nu$
(5)

where
$R_{\mathcal{F}}(\varphi)(x)=\sup_{p\in\mathscr{P}(\mathcal{Y})}\mathcal{F}(x,p)-\int\varphi\,dp$.

This dual formulation can in turn be interpreted in our economic framework:
$\varphi(y)$ represents the wage of the worker with type $y\in\mathcal{Y}$.
Given a wage function $\varphi$, a firm of type $x\in\mathcal{X}$ hires
workers according to a probability distribution $p\in\mathscr{P}(\mathcal{Y})$
chosen to maximize profit defined as output $\mathcal{F}(x,p)$ minus wage bill
$\int\varphi(y)\,dp(y)$. $R_{\mathcal{F}}(\varphi)(x)$ is therefore the
maximum profit the firm type $x\in\mathcal{X}$ can attain given the wage
function $\varphi$, so that $\int_{\mathcal{X}}R_{\mathcal{F}}(\varphi)\,d\mu$
is the total profit in the economy. The salaries are then chosen so as to
minimize the sum of the profits and of the salaries.

When the cost function $\mathcal{F}$ is barycentric, i.e. when
$\mathcal{F}(x,p)=F\left(x,\int y\,dp(y)\right)$ for some
$F:\mathcal{X}\times\operatorname{conv}(\mathcal{Y})\to\mathbb{R}$ [23, Proof
of Theorem 2.11] prove another dual formulation:

$\displaystyle\operatorname{WOT}(\mu,\nu)=\inf_{\begin{subarray}{c}\psi\in\mathcal{C}(\operatorname{conv}(\mathcal{Y}))\\\
\text{convex,
Lipschitz}\end{subarray}}\int_{\mathcal{X}}Q_{F}(\psi)\,d\mu+\int_{\mathcal{Y}}\psi\,d\nu$
(6)

where
$Q_{F}(\psi)(x)=\sup_{y\in\operatorname{conv}(\mathcal{Y})}F(x,y)-\psi(y)$.
[23, Proof of Theorem 2.11] gives a way to construct a minimizer
${\psi_{\star}}$ of the dual (6) from a minimizer ${\varphi_{\star}}$ of the
more general dual problem (5), by simply taking for ${\psi_{\star}}$ the
largest convex function that is smaller than ${\varphi_{\star}}$, i.e.:

$\displaystyle{\psi_{\star}}:z\mapsto\inf_{\begin{subarray}{c}p\in\mathscr{P}(\mathcal{Y})\\\
\int y\,dp(y)=z\end{subarray}}\int_{\mathcal{Y}}{\varphi_{\star}}\,dp.$ (7)

The convexity of dual minimizers in the barycentric case is easy to interpret
in our economic setting. Here the output produced by a firm depends only on
the aggregate skill of its employees. If the wage function is
${\varphi_{\star}}$, ${\psi_{\star}}(z)$ given by (7) represents the lowest
wage bill that a firm must spend to achieve the aggregate skill $z=\int
y\,dp(y)$. The convexity of the wage thus directly results from the firms’
ability to aggregate the skills of their employees.

##  3 Weak Optimal Transport with Unnormalized Kernel (WOTUK)

###  3.1 From WOT to WOTUK

Choné et al. [13] relax the assumption that $\pi_{x}$ in (3) and (4) is a
_probability_ measure. They allow $\pi_{x}$ to be a positive measure. Denoting
by $\mathscr{M}(\mathcal{Y})$ the set of positive measures over $\mathcal{Y}$,
they introduce the weak optimal transport problem with unnormalized kernel
(WOTUK) as

$\displaystyle\operatorname{WOTUK}(\mu,\nu)\operatorname{\overset{\operatorname{def}}{=}}\sup_{\begin{subarray}{c}q\in\mathscr{M}(\mathcal{Y})^{\mathcal{X}}\\\
\int
q_{x}\,d\mu(x)=\nu\end{subarray}}\int_{\mathcal{X}}\mathcal{F}(x,q_{x})\,d\mu(x)$
(8)

where $\mathcal{F}:\mathcal{X}\times\mathscr{M}(\mathcal{Y})\to\mathbb{R}$.
The constraint $\int q_{x}\,d\mu(x)=\nu$ expresses that the unnormalized
kernel $q$ transports $\mu$ onto $\nu$. Choné et al. [13] connect the WOTUK
problem (8) to a WOT problem as follows. Letting

$\Pi(\ll\mu,\nu)\operatorname{\overset{\operatorname{def}}{=}}\\{\pi\in\Pi(\eta,\nu)\,,\,\eta\in\mathscr{P}(\mathcal{X}),\eta\ll\mu\\},$

denote the set of probability measure over $\mathcal{X}$ that are absolutely
continuous with respect to $\mu$, they show that

$\displaystyle\operatorname{WOTUK}(\mu,\nu)$
$\displaystyle=\sup_{\begin{subarray}{c}\eta\in\mathscr{P}(\mathcal{X})\\\
\eta\ll\mu\end{subarray}}\sup_{\pi\in\Pi(\eta,\nu)}\int\mathcal{F}\left(x,\frac{d\eta}{d\mu}(x)\pi_{x}\right)\,d\mu(x)$
(9)
$\displaystyle=\sup_{\pi\in\Pi(\ll\mu,\nu)}\int\mathcal{F}\left(x,\frac{d\pi_{1}}{d\mu}(x)\pi_{x}\right)\,d\mu(x)$
(10)

where $\pi_{x}\in\mathscr{P}(\mathcal{Y})$ is the unique disintegration of
$\pi$ with respect to $\eta$, i.e. such that $d\pi(x,y)=d\eta(x)d\pi_{x}(y)$,
and $\pi_{1}$ is the first marginal of $\pi$. At given $\eta$, we thus get
back the WOT problem studied in Section 2. Instead of constraining the first
marginal of $\pi$ to be $\mu$, the WOTUK problem only imposes that the first
marginal is absolutely continuous with respect to $\mu$. Choné et al. [13]
show that the density of $\eta$ with respect to $\mu$ is nothing else than the
mass of $q_{x}$, i.e.,
$\frac{d\eta}{d\mu}(x)=\frac{d\pi_{1}}{d\mu}(x)=q_{x}(\mathcal{Y})$.

In the economic setting of [12], $q_{x}(\mathcal{Y})$ represents the number of
employees (i.e., the size) of firms with type $x$. Allowing $q_{x}$ to be an
unnormalized positive measure instead of a probability measure avoids having
to assume that all firms have the same size. In contrast to earlier
literature, firms’ sizes are unknowns to be determined rather than given
parameters.

#### Conical WOTUK problem

The conical WOTUK problem corresponds to the case where

$\mathcal{F}(x,p)=F\left(x,\int_{\mathcal{Y}}y\,dp(y)\right)$

for some $F:\mathcal{X}\times\operatorname{cone}(\mathcal{Y})\to\mathbb{R}$,
where the conical hull of $\mathcal{Y}$ is given by

$\operatorname{cone}(\mathcal{Y})\operatorname{\overset{\operatorname{def}}{=}}\Bigg{\\{}\sum_{i=1}^{n}\lambda_{i}y_{i}\,,\,\lambda_{1},\ldots,\lambda_{n}\in\mathbb{R}_{+},y_{1},\ldots,y_{n}\in\mathcal{Y},n\geq
1\Bigg{\\}}.$

In [12], a firm’s output depends on the _conical_ combination of its
employees’ types, $\int y\,dq_{x}(y)$. The combination is said to be “conical”
because the mass of $q_{x}$ is not necessarily equal to one. In other words,
the aggregate skill of the workers hired by a firm is not their average skills
as in the WOT setting, but their average skills _scaled by the positive
factor_ $q_{x}(\mathcal{Y})$ that represents the number of employees.

###  3.2 Duality

The WOTUK problem (8) admits dual formulations that are similar to the dual
WOT formulations (5) and (6). The main difference with the results of
subsection 2.3 lies in the fact that $\mathscr{P}(\mathcal{Y})$ is replaced by
$\mathscr{M}(\mathcal{Y})$ and that $\operatorname{conv}(\mathcal{Y})$ is
replaced by $\operatorname{cone}(\mathcal{Y})$.

Under some technical assumptions on $\mathcal{F}$ detailed in [13], the
theorem 3.2 in the same reference proves that the WOTUK problem (8) admits the
following dual formulation:

$\displaystyle\operatorname{WOTUK}(\mu,\nu)=\inf_{\varphi\in\mathcal{C}_{b}(\mathcal{Y})}\int_{\mathcal{X}}K_{\mathcal{F}}(\varphi)\,d\mu+\int_{\mathcal{Y}}\varphi\,d\nu$
(11)

where
$K_{\mathcal{F}}(\varphi)(x)=\sup_{m\in\mathscr{M}(\mathcal{Y})}\mathcal{F}(x,m)-\int\varphi\,dm$.

Similarly, [13, Theorem 5.1] proves that the conical WOTUK problem admits the
dual formulation:

$\displaystyle\operatorname{WOTUK}(\mu,\nu)=\inf_{\begin{subarray}{c}\psi\in\mathcal{C}(\operatorname{cone}(\mathcal{Y}))\text{
convex, }\\\ \text{positively
homogeneous}\end{subarray}}\int_{\mathcal{X}}J_{F}(\psi)\,d\mu+\int_{\mathcal{Y}}\psi\,d\nu$
(12)

where
$J_{F}(\psi)(x)=\sup_{y\in\operatorname{cone}(\mathcal{Y})}F(x,y)-\psi(y)$.
They show that a minimizer ${\psi_{\star}}$ of the dual problem (12) from a
minimizer ${\varphi_{\star}}$ of the more general dual problem (11) by taking
for ${\psi_{\star}}$ the largest convex and positively homogeneous function
that is smaller than ${\varphi_{\star}}$, i.e.:

$\displaystyle{\psi_{\star}}:z\mapsto\inf_{\begin{subarray}{c}m\in\mathscr{M}(\mathcal{Y})\\\
\int y\,dm(y)=z\end{subarray}}\int_{\mathcal{Y}}{\varphi_{\star}}\,dm.$ (13)

In the economic setting of [12], a dual optimizer $\varphi$ is a wage
function: $\varphi(y)$ represents the wage paid to a worker of type
$y\in\mathcal{Y}$. As $R_{\mathcal{F}}(\varphi)(x)$ and $Q_{F}(\psi)(x)$ in
subsection 2.3, $K_{\mathcal{F}}(\varphi)(x)$ and $J_{F}(\psi)(x)$ are two
forms for the profit function, i.e., for the maximal profit that firms of each
type $x\in\mathcal{X}$ achieve under the wage functions $\varphi$ or $\psi$.

##  4 Algorithms

In this section, we only consider the case of discrete measures
$\mu\in\mathscr{P}(\mathcal{X})$ and $\nu\in\mathscr{P}(\mathcal{Y})$. We will
write $\mu=\sum_{i=1}^{n}a_{i}\delta_{x_{i}}$ where $n\geq 1$ is the number of
firm types, $x_{1},\ldots,x_{n}\in\mathcal{X}$ are the firm types and
$a\in\mathbb{R}^{n}$ represents the proportion of the firm types in the
economy ($a>0$ and $\sum_{i=1}^{n}a_{i}=1$). Likewise, we will write
$\nu=\sum_{j=1}^{m}b_{j}\delta_{y_{j}}$ where $m\geq 1$ is the number of
worker types, $y_{1},\ldots,y_{m}\in\mathcal{Y}$ are the worker types and $b$
represents the proportions of the worker types in the population ($b>0$ and
$\sum_{j=1}^{m}b_{j}=1$). Since we consider here maximization problems, we
will require that the cost function $\mathcal{F}$ is concave and
differentiable with respect to its second argument.

###  4.1 The primal problems

A matching $\pi\in\Pi(\mu,\nu)$ is now represented by a matrix
$P\in\mathbb{R}^{n\times m}$ such that $P_{ij}$ represents the proportion of
the firm type $x_{i}$ which is matched with the worker type $y_{j}$. For the
WOT problem (4), the marginal constraints on $P$ translate into:

$\Pi(\mu,\nu)=\\{P\in\mathbb{R}^{n\times
m}_{+}\,,\,P\mathbf{1}=a,P^{\top}\mathbf{1}=b\\}.$

For the WOTUK problem (9), the set of constraints is simply

$\Pi(\ll\mu,\nu)=\\{P\in\mathbb{R}^{n\times
m}_{+}\,,\,P^{\top}\mathbf{1}=b\\}$

because as explained in Section 3.1 the first marginal $P\mathbf{1}=\eta$ is
unconstrained and $\eta$ is an unknown variable to be determined. The
difference between the WOT problem (4) and the WOTUK problem (9) lies in the
constraint set only. The WOTUK problem corresponds to the WOT problem where
the first marginal constraint has been removed, which establishes an
interesting link with the unbalanced optimal transport theory [10].

Let us now write the objective for the WOT and WOTUK problems in the discrete
setting. The disintegration $\pi_{x_{i}}$ representing the workers hired by
the firm of type $x_{i}$ writes
$\frac{1}{a_{i}}\sum_{j=1}^{m}P_{ij}\delta_{y_{j}}$. For simplicity, we make
the following change of notations: we define
$\widetilde{\mathcal{F}}:\mathcal{X}\times\mathbb{R}^{m}$ by

$\widetilde{\mathcal{F}}:(x,p)\mapsto\mathcal{F}\left(x,\sum_{j=1}^{m}p_{j}\delta_{y_{j}}\right).$

Note that $\widetilde{\mathcal{F}}$ depends on $y_{1},\ldots,y_{m}$ and that
$\frac{\partial\widetilde{\mathcal{F}}}{\partial
p_{j}}(x,p)=\left\langle\delta_{y_{j}},\nabla_{2}\mathcal{F}\left(x,\sum_{k=1}^{m}p_{k}\delta_{y_{k}}\right)\right\rangle$
where
$\langle\cdot,\cdot\rangle:\mathscr{M}(\mathcal{X})\times\mathscr{M}(\mathcal{X})^{*}$
is the duality bracket. We will use the notation $P_{i:}=(P_{ij})_{1\leq j\leq
m}$ for $1\leq i\leq n$.

With these notations, the objective of the WOT problem (4) and of the WOTUK
problem (9) writes:

$\displaystyle\\!\\!\\!f(P)\operatorname{\overset{\operatorname{def}}{=}}\sum_{i=1}^{n}a_{i}\mathcal{F}\left(x_{i},\frac{1}{a_{i}}\sum_{j=1}^{m}P_{ij}\delta_{y_{j}}\right)=\sum_{i=1}^{n}a_{i}\widetilde{\mathcal{F}}\left(x_{i},\frac{P_{i:}}{a_{i}}\right).$

Since the constraints sets $\Pi(\mu,\nu)$ and $\Pi(\ll\mu,\nu)$ are convex and
$\mathcal{F}$ and $\widetilde{\mathcal{F}}$ are convex in their second
argument, both the discrete WOT and WOTUK problems are convex (although not
strictly) optimization problems. In order to solve them, we propose to apply a
mirror ascent on $P$ (using the Kullback-Leibler divergence). The gradient of
the total output writes:

$\frac{\partial f}{\partial
P_{ij}}(P)=\left[\nabla_{2}\widetilde{\mathcal{F}}\left(x_{i},\frac{P_{i:}}{a_{i}}\right)\right]_{j}.$

After each gradient step, the resulting matching $P$ should be projected (for
the KL divergence) onto $\Pi(\mu,\nu)$ (for the WOT problem) or
$\Pi(\ll\mu,\nu)$ (for the WOTUK problem). For the WOT problem, we have to
solve $\min_{Q\in\Pi(\mu,\nu)}\operatorname{KL}(Q|P)$. This problem is
equivalent to the entropic OT problem (17) with cost function $-\log P$ and
regularization strength $\varepsilon=1$, and can therefore be efficiently
solved using the Sinkhorn algorithm [14]. For the WOTUK problem, we have to
solve $\min_{Q\in\Pi(\ll\mu,\nu)}\operatorname{KL}(Q|P)$ which admits the
following closed-form solution (see a proof in Appendix B.1):
${Q_{\star}}=P\odot b/P^{\top}\mathbf{1}$ where $\odot$ and $/$ are the
elementwise multiplication and division.

Algorithm 1 Mirror Ascent Algorithm for WOT and WOTUK (primal)

Input Stepsize $\gamma>0$, tolerance $\varepsilon$

Initialize $P=ab^{\top}$

while $\overline{\mathcal{G}}(P)>\varepsilon f(P)$ do

$P\leftarrow P\exp\left(\gamma\,\nabla f(P)\right)$

For WOT:

$P\leftarrow\text{Sinkhorn}(a,b,\text{kernel}=P)$

For WOTUK:

$P\leftarrow P\odot b/P^{\top}\mathbf{1}$

end while

Return $P$

#### Numerical guarantee

Along the mirror ascent iterations over $P$, we can monitor the convergence by
looking at the gap

$\mathcal{G}(P)\operatorname{\overset{\operatorname{def}}{=}}\sup_{Q\in
K}f(Q)-f(P)$

where $K=\Pi(\mu,\nu)$ for the WOT problem and $K=\Pi(\ll\mu,\nu)$ for the
WOTUK problem. By definition, $\mathcal{G}(P)\geq 0$ and by the concavity of
$f$,

$\mathcal{G}(P)\leq\sup_{Q\in K}\langle\nabla
f(P),Q-P\rangle\operatorname{\overset{\operatorname{def}}{=}}\overline{\mathcal{G}}(P).$

When $K=\Pi(\mu,\nu)$, the upper bound $\overline{\mathcal{G}}(P)$ on
$\mathcal{G}(P)$ corresponds to an optimal transport problem (with cost matrix
$-\nabla f(P)$) and can either be computed exactly or be itself upper bounded
using the solution ${Q_{\star}}$ of an entropic OT problem (efficiently solved
using the Sinkhorn algorithm). When $K=\Pi(\ll\mu,\nu)$, the upper bound
$\overline{\mathcal{G}}(P)$ on $\mathcal{G}(P)$ admits the following closed
form solution: ${Q_{\star}}_{ij}=b_{j}$ if $i=\operatorname*{arg\,max}_{1\leq
k\leq n}[\nabla f(P)]_{kj}$ and ${Q_{\star}}_{ij}=0$ otherwise.

The algorithm stops when $\overline{\mathcal{G}}(P)\leq\varepsilon f(P)$ for a
tolerance $\varepsilon>0$. We summarize the mirror ascent method for the WOT
and WOTUK problems in Algorithm 1.

###  4.2 The dual problems

In the discrete setting, the general dual for WOT (5) writes:

$\displaystyle\min_{\varphi\in\mathbb{R}^{m}_{+}}\langle
b,\varphi\rangle+\max_{\begin{subarray}{c}P\in\mathbb{R}^{n\times m}_{+}\\\
P\mathbf{1}=a\end{subarray}}\sum_{i=1}^{n}a_{i}\widetilde{\mathcal{F}}\left(x_{i},\frac{P_{i:}}{a_{i}}\right)-\mathbf{1}^{\top}P\varphi$
(14)

and likewise for the general dual for the WOTUK problem (11):

$\displaystyle\min_{\varphi\in\mathbb{R}^{m}_{+}}\langle
b,\varphi\rangle+\max_{P\in\mathbb{R}^{n\times
m}_{+}}\sum_{i=1}^{n}a_{i}\widetilde{\mathcal{F}}\left(x_{i},\frac{P_{i:}}{a_{i}}\right)-\mathbf{1}^{\top}P\varphi.$
(15)

Let us define

$h:(\varphi,P)\mapsto\sum_{i=1}^{n}a_{i}\widetilde{\mathcal{F}}\left(x_{i},\frac{P_{i:}}{a_{i}}\right)-\mathbf{1}^{\top}P\varphi.$

Algorithm 2 Mirror Descent Algorithm for WOT and WOTUK (dual)

Input $\gamma_{1}>0$, $\gamma_{2}>0$, $K_{1},K_{2}\in\mathbb{N}$

Initialize $\varphi\in\mathbb{R}^{m}_{+}$ and $P=ab^{\top}$

for $k_{1}=0$ to $K_{1}$ do

for $k_{2}=0$ to $K_{2}$ do

$P\leftarrow P\exp\left(\gamma_{1}\,\nabla_{P}h(\varphi,P)\right)$

For WOT:

$P\leftarrow\operatorname{diag}\left(a/P\mathbf{1}\right)P$

end for

$\varphi\leftarrow\varphi\exp\left(-\gamma_{2}\left[b-P^{\top}\mathbf{1}\right]\right)$

end for

Use a linear programming solver to compute $\psi(z)$ for $z\in\mathcal{Y}$:

For WOT:
$\psi(z)=\max_{\begin{subarray}{c}\lambda\in\mathbb{R}^{q},\mu\in\mathbb{R}\\\
\forall
j,\langle\lambda,y_{j}\rangle+\mu\leq\varphi_{j}\end{subarray}}\langle\lambda,z\rangle+\mu.$

For WOTUK: $\psi(z)=\max_{\begin{subarray}{c}\lambda\in\mathbb{R}^{q}\\\
\forall
j,\langle\lambda,y_{j}\rangle\leq\varphi_{j}\end{subarray}}\langle\lambda,z\rangle.$

Return Dual variables $\varphi_{j},\psi(y_{j})$ for $1\leq j\leq m$.

To solve problems (14) and (15), we propose to run a mirror descent on
$\varphi$ (with the Kullback-Leibler divergence). The objective is itself a
maximization problem (over $P$). The envelope theorem yields the gradient of
the objective, provided the optimal $P$ is given. At each gradient step on
$\varphi$, we therefore propose to run a mirror ascent on $P$ at fixed
$\varphi$. For the WOTUK problem (15), no projection are needed, while for the
WOT problem (14), we need to project $P$ onto $\\{P\in\mathbb{R}^{n\times
m}_{+},P\mathbf{1}=a\\}$ during the ascents. The proof in Appendix B.1
directly adapts to this case, and the projection amounts to reweighting the
rows of $P$.

We can construct dual minimizers for the barycentric WOT (6) and conical WOTUK
(12) from a solution ${\varphi_{\star}}$ of (14) and (15) respectively, using
the results given in (7) and (13) respectively. In the discrete setting, these
linear programs respectively write:

${\psi_{\star}}:z\mapsto\min_{\begin{subarray}{c}p\in\mathbb{R}_{+}^{m}\\\
\sum_{j=1}^{m}p_{j}=1\\\ \sum_{j=1}^{m}p_{j}y_{j}=z\end{subarray}}\langle
p,{\varphi_{\star}}\rangle$

and for WOTUK

${\psi_{\star}}:z\mapsto\min_{\begin{subarray}{c}p\in\mathbb{R}_{+}^{m}\\\
\sum_{j=1}^{m}p_{j}y_{j}=z\end{subarray}}\langle p,{\varphi_{\star}}\rangle.$

Since we may be interested in differentiating those functions
${\psi_{\star}}$, we rather compute the dual problems of the above linear
programs (see a proof in the Appendix B.2):

${\psi_{\star}}:z\mapsto\max_{\begin{subarray}{c}\lambda\in\mathbb{R}^{q},\mu\in\mathbb{R}\\\
\forall
j,\langle\lambda,y_{j}\rangle+\mu\leq{\varphi_{\star}}_{j}\end{subarray}}\langle\lambda,z\rangle+\mu\quad\text{and}\quad{\psi_{\star}}:z\mapsto\max_{\begin{subarray}{c}\lambda\in\mathbb{R}^{q}\\\
\forall
j,\langle\lambda,y_{j}\rangle\leq{\varphi_{\star}}_{j}\end{subarray}}\langle\lambda,z\rangle.$

We summarize the mirror ascent method for the WOT and WOTUK problems in
Algorithm 2.

##  5 Experiments and Applications in Economics

All the experiments have been run on a Google colab notebook using JAX. We
used the following packages: OTT [16] for entropic optimal transport, POT [19]
for exact optimal transport, GLPK [31] through CVXPY [2] for linear
programming.

Figure 1: An example of discrete probability distributions of firm types $\mu$
(left) and worker types $\nu$ (right). Each dot represents a Dirac mass, and
its size represents its weight. Here, all firm types have the same weight
while there are more specialist workers in the economy than generalists.

From now on, we focus on the setting of [12] and will therefore consider the
distributions of firm types and worker types and the cost function they use in
their model. We focus on the case where there are two skills in the economy.
The set of firm types is
$\mathcal{X}=\\{(z,\alpha_{1},\alpha_{2})\in\mathbb{R}_{+}^{3},\alpha_{1}+\alpha_{2}=1\\}$
where $z$ represents the productivity of the firm type and $\alpha_{i}$,
$i\in\\{1,2\\}$ its demand in skill $i$. The set of worker types is
$\mathcal{Y}=\\{(x_{1},x_{2})\in\mathbb{R}_{+}^{2}\\}$, where $x_{i}$,
$i\in\\{1,2\\}$, represents the proficiency of the worker type in skill $i$.
We will denote the skill profile of a worker type by
$\theta=\arctan\left(x_{2}/x_{1}\right)\in[0,\frac{\pi}{2}]$, where $\theta=0$
represents an expert worker in skill $1$, $\theta=\frac{\pi}{2}$ an expert
worker in skill $2$ and $\theta=\frac{\pi}{4}$ a generalist worker. The skill
profile represents the worker’s comparative advantage in skill 2 over skill 1.
We plot in Figure 1 an example of such distributions.

As for the production function, we consider a Constant Elasticity of
Substitution (CES) function:

$F\left((z,\alpha_{1},\alpha_{2}),(x_{1},x_{2})\right)=\frac{z}{\zeta}\left(\alpha_{1}x_{1}^{\sigma}+\alpha_{2}x_{2}^{\sigma}\right)^{\zeta/\sigma}$

with parameters $\zeta=\sigma=\frac{1}{2}$.

###  5.1 Comparison between OT, EOT, WOT and WOTUK

(a) OT

(b) EOT ($\varepsilon=0.1$)

(c) WOT

(d) WOTUK

Figure 2: Optimal transport plans for different variants of OT, between the
marginals depicted in Figure 1. Each row represents a firm type, each column
represents a worker type, and each cell represents the matching strength
between the corresponding firm and worker types. The darker the cell, the
stronger the matching. The color scale is the same in the four pictures.
Figure 3: The firm types’ size
$N(z,\alpha_{1},\alpha_{2})=\frac{d\eta}{d\mu}(z,\alpha_{1},\alpha_{2})$ in
function of $\alpha_{2}$ in the WOTUK model. See equation (9) for the
definition of $\eta$. Specialist firms get bigger while generalist firms get
smaller.

In Figure 2, we showcase the empirical differences between four different
optimal transport problems, when computed on the economic setting described
above: the OT problem (1), the EOT problem (17), the WOT problem (4) and the
WOTUK problem (8).

While all four methods agree that firms with a higher demand in skill $2$
(i.e. with a higher value of $\alpha_{2}$) employ workers with better
proficiency in skill $2$ (i.e. with higher $\theta$), we can spot some
differences.

The OT and WOTUK plans (Figures 2(a) and 2(d)) are sparse, i.e. firms do not
tend to employ workers of different types, while the EOT and WOT plans (Figure
2(b) and 2(c)) show that in these models, generalist firms (i.e. firms such
that $\alpha_{1}\approx\alpha_{2}$) mix different types of workers. Since the
EOT plan is a blurry version of the OT plan, firms in this model mix workers
of different but similar types. On the other hand, generalist firms in the WOT
model employ both specialist workers of both skills, and generalist workers.

Since there are ten firm types with uniform proportion in the economy, the
maximal value in the transport matrix is $1/10$ in the OT, EOT and WOT models.
Since the WOTUK model is unnormalized, its optimal plan may (and does) display
higher values. We depict in Figure 3 the size of firm types in function of
their $\alpha_{2}$. Specialist firms get bigger and generalist firms get
smaller, which was to be expected since there are more specialist workers and
less generalist workers on the job market.

###  5.2 Simulation of the economic model

(a) Scenario A

(b) Scenario B

Figure 4: The “mapping” $\alpha_{2}\mapsto\theta$.

We now consider the same 2D model, but with 200 firms and workers. Like
before, the firms are distributed uniformly along $z$, while we consider the
two cases: many specialist workers (scenario A) and many generalist workers
(scenario B).

Figure 4 shows the aggregate skill profile of employees (the ratio of skill 2
over skill 1) as a function of the technical intensity in skill 2 of their
employing firm. As predicted by [12], firms whose technology is very intensive
in skill 2 (i.e., high $\alpha_{2}$) use more skill 2 relative to skill 1 in
scenario A compared to scenario B. In other words, firms are able to
specialize to their “core business” in scenario A. This is because in that
scenario (with many specialist workers), the salary tends to become linear
(see Figure 5) and firms freely adjust the proportion of specialists they hire
to achieve their optimal mix of skills. In contrast, in scenario B, the
workers’ salary is strictly convex, implying that specialist workers are
expensive (compared to generalists) and hence it is too costly for firms to
hire the specialists they would need to take full advantage of their
technology.

(a) When there are many specialists workers in the economy, the salary becomes
linear.

(b) When there are many generalist workers in the economy, the salary is
strictly convex.

Figure 5: The isowage curves $\psi(x_{1},x_{2})=\text{constant}$ in both
scenarios.

##  6 Conclusion and Future Work

We have proposed the first algorithms to compute the weak optimal transport
and weak optimal transport with unnormalized kernel problems in the discrete
setting, both in their primal and dual formulations, as well as in the special
case of barycentric cost functions. We have illustrated the interest of the
WOT and WOTUK problems on the economic application of [12]. Future work
includes the statistical study of WOT ad WOTUK, as well as applications of
these models in other applied sciences.

#### Acknowledgments

The three authors acknowledge support from the European Research Council
(ERC), advanced grant FIRMNET (project ID 741467).

## References

  * Acciaio et al. [2020] Beatrice Acciaio, Mathias Beiglboeck, and Gudmund Pammer.  Weak transport for non-convex costs and model-independence in a fixed-income market.  _arXiv preprint arXiv:2011.04274_ , 2020.
  * Agrawal et al. [2018] Akshay Agrawal, Robin Verschueren, Steven Diamond, and Stephen Boyd.  A rewriting system for convex optimization problems.  _Journal of Control and Decision_ , 5(1):42–60, 2018.
  * Alaux et al. [2019] Jean Alaux, Edouard Grave, Marco Cuturi, and Armand Joulin.  Unsupervised hyper-alignment for multilingual word embeddings.  In _International Conference on Learning Representations_ , 2019.
  * Alibert et al. [2018] Jean-Jacques Alibert, Guy Bouchitté, and Thierry Champion.  A new class of cost for optimal transport planning.  working paper or preprint, March 2018.  URL https://hal-univ-tln.archives-ouvertes.fr/hal-01741688.
  * Arjovsky et al. [2017] Martin Arjovsky, Soumith Chintala, and Léon Bottou.  Wasserstein generative adversarial networks.  _Proceedings of the 34th International Conference on Machine Learning_ , 70:214–223, 2017.
  * Backhoff-Veraguas and Pammer [2020] Julio Daniel Backhoff-Veraguas and Gudmund Pammer.  Applications of weak transport theory.  _arXiv preprint arXiv:2003.05338_ , 2020.
  * Beiglböck et al. [2013] Mathias Beiglböck, Pierre Henry-Labordere, and Friedrich Penkner.  Model-independent bounds for option prices—a mass transport approach.  _Finance and Stochastics_ , 17(3):477–501, 2013\.
  * Bonneel et al. [2016] Nicolas Bonneel, Gabriel Peyré, and Marco Cuturi.  Wasserstein barycentric coordinates: histogram regression using optimal transport.  _ACM Transactions on Graphics_ , 35(4):71:1–71:10, 2016.
  * Cazelles et al. [2021] Elsa Cazelles, Felipe Tobar, and Joaquin Fontbona.  Streaming computation of optimal weak transport barycenters.  _arXiv preprint arXiv:2102.13380_ , 2021.
  * Chizat et al. [2018] Lenaic Chizat, Gabriel Peyré, Bernhard Schmitzer, and François-Xavier Vialard.  Unbalanced optimal transport: geometry and Kantorovich formulation.  _Journal of Functional Analysis_ , 274(11):3090–3123, 2018.
  * Chizat [2017] Lénaïc Chizat.  _Unbalanced optimal transport: models, numerical methods, applications_.  PhD thesis, 2017.
  * Choné and Kramarz [2021] Philippe Choné and Francis Kramarz.  Matching workers’ skills and firms’ technologies: From bundling to unbundling.  Technical report, 2021.
  * Choné et al. [2022] Philippe Choné, Nathael Gozlan, and Francis Kramarz.  Weak optimal transport with unnormalized kernels.  _arXiv preprint arXiv: preprint: 2203.16227_ , 2022.
  * Cuturi [2013] Marco Cuturi.  Sinkhorn distances: lightspeed computation of optimal transport.  In _Advances in Neural Information Processing Systems 26_ , pages 2292–2300, 2013.
  * Cuturi and Peyré [2016] Marco Cuturi and Gabriel Peyré.  A smoothed dual approach for variational Wasserstein problems.  _SIAM Journal on Imaging Sciences_ , 9(1):320–343, 2016.
  * Cuturi et al. [2022] Marco Cuturi, Laetitia Meng-Papaxanthos, Yingtao Tian, Charlotte Bunne, Geoff Davis, and Olivier Teboul.  Optimal transport tools (ott): A jax toolbox for all things wasserstein, 2022.
  * Eeckhout and Kircher [2018] Jan Eeckhout and Philipp Kircher.  Assortative matching with large firms.  _Econometrica_ , 86(1):85–132, 2018.
  * Figalli [2010] Alessio Figalli.  The optimal partial transport problem.  _Archive for Rational Mechanics and Analysis_ , 195(2):533–560, 2010.
  * Flamary et al. [2021] Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy, Ievgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer.  Pot: Python optimal transport.  _Journal of Machine Learning Research_ , 22(78):1–8, 2021.  URL http://jmlr.org/papers/v22/20-451.html.
  * Galichon [2016] Alfred Galichon.  _Optimal Transport Methods in Economics_.  Princeton University Press, 2016.
  * Galichon and Salanié [2015] Alfred Galichon and Bernard Salanié.  Cupid’s invisible hand: Social surplus and identification in matching models.  _Available at SSRN 1804623_ , 2015.
  * Galichon et al. [2014] Alfred Galichon, Pierre Henry-Labordère, and Nizar Touzi.  A stochastic control approach to no-arbitrage bounds given marginals, with an application to lookback options.  _Annals of Applied Probability_ , 24(1):312–336, 2014.
  * Gozlan et al. [2017] Nathael Gozlan, Cyril Roberto, Paul-Marie Samson, and Prasad Tetali.  Kantorovich duality for general transport costs and applications.  _Journal of Functional Analysis_ , 273(11):3327–3405, 2017.  ISSN 0022-1236.  doi: https://doi.org/10.1016/j.jfa.2017.08.015.  URL https://www.sciencedirect.com/science/article/pii/S0022123617303294.
  * Grave et al. [2019] Edouard Grave, Armand Joulin, and Quentin Berthet.  Unsupervised alignment of embeddings with wasserstein procrustes.  2019\.
  * Hashimoto et al. [2016] Tatsunori Hashimoto, David Gifford, and Tommi Jaakkola.  Learning population-level diffusions with generative RNNs.  In _International Conference on Machine Learning_ , pages 2417–2426, 2016.
  * Heckman and Scheinkman [1987] James Heckman and Jose Scheinkman.  The importance of bundling in a gorman-lancaster model of earnings.  _The Review of Economic Studies_ , 54(2):243–255, 1987.
  * Kantorovich [1942] Leonid Kantorovich.  On the transfer of masses (in russian).  _Doklady Akademii Nauk_ , 37(2):227–229, 1942\.
  * Kelso and Crawford [1982] Alexander S Kelso and Vincent P Crawford.  Job matching, coalition formation, and gross substitutes.  _Econometrica: Journal of the Econometric Society_ , pages 1483–1504, 1982.
  * Korotin et al. [2022] Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev.  Neural optimal transport, 2022.
  * Lindenlaub [2017] Ilse Lindenlaub.  Sorting multidimensional types: Theory and application.  _The Review of Economic Studies_ , 84(2):718–789, 2017.
  * Makhorin [2008] Andrew Makhorin.  Glpk (gnu linear programming kit).  _http://www. gnu. org/s/glpk/glpk. html_ , 2008.
  * Paty et al. [2020] François-Pierre Paty, Alexandre d’Aspremont, and Marco Cuturi.  Regularity as regularization: Smooth and strongly convex brenier potentials in optimal transport.  In Silvia Chiappa and Roberto Calandra, editors, _Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics_ , volume 108 of _Proceedings of Machine Learning Research_ , pages 1222–1232. PMLR, 26–28 Aug 2020.  URL https://proceedings.mlr.press/v108/paty20a.html.
  * Paty and Cuturi [2020] François-Pierre Paty and Marco Cuturi.  Regularized optimal transport is ground cost adversarial.  In Hal Daumé III and Aarti Singh, editors, _Proceedings of the 37th International Conference on Machine Learning_ , volume 119 of _Proceedings of Machine Learning Research_ , pages 7532–7542. PMLR, 13–18 Jul 2020.  URL https://proceedings.mlr.press/v119/paty20a.html.
  * Rabin and Papadakis [2015] Julien Rabin and Nicolas Papadakis.  Convex color image segmentation with optimal transport distances.  In _International Conference on Scale Space and Variational Methods in Computer Vision_ , pages 256–269. Springer, 2015.
  * Salimans et al. [2018] Tim Salimans, Han Zhang, Alec Radford, and Dimitris Metaxas.  Improving GANs using optimal transport.  In _International Conference on Learning Representations_ , 2018.  URL https://openreview.net/forum?id=rkQkBnJAb.
  * Schiebinger et al. [2019] Geoffrey Schiebinger, Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Aryeh Solomon, Joshua Gould, Siyan Liu, Stacie Lin, Peter Berube, et al.  Optimal-transport analysis of single-cell gene expression identifies developmental trajectories in reprogramming.  _Cell_ , 176(4):928–943, 2019.
  * Solomon et al. [2015] Justin Solomon, Fernando De Goes, Gabriel Peyré, Marco Cuturi, Adrian Butscher, Andy Nguyen, Tao Du, and Leonidas Guibas.  Convolutional Wasserstein distances: efficient optimal transportation on geometric domains.  _ACM Transactions on Graphics_ , 34(4):66:1–66:11, 2015.
  * Wilson [1969] Alan Geoffrey Wilson.  The use of entropy maximizing models, in the theory of trip distribution, mode split and route split.  _Journal of Transport Economics and Policy_ , pages 108–126, 1969\.

##  Appendix A Other examples of WOT problems

#### Entropic optimal transport

The entropic optimal transport (EOT) problem [38] is a variant of the
Kantorovich problem in which an entropic regularization term is added:

$\displaystyle\mathscr{S}_{\varepsilon}(\mu,\nu)$
$\displaystyle\operatorname{\overset{\operatorname{def}}{=}}\inf_{\pi\in\Pi(\mu,\nu)}\iint_{\mathcal{X}\times\mathcal{Y}}c\,d\pi+\varepsilon\operatorname{KL}(\pi|\mu\otimes\nu)$
(16) $\displaystyle=\inf_{\pi\in\Pi(\mu,\nu)}\int
c\,d\pi+\varepsilon\int\log\frac{d\pi}{d\mu\otimes\nu}\,d\pi$ (17)

where $c\in\mathcal{C}(\mathcal{X}\times\mathcal{Y})$ is the cost function,
$\varepsilon>0$ is the regularization strength, and
$\operatorname{KL}(\cdot|\cdot)$ is the relative entropy (or Kullback-Leibler
divergence). Considering the disintegration
$\left(\pi_{x}\right)_{x\in\mathcal{X}}$ of $\pi$ with respect to $\mu$, and
noting that $\frac{d\pi}{d\mu\otimes\nu}(x,y)=\frac{d\pi_{x}}{d\nu}(y)$,
problem (17) rewrites:

$\inf_{\pi\in\Pi(\mu,\nu)}\int_{\mathcal{X}}\left[\int_{\mathcal{Y}}\left(c(x,y)+\varepsilon\log\frac{d\pi_{x}}{d\nu}(y)\right)d\pi_{x}(y)\right]d\mu(x)$

which corresponds to the WOT problem (4) with

$\displaystyle\mathcal{F}(x,p)$
$\displaystyle=\int_{\mathcal{Y}}\left(c(x,y)+\varepsilon\log\frac{dp}{d\nu}(y)\right)dp(y)$
$\displaystyle=\int_{\mathcal{Y}}c(x,y)\,dp(y)+\varepsilon\operatorname{KL}(p|\nu).$

#### Martingale optimal transport

The martingale optimal transport (MOT) problem [7] is a variant of the
Kantorovich problem in which the optimal transport plan is constrained to be a
martingale:

$\displaystyle\sup_{\begin{subarray}{c}\pi\in\Pi(\mu,\nu)\\\
\mu\,\text{a.e.},\int
y\,d\pi_{x}(y)=x\end{subarray}}\iint_{\mathcal{X}\times\mathcal{Y}}F(x,y)\,d\pi(x,y).$

Up to the fact that $\mathcal{F}$ is now allowed to take value $+\infty$, this
problem corresponds to the WOT problem (4) with $\mathcal{F}(x,p)=\int
c(x,y)\,dp(y)-\iota\left(\mu\,\text{a.e.},\int y\,dp(y)=x\right)$ where
$\iota(a)$ equals $0$ if the assertion $a$ is true, and $+\infty$ if $a$ is
false.

##  Appendix B Proofs

###  B.1 Proof for the projection onto $\Pi(\ll\mu,\nu)$

Projecting a matrix $P\in\mathbb{R}^{n\times m}_{+}$ onto $\Pi(\ll\mu,\nu)$
means solving the following optimization problem:

$\min_{Q\in\Pi(\ll\mu,\nu)}\operatorname{KL}(Q|P)=\min_{\begin{subarray}{c}Q\in\mathbb{R}^{n\times
m}_{+}\\\
Q^{\top}\mathbf{1}=b\end{subarray}}\operatorname{KL}(Q|P)=\min_{\begin{subarray}{c}Q\in\mathbb{R}^{n\times
m}\\\
Q^{\top}\mathbf{1}=b\end{subarray}}\sum_{i,j}Q_{ij}\left(\log\frac{Q_{ij}}{P_{ij}}-1\right)$

where we can drop the non-negativity constraint over $Q$ for it is already
constrained by the $\log$ in the objective.

The Lagrangian of the problem is

$L(Q,\lambda)=\sum_{i,j}Q_{ij}\left(\log\frac{Q_{ij}}{P_{ij}}-1\right)+\langle\lambda,b-Q^{\top}\mathbf{1}\rangle$

where $\lambda\in\mathbb{R}^{m}$ is the Lagrange multiplier for the constraint
$Q^{\top}\mathbf{1}=b$. The problem is convex, so the first-order condition is
sufficient for optimality. So the solution should verify

$\log\frac{Q_{ij}}{P_{ij}}=\lambda_{j}$

hence

$Q_{ij}=\frac{P_{ij}b_{j}}{\sum_{i^{\prime}}P_{i^{\prime}j}}.$

###  B.2 Proof for the dual formulation of ${\psi_{\star}}$

Let us first compute the dual problem corresponding to

$\min_{\begin{subarray}{c}p\in\mathbb{R}_{+}^{m}\\\ \sum_{j=1}^{m}p_{j}=1\\\
\sum_{j=1}^{m}p_{j}y_{j}=z\end{subarray}}\langle p,\varphi\rangle.$

One has:

$\displaystyle\min_{\begin{subarray}{c}p\in\mathbb{R}_{+}^{m}\\\
\sum_{j=1}^{m}p_{j}=1\\\ \sum_{j=1}^{m}p_{j}y_{j}=z\end{subarray}}\langle
p,\varphi\rangle$ $\displaystyle=\min_{p\in\mathbb{R}_{+}^{m}}\langle
p,\varphi\rangle+\sup_{\lambda\in\mathbb{R}^{q},\mu\in\mathbb{R}}\left\langle\lambda,z-\sum_{j=1}^{m}p_{j}y_{j}\right\rangle+\mu\left(1-\sum_{j=1}^{m}p_{j}\right)$
$\displaystyle=\sup_{\lambda\in\mathbb{R}^{q},\mu\in\mathbb{R}}\langle\lambda,z\rangle+\mu+\inf_{p\in\mathbb{R}_{+}^{m}}\sum_{j=1}^{m}p_{j}\left(\varphi_{j}-\mu-\langle\lambda,y_{j}\rangle\right)$
$\displaystyle=\sup_{\begin{subarray}{c}\lambda\in\mathbb{R}^{q},\mu\in\mathbb{R}\\\
\forall
j,\langle\lambda,y_{j}\rangle+\mu\leq\varphi_{j}\end{subarray}}\langle\lambda,z\rangle+\mu$

where we have swapped the $\min$ and the $\max$ using the strong duality
theorem for linear programs.

Likewise, the dual of

$\min_{\begin{subarray}{c}p\in\mathbb{R}_{+}^{m}\\\
\sum_{j=1}^{m}p_{j}y_{j}=z\end{subarray}}\langle p,\varphi\rangle$

will be the same as before but without $\mu\in\mathbb{R}$, because we have
dropped the associated constraint, i.e.

$\min_{\begin{subarray}{c}p\in\mathbb{R}_{+}^{m}\\\
\sum_{j=1}^{m}p_{j}y_{j}=z\end{subarray}}\langle
p,\varphi\rangle=\sup_{\begin{subarray}{c}\lambda\in\mathbb{R}^{q}\\\ \forall
j,\langle\lambda,y_{j}\rangle\leq\varphi_{j}\end{subarray}}\langle\lambda,z\rangle.$