\section{Dual Stochastic Method}\label{sec:met1}

In this section, we present a stochastic method for unbalanced OT based on the regularized dual formulation of \citep{chizat2015unbalanced}, which can be considered a natural generalization of \cite{seguy2017large}. The dual formulation of (\ref{eq:transport}) is given by
\[
\sup_{u, v} - \int \psi^*_1 (-u) d\mu - \int \psi^*_2(-v) d\nu
\]
subject to $u \oplus v \leq c$, where the supremum is taken over functions $u: \pazocal{X} \rightarrow [-\psi'_{1\infty}, \infty]$ and $v: \pazocal{Y} \rightarrow [-\psi'_{2\infty}, \infty]$. This is a constrained optimization problem that is challenging to solve. A standard technique for making the dual problem unconstrained is to add a strongly convex regularization term to the primal objective \citep{blondel2017smooth}, such as an entropic regularization term \citep{cuturi2013sinkhorn}:
\[
R_e(\gamma) = {\epsilon} D_{\psi_{KL}}(\gamma | \mu \otimes \nu)
\]
where ${\epsilon} > 0$. Concretely, this term has a ``smoothing" effect on the transport plan, in the sense that it encourages plans with high entropy. By the Fenchel-Rockafellar theorem, the dual of the regularized problem is given by,
\begin{equation}\label{eq:reg-dual}
W_{ub}^{\epsilon}(\mu, \nu):= \sup_{u, v} - \int \psi^*_1 (-u) d\mu
- \int \psi^*_2(-v) d\nu - \epsilon \int  e^{(u + v - c)/\epsilon} d(\mu \otimes \nu),
\end{equation}

where the supremum is taken over functions $u: \pazocal{X} \rightarrow [-\psi'_{1\infty}, \infty]$ and $v: \pazocal{Y} \rightarrow [-\psi'_{2\infty}, \infty]$, and the relationship between the primal optimizer $\gamma^*$ and dual optimizer $(u^*, v^*)$ is given by
\begin{equation} \label{eq:primal-dual}
d\gamma^* = e^{(u^* + v^* - c)/\epsilon} d(\mu \otimes \nu).
\end{equation}

Next, we rewrite (\ref{eq:reg-dual}) in terms of expectations. We assume that one has access to samples from $\mu, \nu$, and in the setting where $\mu, \nu$ are not normalized, then samples to the normalized measures $\tilde \mu, \tilde \nu$ as well as the normalization constants. Based on these assumptions, we have
\begin{equation} \label{eq:met1-exp}
W_{ub}^{\epsilon}(\mu, \nu) = \sup_{u, v} - c_\mu \mathbb{E}_{x \sim \tilde \mu} \psi^*_1 (-u(x))
- c_\nu \mathbb{E}_{y \sim \tilde \nu} \psi^*_2(-v(y)) - \epsilon c_\mu c_\nu \mathbb{E}_{x \sim \tilde \mu, y \sim \tilde \nu}  e^{(u(x) + v(y) - c(x,y))/\epsilon}.
\end{equation}

If $\psi_1^*, \psi_2^*$ are differentiable, we can parameterize $u$, $v$ with neural networks $u_\theta$, $v_\phi$ and optimize $\theta, \phi$ using stochastic gradient descent. This is described in Algorithm \ref{alg:1}. Note that this algorithm is a generalization of the algorithm of \citep{seguy2017large} from classical OT to unbalanced OT. Indeed, taking $\psi_1, \psi_2$ to be equality constraints, (\ref{eq:reg-dual}) becomes

\[
\sup_{u, v} \int u ~d\mu + \int v ~d\nu - \epsilon \int  e^{(u + v - c)/\epsilon} d(\mu \otimes \nu),
\]

which is the dual of the entropy-regularized classical OT problem. 

\begin{algorithm}[H]
\caption{SGD for Unbalanced OT}
\begin{algorithmic}
   \STATE {\bfseries Input:} Initial parameters $\theta$, $\phi$; step size $\eta$; regularization parameter $\epsilon$; constants $c_\mu, c_\nu$ and normalized measures $\tilde \mu$, $\tilde \nu$
   \STATE {\bfseries Output:} Updated parameters $\theta$, $\phi$
   \WHILE{($\theta$, $\phi$) not converged}
   
   \STATE Sample $(x_1,y_1), \cdots, (x_n, y_n)$ from $\tilde \mu \otimes \tilde \nu$
   
   \begin{align}
   \ell(\theta, \phi):= \frac{1}{n} \sum_{i=1}^n [ c_\mu \psi^*_1 (-u(x_i))
+ c_\nu \psi^*_2(-v(y_i)) + \epsilon c_\mu c_\nu  e^{(u(x_i) + v(y_i) - c(x_i,y_i))/\epsilon} ] \nonumber
   \end{align}   
   
   \STATE Update $\theta, \phi$ by gradient descent on $\ell(\theta, \phi)$
   \ENDWHILE
\end{algorithmic}
\label{alg:1}
\end{algorithm}

The dual solution ($u^*, v^*$) learned from Algorithm \ref{alg:1} can be used to reconstruct the primal solution $\gamma^*$ based on the relation in (\ref{eq:primal-dual}). Concretely, $\gamma^*$ is a transport map that indicates the amount of mass transported between every pair of points in $\pazocal{X}$ and $\pazocal{Y}$. Note that the marginals of $\gamma^*$ with respect to $\pazocal{X}$ and $\pazocal{Y}$ are not necessarily $\mu$ and $\nu$, which is where mass variation is implicitly built into the problem. Given $\gamma^*$, it is possible to also learn an ``averaged" deterministic mapping from $\pazocal{X}$ to $\pazocal{Y}$. A standard approach is to take the \emph{barycentric projection} $T: \pazocal{X} \rightarrow \pazocal{Y}$, defined as,
\[
T(x) = \min_{z \in \pazocal{Y}} \mathbb{E}_{y \sim \gamma^*(\cdot | x)} d(z, y),
\]
with respect to some distance $d: \pazocal{Y} \times \pazocal{Y} \rightarrow \mathbb{R}^+$. \cite{seguy2017large} proposed a stochastic algorithm for learning such a map from the dual solution, which we reproduce in Algorithm \ref{alg:BP}.

\begin{algorithm}[H]
\caption{Learning Barycentric Projection}
\begin{algorithmic}
   \STATE {\bfseries Input:} Learned functions $u, v$; initial $T_\theta$; distance function $d$
   \STATE {\bfseries Output:} Updated $T_\theta$
   \WHILE{$T_\theta$ not converged}
   \STATE Sample $(x_1,y_1), \cdots, (x_n, y_n)$ from $\tilde \mu \otimes \tilde \nu$
   
   \[
   \ell(\theta):= \frac{1}{n} \sum_{i=1}^n d(x_i, T_\theta(y_i)) e^{(u(x_i) + v(y_i) - c(x_i,y_i))/\epsilon} 
   \] 
   
   \STATE Update $T_\theta$ by gradient descent on $\ell(\theta)$
   \ENDWHILE
\end{algorithmic}
\label{alg:BP}
\end{algorithm}








\vspace{0.5cm}




\section{Supplement to Section \ref{sec:met2}}

\subsection{Relaxation to optimal-entropy transport}
The \emph{objectives} in (\ref{eq:ubOT}) and (\ref{eq:transport}) are equivalent if one reformulates (\ref{eq:ubOT}) in terms of $\gamma$ instead of $(T, \xi)$, where $\gamma \in \mathcal{M}_+(\pazocal{X} \times \pazocal{Y})$ is a joint measure given by
\begin{equation}\label{eq:gamma}
\gamma(C) := \int_\pazocal{X} \int_\pazocal{Z} \mathbbm{1}_C(x,T(x,z)) d\lambda(z) \xi(x) d\mu(x), \quad \forall C \in \pazocal{B}(\pazocal{X}) \times \pazocal{B}(\pazocal{Y}).
\end{equation}
Furthermore, the formulations are equivalent if one restricts the search space of (\ref{eq:transport}) to contain only those joint measures that can be specified by some $(T, \xi)$. This relation between the formulations is formalized by Lemma \ref{lem:opt-entropy}.

\begin{proof}[Proof of Lemma \ref{lem:opt-entropy}]
First we show $L_\psi(\mu, \nu) \geq \tilde W_{c_1, c_2, \psi}(\mu, \nu)$. If $L_\psi(\mu, \nu) = \infty$, this is trivial, so assume $L_\psi(\mu, \nu) < \infty$. Let $(T, \xi)$ be any solution and define $\gamma$ by (\ref{eq:gamma}). Note by this definition that
\[
\pi_{\#}^\pazocal{X} \gamma (A) = \int_\pazocal{X} \mathbbm{1}_A(x) \xi(x) d\mu(x), \quad \forall A \in \pazocal{B}(\pazocal{X}),
\]
i.e. $\xi$ is the Radon-Nikodym derivative of $\pi_{\#}^\pazocal{X} \gamma$ with respect to $\mu$. Also 
\[
\pi_{\#}^\pazocal{Y} \gamma (B) = \int_\pazocal{X} \int_\pazocal{Z} \mathbbm{1}_B(T(x,z)) d\lambda(z) \xi(x) d\mu(x) = T_\#(\xi\mu \times \lambda)(B), \quad \forall B \in \pazocal{B}(\pazocal{Y}).
\]
It follows that
\begin{align*}
&\int_\pazocal{X} \left(\int_\pazocal{Z} c_1(x, T(x,z)) d\lambda(z)\right) \xi(x) d\mu(x) + \int_\pazocal{X} c_2(\xi(x)) d\mu(x) + D_{\psi}(T_{\#}(\xi\mu \times \lambda) | \nu)\\
&\qquad= \int_{\pazocal{X} \times \pazocal{Y}} c_1(x,y) d\gamma(x,y) + \int_\pazocal{X} c_2(\xi(x)) d\mu(x) + D_{\psi}(T_{\#}(\xi\mu \times \lambda) | \nu)\\
& \qquad\quad\;\; (\text{by definition of $\gamma$, linearity and monotone convergence})\\
&\qquad= \int_{\pazocal{X} \times \pazocal{Y}} c_1(x,y) d\gamma(x,y) + \int_\pazocal{X} c_2(\frac{d\pi_{\#}^\pazocal{X} \gamma}{d\mu}) d\mu(x) + D_{\psi}(T_{\#}(\xi\mu \times \lambda) | \nu)\\
& \qquad\quad\;\;(\text{since $\xi$ is the Radon-Nikodym derivative})\\
&\qquad= \int_{\pazocal{X} \times \pazocal{Y}} c_1(x,y) d\gamma(x,y) + D_{c_2}(\pi_{\#}^\pazocal{X} \gamma | \mu) + D_{\psi}(T_{\#}(\xi\mu \times \lambda) | \nu)\\
&\qquad= \int_{\pazocal{X} \times \pazocal{Y}} c_1(x,y) d\gamma(x,y) + D_{c_2}(\pi_{\#}^\pazocal{X} \gamma | \mu) + D_{\psi}(\pi_{\#}^\pazocal{Y} \gamma | \nu)\\
&\qquad\geq \tilde W_{c, \psi_1, \psi_2}(\mu,\nu).
\end{align*}
Since this inequality holds for any $(T, \xi)$, taking the infimum over the left-hand side yields $L_\psi(\mu, \nu) \geq \tilde W_{c, \psi_1, \psi_2}(\mu,\nu)$.\\


To show $L_\psi(\mu, \nu) \leq \tilde W_{c, \psi_1, \psi_2}(\mu,\nu)$, assume $\tilde W_{c, \psi_1, \psi_2}(\mu,\nu) < \infty$ and let $\gamma$ be any solution. By the disintegration theorem, there exists a family of probability measures $\{\gamma_{y|x}\}_{x \in \pazocal{X}}$ in $\mathcal{M}_+^1(\pazocal{Y})$ such that
\[
\gamma(C) = \int_\pazocal{X} \int_\pazocal{Y} \mathbbm{1}_C(x,y) d\gamma_{y|x}(y) d\pi^\pazocal{X}_\#\gamma, \quad \forall C \in \pazocal{B}(\pazocal{X}) \times \pazocal{B}(\pazocal{Y}).
\]
Since $(\pazocal{Z}, \pazocal{B}(\pazocal{Z}), \lambda)$ is atomless, it follows from Proposition 9.1.2 and Theorem 13.1.1 in \citep{dudley2018real} that there exists a family of measurable functions $\{T_x:\pazocal{Z} \rightarrow \pazocal{Y}\}_{x \in \pazocal{X}}$ such that $\gamma_{y|x}$ is the pushforward measure of $\lambda$ under $T_x$ for all $x \in \pazocal{X}$. %
Denoting $T(x,z): (x,z) \mapsto T_x(z)$, then by a change of variables,
\[
\gamma(C) = \int_\pazocal{X} \int_\pazocal{Z} \mathbbm{1}_C(x,T(x,z)) d\lambda(z) d\pi^\pazocal{X}_\#\gamma, \quad \forall C \in \pazocal{B}(\pazocal{X}) \times \pazocal{B}(\pazocal{Y}).
\]

By hypothesis, $\pi^\pazocal{X}_\#\gamma$ is restricted to the support of $\mu$, i.e. $\pi^\pazocal{X}_\#\gamma \lll \mu$.  Let $\xi$ be the Radon-Nikodym derivative $\frac{d\pi^\pazocal{X}_\#\gamma}{d\mu}$. It follows from the Radon-Nikodym theorem that $(T, \xi)$ satisfy
\[
\gamma(C) = \int_\pazocal{X} \int_\pazocal{Z} \mathbbm{1}_C(x,T(x,z)) d\lambda(z) \xi(x) d\mu(x), \quad \forall C \in \pazocal{B}(\pazocal{X}) \times \pazocal{B}(\pazocal{Y}),
\]

which is the same relation as in (\ref{eq:gamma}). Same as before,
\[
\pi_{\#}^\pazocal{X} \gamma (A) = \int_\pazocal{X} \mathbbm{1}_A(x) \xi(x) d\mu(x), \quad \forall A \in \pazocal{B}(\pazocal{X}),
\]
and 
\[
\pi_{\#}^\pazocal{Y} \gamma (B) = \int_\pazocal{X} \int_\pazocal{Z} \mathbbm{1}_B(T(x,z)) d\lambda(z) \xi(x) d\mu(x) = T_\#(\xi\mu \times \lambda)(B), \quad \forall B \in \pazocal{B}(\pazocal{Y}).
\]
It then follows that
\begin{align*}
& \int_{\pazocal{X} \times \pazocal{Y}} c_1 d\gamma + D_{c_2}(\pi_{\#}^\pazocal{X} \gamma | \mu) + D_{\psi}(\pi_{\#}^\pazocal{Y} \gamma  | \nu)\\
&= \int_\pazocal{X} \left(\int_\pazocal{Y} c_1(x, y) d\gamma_{y|x}(y)\right) d\pi^\pazocal{X}_\#\gamma + D_{c_2}(\pi_{\#}^\pazocal{X} \gamma | \mu) + D_{\psi}(\pi_{\#}^\pazocal{Y} \gamma  | \nu)\\
&\text{(by Fubini's Theorem for Markov kernels)}\\
&=\int_\pazocal{X} \left(\int_\pazocal{Z} c(x, T(x,z)) d\lambda(z)\right) d\pi^\pazocal{X}_\#\gamma + D_{c_2}(\pi_{\#}^\pazocal{X} \gamma | \mu) + D_{\psi}(\pi_{\#}^\pazocal{Y} \gamma  | \nu)\\
&\text{(by change of variables)}\\
&=\int_\pazocal{X} \left(\int_\pazocal{Z} c(x, T(x,z)) d\lambda(z)\right) d\pi^\pazocal{X}_\#\gamma + \int c_2(\frac{d\pi_{\#}^\pazocal{X} \gamma }{d\mu}) d\mu(x) + D_{\psi}(\pi_{\#}^\pazocal{Y} \gamma  | \nu)\\
&=\int_\pazocal{X} \left(\int_\pazocal{Z} c(x, T(x,z)) d\lambda(z)\right) \xi(x) d \mu(x) + \int c_2(\xi(x)) d\mu(x) + D_{\psi}(\pi_{\#}^\pazocal{Y} \gamma  | \nu)\\
&\text{(by definition of $\xi$)}\\
&=\int_\pazocal{X} \left(\int_\pazocal{Z} c(x, T(x,z)) d\lambda(z)\right) d \mu(x) + \int c_2(\xi(x)) d\mu(x) + D_{\psi}(T_{\#}(\xi\mu \times \lambda)  | \nu)\\
&\geq L_\psi( \mu, \nu).
\end{align*}
Since this inequality holds for any $\gamma$, this implies that $\tilde W_{c_1, c_2, \psi}(\mu, \nu) \geq L_\psi(\mu, \nu)$, which completes the proof.
\end{proof}

Due to the near equivalence of the formulations, several theoretical results for (\ref{eq:ubOT}) follow from the analysis of optimal entropy-transport by \citet{liero2018optimal}, such as the following existence and uniqueness result: 

\begin{proposition} \label{prop:wellposed} Suppose $L_\psi(\mu,\nu) < \infty$, $c_2$ is convex and lower semi-continuous, and $(\pazocal{Z}, \pazocal{B}(\pazocal{Z}), \lambda)$ is atomless. If (i) $c_1$ has compact sublevel sets in $\pazocal{X} \times \pazocal{Y}$ and $c_{2\infty}' + \psi_\infty' > 0$, or (ii) $c_{2\infty}' = \psi_\infty' = \infty$ then a minimizer of $L_\psi(\mu,\nu)$ exists. If $c_2, \psi$ are strictly convex, $\psi_\infty' = \infty$ and $c_1$ satisfies Corollary 3.6 \citep{liero2018optimal}, then the joint measure $\gamma$ specified by any minimizer of $L_\psi(\mu,\nu)$ is unique.
\end{proposition}

\begin{proof}[Proof of Proposition \ref{prop:wellposed}]
Note that $\tilde W_{c_1, c_2, \psi}(\mu, \nu)$ is equivalent to $W_{c_1, c_2, \psi}(\mu, \nu)$ when $\pazocal{X}$ is restricted to the support of $\mu$. If $c_1, c_2, \psi$ satisfy (i) or (ii), they also satisfy (i) and (ii) when $\pazocal{X}$ is restricted to the support of $\mu$. By Theorem 3.3 of \citep{liero2018optimal}, $\tilde W_{c_1, c_2, \psi}(\mu, \nu)$ has a minimizer. It follows from the construction of the proof of Lemma \ref{lem:opt-entropy} that a minimizer of $L_\psi(\mu, \nu)$ also exists. For uniqueness, if $\psi_\infty'= \infty$, then it follows from Lemma 3.5 of \citep{liero2018optimal} and the fact that minimizers are restricted to $\pazocal{G}$ that the marginals $\pi^\pazocal{X}_\#\gamma$, $\pi^\pazocal{Y}_\#\gamma$ are uniquely determined for any solution $\gamma$ of $\tilde W_{c_1, c_2, \psi}(\mu, \nu)$. The uniqueness of $\gamma$ then follows from the proof of Corollary 3.6 in \citep{liero2018optimal}. It follows from the construction of the proof of Lemma \ref{lem:opt-entropy} that the product measure generated by the minimizers of $L_\psi(\mu, \nu)$ is unique, which completes the proof.
\end{proof}

For certain cost functions and divergences, it can be shown that $L_{\psi}$ defines a proper metric between positive measures $\mu$ and $\nu$, i.e. taking $c_2, \psi$ to be entropy functions corresponding to the KL-divergence and $c_1 = −\log \cos^2_+(d(x,y))$, then $L_\psi(\mu, \nu)$ corresponds to the Hellinger-Kantorovich \citep{liero2018optimal} or the Wasserstein-Fisher-Rao \citep{chizat2018interpolating} metric between positive measures $\mu$ and $\nu$. 

Based on Lemma \ref{lem:opt-entropy}, the theoretical analysis of \cite{liero2018optimal}, and standard results on constrained optimization, it can be shown that for an appropriate and sufficiently large choice of divergence penalty, solutions of the relaxed problem (\ref{eq:ubOT}) converge to solutions of the original problem (\ref{eq:ubOT_exact}) (Theorem \ref{prop:relaxed}). 

\begin{proof}[Proof of Theorem \ref{prop:relaxed}] Since $\zeta^k \psi(s)$ converges pointwise to the equality constraint $\iota_{=}(s)$, which is $0$ for $s=1$ and $\infty$ otherwise, by Lemma 3.9 in \citep{liero2018optimal}, we have that $\liminf_{k \rightarrow \infty} \tilde W_{c_1, c_2, \zeta^k \psi}(\mu, \nu) \geq \tilde W_{c_1, c_2, \iota_{=}}(\mu, \nu)$. Additionally, $\tilde W_{c_1, c_2, \zeta^k \psi}(\mu, \nu) \leq \tilde W_{c_1, c_2, \iota_{=}}(\mu, \nu)$ for any value of $k$ since for any minimizer~$\gamma$ of $\tilde W_{c_1, c_2, \iota_{=}}(\mu, \nu)$, it holds that $\pi^\pazocal{Y}_\#\gamma = \nu$. Hence
\begin{align*}
\tilde W_{c_1, c_2, \iota_{=}}(\mu, \nu) = \int c_1 d\gamma + D_{c_2}(\pi^\pazocal{X}_\#\gamma | \mu) + D_{\zeta^k \psi_2}(\pi^\pazocal{Y}_\#\gamma| \nu) \geq \tilde W_{c_1, c_2, \zeta^k \psi}(\mu, \nu),
\end{align*}
for all $k$. Therefore, $\lim_{k \rightarrow \infty} \tilde W_{c_1, c_2, \zeta^k \psi}(\mu, \nu) = \tilde W_{c_1, c_2, \iota_{=}}(\mu, \nu)$, which then by Lemma \ref{lem:opt-entropy} implies the first part of the proposition.

For the second part, by the hypothesis we have that $\tilde W_{c_1, c_2, \iota_{=}}(\mu, \nu) = C < \infty$ and as a consequence $\tilde W_{c_1, c_2, \zeta^k \psi}(\mu, \nu) \leq C$ for all $k$. Hence, by Proposition 2.10 in \citep{liero2018optimal}, the sequence of minimizers $\gamma^k$ is bounded. If the assumptions of Proposition \ref{prop:wellposed} are satisfied, then the sequence $\gamma^k$ is equally tight. For assumption (ii) this follows by Proposition 2.10 in \citep{liero2018optimal} and for assumption (i) this follows by the Markov inequality: for any $\lambda>0$, 
\[
\gamma^k(\{ (x,y) \in \pazocal{X} \times \pazocal{Y} | c_1(x,y) > \lambda \}) \leq \frac{1}{\lambda} \int c_1 d\gamma^k \leq \frac{C}{\lambda}.
\]

Since $\gamma^k$ are bounded and equally tight, by an extension of Prokhorov's theorem (Theorem 2.2 of \citep{liero2018optimal}), there exists a subsequence of $\gamma^k$ that is weakly convergent to some $\bar{\gamma}$. Then by lower semicontinuity, we obtain that
\[
\int c_1 d\bar{\gamma} + D_{c_2}(\pi^\pazocal{X}_\#\bar{\gamma} | \mu) + \limsup_{k \rightarrow \infty} D_{\zeta^k \psi_2}(\pi^\pazocal{Y}_\#\gamma^k | \nu) \leq \tilde W_{c_1, c_2, \iota_{=}}(\mu, \nu) = C
\]
Since $D_{\zeta^k \psi}(\pi^\pazocal{Y}_\#\gamma^k | \nu) = \zeta^k D_{\psi}(\pi^\pazocal{Y}_\#\gamma^k | \nu) \geq 0$ and $\zeta^k \rightarrow \infty$, for the left side to be finite, $D_{\psi}(\pi^\pazocal{Y}_\#\gamma^k | \nu)$ must converge to $0$, so $D_{\psi}(\pi^\pazocal{Y}_\#\bar{\gamma} | \nu) = 0$ by lower semicontinuity. Therefore, $\bar{\gamma}$ is a minimizer of $\tilde W_{c_1, c_2, \iota_{=}}(\mu, \nu)$. By construction of the proof of Lemma \ref{lem:opt-entropy}, $\gamma^k$ is equivalent to the product measure induced by minimizers of $L_{\zeta^k \psi}(\mu, \nu)$, which implies the second part of the proposition. 
\end{proof}


\subsection{Convex conjugate form of divergences}

In this section, we present the convex conjugate form of $\psi$-divergence used to rewrite the main objective as a min-max problem.

\begin{lemma} \label{lem:vc2}
 For non-negative finite measures $P,Q$ over $\pazocal{T} \subset \mathbb{R}^d$, it holds that
\begin{equation}
D_{\psi}(P|Q) \geq \sup_{f \in \pazocal{F}} \int f dP - \psi^*(f) dQ
\end{equation}
where $\pazocal{F}$ is a subset of measurable functions $\{f: \pazocal{T} \rightarrow (-\infty, \psi'_{\infty}] \}$. Equality holds if and only if $\exists f \in \pazocal{F}$ such that (i) the restriction of $f$ to the support of $Q$ belongs to the subdifferential of $\psi(\frac{dP}{dQ})$, i.e. the Radon-Nikodym derivative of $P$ with respect to $Q$ and (ii) $f=\psi'_{\infty}$ over the support of $P^\perp$.
\end{lemma}

We provide a simple proof of this result. A similar result under stronger assumptions was shown in \cite{nguyen2008estimating} and used by \cite{nowozin2016f} for generative modeling. A rigorous proof can be found in \cite{liero2018optimal}.

\begin{proof}[Proof of Lemma \ref{lem:vc2}]
	Note that
\begin{align*}
D_\psi (P | Q) &= \psi_\infty' P_\perp (\pazocal{T}) + \int_{\pazocal{T}} \psi \left( \frac{dP}{dQ} \right) dQ\\
&= \psi_\infty' P_\perp (\pazocal{T}) + \int_{\pazocal{T}} \sup_{\xi \in \mathbb{R}} \{ \xi  \frac{dP}{dQ} - \psi^*(\xi) \} dQ\\ 
&\text{(by defintion of convex conjugate)}\\
&= \psi_\infty' P_\perp (\pazocal{T}) + \int_{\pazocal{T}} \sup_{\xi \in (-\infty, \psi_\infty']} \{ \xi  \frac{dP}{dQ} - \psi^*(\xi) \} dQ\\ 
&\text{(by Lemma \ref{lem:psi}) below}\\
&= \int_{\pazocal{T}} \sup_{\xi \in (-\infty, \psi_\infty']} \{\xi dP_\perp + \xi  \frac{dP}{dQ}dQ - \psi^*(\xi)dQ \} \\ 
&= \sup_{f \in \pazocal{F}} \int_{\pazocal{T}}  f dP - \psi^*(f) dQ.\\
\end{align*}
By first-order optimality conditions, the optimal $f$ over the support of $Q$ is obtained when $\frac{dP}{dQ}$ belongs to the subdifferential of $\psi^*(f)$, or equivalently when $f$ belongs to the subdifferential of $\psi(\frac{dP}{dQ})$. It is straightforward to see that the optimal $f$ over the support of $P_\perp$ is equal to $\psi_\infty'$, which completes the proof.
\end{proof}

\begin{lemma}\label{lem:psi} If $\xi > \psi_\infty':= \lim_{s \rightarrow \infty} \frac{\psi(s)}{s}$, then $\psi^*(\xi) = \infty$.
\end{lemma}
\begin{proof}
\begin{align*}
\psi^*(\xi) &= \sup_{s \in \mathbb{R}} s\xi - \psi(s)\\
&\geq \lim_{s \rightarrow \infty} s(\xi-\frac{\psi(s)}{s})\\
&= \infty ~\text{if} ~\xi > \psi_\infty'
\end{align*}
\end{proof}

\iffalse
\subsection{Discussion of related work}

Since the original GAN paper by \cite{goodfellow2014generative}, numerous variants of GANs have been proposed for learning transport maps between distributions. For example, numerous transport problems in image translation \citep{mirza2014conditional,zhu2017unpaired, yi2017dualgan}, natural language translation \citep{he2016dual}, domain adaptation \citep{bousmalis2017unsupervised} and biological data integration \citep{amodio2018magan} have been tackled using variants of GANs, with strategies such as conditioning or cycle-consistency employed to enforce correspondence between original and transported samples. However, all these methods conserve mass between the source and target and therefore cannot handle mass variation. This limits applications to large-scale problems in population modeling, particularly in computational biology \citep{schiebinger2017reconstruction}.

Algorithm \ref{alg:TAG} proposes a generative-adversarial strategy that can be used for unbalanced transport, based on the idea of learning both a transport map and a scaling factor. Our derivation of Algorithm \ref{alg:TAG} uses a convex conjugate form of divergences presented in Lemma \ref{lem:vc2}, which enables us to minimize many divergences using alternating gradient updates similar to GAN training.  A similar inequality with stronger assumptions proposed by \cite{nguyen2008estimating}  was used by \cite{nowozin2016f} to propose a large class of generative models that generalize GANs. An important distinction of our work is that we consider divergence minimization between measures that are not probability measures, which requires additional assumptions on the choice of divergence; see Lemma \ref{lem:div}. 
\fi


\section{Practical Considerations for Numerical Experiments} \label{sec:prac}

{\bf Choice of cost functions.} Proposition \ref{prop:wellposed} gives sufficient conditions on $c_1, c_2$ for the problem to be well-posed. In practice, it is often convenient the cost of transport, $c_1$, to be some measurement of correspondence between $\pazocal{X}$ and $\pazocal{Y}$. For example, we can take $c_1(x,y)$ to be the Euclidean distance between $x$ and $y$ after mapping them to some common feature space. For the cost of mass adjustment, $c_2$, it is generally sensible to choose some convex function that vanishes at $1$ (i.e. no mass adjustment) and such that $\lim_{x \rightarrow 0} c_2(x) = \lim_{x \rightarrow \infty} c_2(x) = \infty$ to prevent $\xi$ from becoming too small or too large. Any of the entropy functions shown in Table \ref{tab:div} are reasonable choices. 

{\bf Choice of $\psi$.}
In \cite{nowozin2016f}, it was shown that any $\psi$-divergence could be used to train generative models, i.e. to match a generated distribution $P$ to a true data distribution $Q$. This is due to Jensen's inequality: for any convex lower semi-continous entropy function $\psi$, $D_{\psi}(P|Q)$ is uniquely minimized when $P=Q$, where $P,Q$ are probability measures. However, this does not generally hold when $P,Q$ are not probability measures, as illustrated by the following example. 

\begin{example} In the original GAN paper, the discrminative objective,

\[
\sup_f \int  \log f(x) dP(x) -  \int \log(1-f(x)) dQ(x),
\]
corresponds to $D_\psi(P|Q)$ with $\psi(s) = s \log s - (s+1) \log(s + 1)$ \citep{nowozin2016f}. If $P,Q$ are probability measures, this divergence is equivalent to the Jensen-Shannon divergence and is minimized when $P=Q$. If $P,Q$ are non-negative measures with unconstrained total mass, the divergence is minimized when $P = \infty$ and $Q = 0$.
\end{example}

When $P,Q$ are not probability measures, we require an additional constraint on $\psi$ to ensure that divergence minimization matchces $P$ to $Q$:

\begin{lemma}\label{lem:div} Suppose $P, Q$ are non-negative finite measures over $T \subseteq \mathbb{R}^n$. If $\psi(s)$ attains a unique minimum at $s=1$ with $\psi(1)=0$ and $\psi'_\infty > 0$, then $D_\psi(P|Q) = 0 \Rightarrow P=Q$. Otherwise, then $P \neq Q$ in general when $D_\psi(P|Q)$ is minimized.
\end{lemma}

\begin{proof} Suppose $\psi(s)$ attains a unique minimum at $s=1$ with $\psi(1)=0$, $\psi'_\infty > 0$, and $P \neq Q$ over a region with positive measure. It is straightforward to see by the definition in (\ref{eq:div}) that $D_\psi(P|Q) > 0$, since at least one of the two terms will be strictly positive. Therefore, the first statement holds. For the second statement, suppose either $\psi(s)$ does not attain a unique minimum at $s=1$ or $\psi'_\infty \leq 0$. If $\psi(s)$ attains a minimum at some $s' \neq 1$, then taking $P = s'Q$ results in a divergence that is equal to or less than $P = Q$. If $\psi'_\infty \leq 0$, then letting $P = Q + P_\perp$ where $P_\perp$ is a positive measure orthogonal to $Q$ results in a divergence that is equal to or less than $P=Q$.
\end{proof}

Table \ref{tab:div} provides some examples of $\psi$ corresponding to common divergences that can be used for unbalanced OT.

{\bf Choice of $f$.} According to Lemma \ref{lem:vc2}, $f$ should belong to a class of functions that maps from $\pazocal{Y}$ to $(-\infty, \psi_\infty']$. In practice, this can be enforced by parameterizing $f$ using a neural network with a final layer that maps to the correct range, also known as an output activation layer \citep{nowozin2016f}. Table \ref{tab:div} provides some examples of activation layers that can be used. 

\begin{table}
\tiny
\begin{center}
 \begin{tabular}{| c | c | c | c | c | c | c |} 
 \hline
 Name & $\psi(s) $ & $D_\psi(P|Q)$ & $\psi^*(s)$ & $\psi_\infty'$ & Activation Layer \\
 \hline
 Kullback-Leibler (KL) & $s \log s - s + 1$ & $\int \log\frac{dP}{dQ} dP - \int dP + \int dQ$ & $e^{s}$ & $\infty$ & N/A \\ 
 \hline
 Pearson $\chi^2$ & $ (s - 1)^2$ & $\int (\frac{dP}{dQ} - 1)^2 dQ$ & $\frac{s^2}{4} + s$ & $\infty$ & N/A\\ 
 \hline
 Hellinger & $ (\sqrt{s} - 1)^2$ & $\int (\sqrt{\frac{dP}{dQ}} - 1)^2 dQ$ & $\frac{s}{1-s}$ & $1$ & $1-e^s$ \\ 
 \hline
 Jensen-Shannon & $s \log s - (s+1) \log \frac{s+1}{2}$ & $\frac{1}{2}D_{KL}(P | \frac{P+Q}{2}) + \frac{1}{2} D_{KL}(Q | \frac{P+Q}{2})$ & $-\log(2-e^s)$ & $\log 2$ & $\log(2) - \log(1 + e^{-s}))$ \\ 
 \hline

\end{tabular}

\end{center}
 \caption{Table of some common $\psi$-divergences, associated entropy functions $\psi$, and convex conjugates $\psi^*$ for Algorithm \ref{alg:TAG}, partly adapted from \citep{nowozin2016f}.}
 \label{tab:div}
\end{table}

{\bf Choice of neural architectures.} For our experiments in Section \ref{sec:exp}, we used fully-connected feedforward networks with 3 hidden layers and ReLU activations. For $T$, the output activation layer was a sigmoid function to map the final pixel brightness to the range $(0,1)$. For $\xi$, the output activation layer was a softplus function to map the scaling factor weight to the range $(0, 
\infty)$.

{\bf Gradient penalties.} The training of GANs using alternating stochastic gradient descent is not guaranteed to converge locally \citep{mescheder2018training}. Stability can be improved using regularization in the form of added instance noise \citep{arjovsky2017towards} or gradient penalties \citep{gulrajani2017improved}. We observed that enforcing gradient penalities on $\xi$ and $f$, while not necessary, improved the stability of Algorithm~\ref{alg:TAG}. A gradient penalty on $\xi$ effectively restricts the search to sufficiently smooth $\xi$, which is reasonable in practice considering that similar regions of the source measure often have similar rates of growth or contraction. A gradient penalty on $f$ changes the nature of the relaxation from (\ref{eq:ubOT_exact}) to (\ref{eq:ubOT}): the right-hand side of (\ref{eq:var-div}) is no longer equivalent to the $\psi$-divergence, but is rather a lower-bound with a relation to bounded Lipschitz metrics \citep{gulrajani2017improved}. In this case, while the problem formulation is not equivalent to optimal entropy-transport, it is still a valid relaxation of unbalanced Monge OT in (\ref{eq:ubOT_exact}).

{\bf Choice of $\lambda$.} One can take $\lambda$ to be the standard Gaussian measure if a stochastic mapping is desired, similar to \cite{almahairi2018augmented}. If a deterministic mapping is desired, then $\lambda$ is set to a deterministic distribution. 

{\bf Improved training dynamics.}
Recall that the objective function for our alternating gradient updates is
\begin{align*}
   \ell(\theta, \phi, \omega):= \frac{1}{n} \sum_{i=1}^n [ c_\mu c_1(x_i, T_\theta(x_i, z_i))\xi_\phi(x_i) &+ c_\mu c_2(\xi_\phi(x_i)) \\
   &+ c_\mu \xi_\phi(x_i) f_\omega(T_\theta(x_i, z_i)) - c_\nu \psi^*(f_\omega(y_i)) ]. \nonumber
\end{align*}
Early in training, $\xi_\phi$ can become very small for some $x_i$ as none of the transported samples resemble samples from the target distribution. As a result, $T_\theta$ may improve very slowly for some inputs $x_i$. One way to address this issue without changing the fixed point is to update $\xi_\phi, f_\omega$ using the above objective and update $T_\theta$ using the following objective.
\begin{align*}
   \ell(\theta, \phi, \omega):= \frac{1}{n} \sum_{i=1}^n c_\mu [ c_1(x_i, T_\theta(x_i, z_i)) + f_\omega(T_\theta(x_i, z_i))] \nonumber
\end{align*}
Note that we have omitted terms in the original objective that do not include $T_\theta$, and which therefore do not affect the gradient update. For the terms that remain, the difference is that we are rescaling the contribution of each sample $x_i$ by $1/\xi(x_i)$. As long as $\xi(x_i) > 0$, this has the effect of rescaling the gradient update of the loss function with respect to each $T_\theta(x_i, z_i)$ \emph{without} changing the direction of the update.




\section{Introduction}

We consider the problem of unbalanced optimal transport: given two measures, find a cost-optimal way to transform one measure to the other using a combination of mass variation and transport. Such problems arise, for example, when modeling the transformation of a source population into a target population (Figure \ref{fig:pop_example}a). In this setting, one needs to model mass transport to account for the features that are evolving, as well as local mass variations to allow sub-populations to become more or less prominent in the target population \citep{schiebinger2017reconstruction}.

Classical optimal transport (OT) considers the problem of pushing a source to a target distribution in a way that is optimal with respect to some transport cost without allowing for mass variations. Modern approaches are based on the Kantorovich formulation \citep{kantorovich1942translocation}, which seeks the optimal probabilistic coupling between measures and can be solved using linear programming methods for discrete measures. Recently, 
\cite{cuturi2013sinkhorn} showed that regularizing the objective using an entropy term allows the dual problem to be solved more efficiently using the Sinkhorn algorithm. Stochastic methods based on the dual objective have been proposed for the continuous setting \citep{genevay2016stochastic, seguy2017large, arjovsky2017wasserstein}. Optimal transport has been applied to many areas, such as computer graphics \citep{ferradans2014regularized, solomon2015convolutional} and domain adaptation \citep{courty2014domain, courty2017optimal}. 

In many applications where a transport cost is not available, transport maps can also be learned using generative models such as generative adversarial networks (GANs) \citep{goodfellow2014generative}, which push a source distribution to a target distribution by training against an adversary. Numerous transport problems in image translation \citep{mirza2014conditional,zhu2017unpaired, yi2017dualgan}, natural language translation \citep{he2016dual}, domain adaptation \citep{bousmalis2017unsupervised} and biological data integration \citep{amodio2018magan} have been tackled using variants of GANs, with strategies such as conditioning or cycle-consistency employed to enforce correspondence between original and transported samples. However, all these methods conserve mass between the source and target and therefore cannot handle mass variation.

Several formulations have been proposed for extending the theory of OT to the setting where the measures can have unbalanced masses \citep{chizat2015unbalanced, chizat2018interpolating, kondratyev2016new, liero2018optimal, frogner2015learning}. In terms of numerical methods, a class of scaling algorithms \citep{chizat2016scaling} that generalize the Sinkhorn algorithm for balanced OT have been developed for approximating the solution to \emph{optimal entropy-transport} problems; this formulation of unbalanced OT by \citet{liero2018optimal} corresponds to the Kantorovich OT problem in which the hard marginal constraints are relaxed using divergences to allow for mass variation. In practice, these algorithms have been used to approximate unbalanced transport plans between discrete measures for applications such as computer graphics \citep{chizat2016scaling}, tumor growth modeling \citep{chizat2017tumor} and computational biology \citep{schiebinger2017reconstruction}.  However, while optimal entropy-transport allows mass variation, it cannot  explicitly model it, and there are currently no methods that can perform unbalanced OT between \emph{continuous} measures. 

\begin{figure}[t]\label{fig:pop_example}
\centering
\vspace{-0.2cm}
\subfigure[]{\includegraphics[scale=.15]{pop_example.png}}
\subfigure[]{\includegraphics[scale=.15]{monge-transport.png}}
\vspace{-0.2cm}
\caption{(a) Illustration of the problem of modeling the transformation of a source population $\mu$ to a target population $\nu$. In this example, one sub-population is growing more rapidly than the others. (b) Schematic of Monge-like formulations of unbalanced optimal transport. The objective is to learn a transport map $T$ (for transporting mass) and scaling factor $\xi$ (for mass variation) to push the source $\mu$ to the target $\nu$, using a deterministic transport map (top) \citep{chizat2015unbalanced} or a stochastic transport map (bottom). }
\end{figure}

{\bf Contributions.} Inspired by the recent successes of GANs for high-dimensional transport problems, we present a novel framework for unbalanced optimal transport that directly models mass variation in addition to transport. Concretely, our contributions are the following:
\begin{itemize}
\item We propose to solve a Monge-like formulation of unbalanced OT, in which the goal is to learn a stochastic transport map and scaling factor to push a source to a target measure in a cost-optimal manner. This generalizes the unbalanced Monge OT problem by \citet{chizat2015unbalanced}.
\item By relaxing this problem, we obtain an alternative form of the optimal entropy-transport problem by \citet{liero2018optimal}, which confers desirable theoretical properties. 
\item We develop scalable methodology for solving the relaxed problem. Our derivation uses a convex conjugate representation of divergences, resulting in an alternating gradient descent method similar to GANs \citep{goodfellow2014generative}.
\item We demonstrate in practice how our methodology can be applied towards population modeling using the MNIST and USPS handwritten digits datasets, the CelebA dataset, and 
a recent single-cell RNA-seq dataset from zebrafish embrogenesis.
\end{itemize}

In addition to these main contributions, for completeness we also propose a new scalable method (Algorithm~\ref{alg:1}) in the Appendix for solving the optimal-entropy transport problem by \citet{liero2018optimal} in the continuous setting. The algorithm extends the work of \citet{seguy2017large} to unbalanced OT and is a scalable alternative to the algorithm of \citet{chizat2016scaling} for very large or continuous datasets.


\section{Preliminaries}

{\bf Notation.} Let $\pazocal{X}, \pazocal{Y} \subseteq \mathbb{R}^n$ be topological spaces and let $\pazocal{B}$ denote the Borel $\sigma$-algebra. 
Let $\mathcal{M}_+^1(\pazocal{X}), \mathcal{M}_+(\pazocal{X})$ denote respectively the space of probability measures and finite non-negative measures over $\pazocal{X}$. For a measurable function $T$, let $T_{\#}$ denote its pushforward operator: if $\mu$ is a measure, then $T_{\#} \mu$ is the pushforward measure of $\mu$ under $T$. Finally, let $\pi^\pazocal{X}$, $\pi^\pazocal{Y}$ be functions that project onto $\pazocal{X}$ and $\pazocal{Y}$; for a joint measure $\gamma \in \mathcal{M}_+(\pazocal{X} \times \pazocal{Y})$, $\pi^\pazocal{X}_\#\gamma$ and $\pi^\pazocal{Y}_\#\gamma$ are its marginals with respect to $\pazocal{X}$ and $\pazocal{Y}$ respectively.

{\bf Optimal transport (OT)} addresses the problem of transporting between measures in a cost-optimal manner. \citet{monge1781memoire} formulated this problem as a search over deterministic transport maps.
Specifically, given $\mu \in \mathcal{M}_+^1(\pazocal{X}), \nu \in \mathcal{M}_+^1(\pazocal{Y})$ and a cost function $c: \pazocal{X} \times \pazocal{Y} \rightarrow \mathbb{R}^+$, Monge OT seeks a measurable function $T: \pazocal{X} \rightarrow \pazocal{Y}$ minimizing
\begin{equation}
\inf_{T} \int_{\pazocal{X}} c(x,T(x)) ~d\mu(x)
\end{equation}
subject to the constraint $T_{\#}\mu = \nu$. While the optimal $T$ has an intuitive interpretation as an optimal transport map, the Monge problem is non-convex and not always feasible depending on the choices of $\mu$ and $\nu$. The \emph{Kantorovich OT} problem is a convex relaxation of the Monge problem that formulates OT as a search over probabilistic transport plans. Given $\mu \in \mathcal{M}_+^1(\pazocal{X}), \nu \in \mathcal{M}_+^1(\pazocal{Y})$ and a cost function $c: \pazocal{X} \times \pazocal{Y} \rightarrow \mathbb{R}^+$, Kantorovich OT seeks a joint measure $\gamma \in \mathcal{M}_+^1(\pazocal{X} \times \pazocal{Y})$ subject to $\pi^\pazocal{X}_\#\gamma=\mu$ and $\pi^\pazocal{Y}_\#\gamma=\nu$ minimizing
\begin{equation}\label{eq:MK}
\begin{split}
W(\mu, \nu) := \inf_{\gamma} \int_{\pazocal{X} \times \pazocal{Y}} c(x,y) ~d\gamma(x,y).
\end{split}
\end{equation}
Note that the conditional probability distributions $\gamma_{y|x}$ specify stochastic maps from $\pazocal{X}$ to $\pazocal{Y}$ and can be considered a ``one-to-many" version of the deterministic map from the Monge problem. In terms of numerical methods, the relaxed problem is a linear program that is always feasible and can be solved in $O(n^3)$ time for discrete $\mu, \nu$. \cite{cuturi2013sinkhorn} recently showed that introducing entropic regularization results in a simpler dual optimization problem that can be solved efficiently using the Sinkhorn algorithm. Based on the entropy-regularized dual problem, \cite{genevay2016stochastic} and \cite{seguy2017large} proposed stochastic algorithms for computing transport plans that can handle continuous measures. 

{\bf Unbalanced OT.} Several formulations that extend classical OT to handle mass variation have been proposed \citep{chizat2015unbalanced, chizat2018interpolating, kondratyev2016new}. Existing numerical methods are based on a Kantorovich-like formulation known as \emph{optimal-entropy transport} \citep{liero2018optimal}. This formulation is obtained by relaxing the marginal constraints of (\ref{eq:MK}) using divergences as follows: given two positive measures $\mu \in \mathcal{M}_+(\pazocal{X})$ and $\nu \in \mathcal{M}_+(\pazocal{Y})$ and a cost function $c: \pazocal{X} \times \pazocal{Y} \rightarrow \mathbb{R}^+$, optimal entropy-transport finds a measure $\gamma \in \mathcal{M}_+(\pazocal{X} \times \pazocal{Y})$ that minimizes
\begin{equation}\label{eq:transport}
\begin{split}
W_{ub}(\mu, \nu) := \inf_\gamma \int_{\pazocal{X} \times \pazocal{Y}} c(x,y) ~d\gamma(x,y) + D_{\psi_1}(\pi^\pazocal{X}_\#\gamma| \mu) 
+ D_{\psi_2}(\pi^\pazocal{Y}_\#\gamma| \nu),
\end{split}
\end{equation}
where $D_{\psi_1}$, $D_{\psi_2}$ are $\psi$-divergences induced by $\psi_1, \psi_2$. The $\psi$-divergence between non-negative finite measures $P,Q$ over $\pazocal{T} \subseteq \mathbb{R}^d$ induced by a lower semi-continuous, convex entropy function $\psi: \mathbb{R} \rightarrow \mathbb{R} \cup \{ \infty \}$~is
\begin{equation}\label{eq:div}
D_{\psi}(P | Q) := \psi'_{\infty}P^\perp(\pazocal{T}) + \int_{\pazocal{T}} \psi \left(\frac{dP}{dQ} \right) dQ, 
\end{equation} 
where $\psi'_{\infty} := \lim_{s \rightarrow \infty} \frac{\psi(s)}{s}$ and $\frac{dP}{dQ} Q + P^\perp$ is the Lebesgue decomposition of $P$ with respect to $Q$. Note that mass variation is allowed since the marginals of $\gamma$ are not constrained to be $\mu$ and $\nu$. 
In terms of numerical methods,
the state-of-the-art in the discrete setting is a class of iterative scaling algorithms \citep{chizat2016scaling} that generalize the Sinkhorn algorithm for computing regularized OT plans \citep{cuturi2013sinkhorn}.
There are no practical algorithms for unbalanced OT between continuous measures, especially in high-dimensional spaces.

\begin{figure}[t]
\centering
\vspace{-0.3cm}
\subfigure[]{\includegraphics[scale=.25]{ubOT_gauss_1.png}}
\subfigure[]{\includegraphics[scale=.25]{ubOT_gauss_2.png}}
\vspace{-0.2cm}
\caption{{\bf Motivating examples for Unbalanced Monge OT.} }
\vspace{-0.1cm}
\label{fig:syn_data}
\end{figure}

\section{Scalable Unbalanced OT using GANs} \label{sec:met2}
In this section, we propose the first algorithm for unbalanced OT that directly models mass variation and can be applied towards transport between high-dimensional continuous measures. The starting point of our development is the following Monge-like formulation of unbalanced OT, in which the goal is to learn a stochastic transport map and scaling factor to push a source to a target measure in a cost-optimal manner.

{\bf Unbalanced Monge OT.} Let $c_1: \pazocal{X} \times \pazocal{Y} \rightarrow \mathbb{R}^+$ be the cost of transport and $c_2: \mathbb{R}^+ \rightarrow \mathbb{R}^+$ the cost of mass variation. Let the probability space $(\pazocal{Z}, \pazocal{B}(\pazocal{Z}), \lambda)$ be the source of randomness in the transport map $T$. Given two positive measures $\mu \in \mathcal{M}_+(\pazocal{X})$ and $\nu \in \mathcal{M}_+(\pazocal{Y})$, we seek a transport map $T: \pazocal{X} \times \pazocal{Z} \rightarrow \pazocal{Y}$ and a scaling factor $\xi: \pazocal{X} \rightarrow \mathbb{R}^+$ minimizing
\begin{align}\label{eq:ubOT_exact}
&L(\mu, \nu) := \inf_{T, \xi} \int_\pazocal{X} \int_\pazocal{Z} c_1(x, T(x,z)) d\lambda(z) \xi(x) d\mu(x) + \int_\pazocal{X} c_2(\xi(x)) d\mu(x),
\end{align}
subject to the constraint $T_{\#}(\xi\mu \times \lambda) = \nu$. Concretely, the first and second terms of (\ref{eq:ubOT_exact}) penalize the cost of mass transport and variation respectively, and the equality constraint ensures that $(T, \xi)$ pushes $\mu$ to $\nu$ exactly. A special case of (\ref{eq:ubOT_exact}) is the unbalanced Monge OT problem by \citet{chizat2015unbalanced}, which employs a deterministic transport map (Figure \ref{fig:pop_example}b). We consider the more general case of stochastic (i.e. one-to-many) maps because it is a more suitable model for many practical problems. For example, in cell biology, it is natural to think of one cell in a source population as potentially giving rise to multiple cells in a target population. In practice, one can take $\pazocal{Z} = \mathbb{R}^n$ and $\lambda$ to be the standard Gaussian measure if a stochastic map is desired; otherwise $\lambda$ can be set to a deterministic distribution. The following are examples of problems that can be modeled using unbalanced Monge OT.

\begin{example}[Figure \ref{fig:syn_data}a] Suppose the objective is to model the transformation from a source measure (Column 1) to the target measure (Column 2), which represent a population of interest at two distinct time points. The transport map $T$ models the transport/movement of points from the source to the target, while the scaling factor $\xi$ models the growth (replication) or shrinkage (death) of these points. Different models of transformation are optimal depending on the relative costs of mass transport and variation (Columns 3-6). 
\end{example}
\begin{example}[Figure \ref{fig:syn_data}b] Suppose the objective is to transport points from a source measure (1st panel, color) to a target measure (1st panel, grey) in the presence of class imbalances. A pure transport map would muddle together points from different classes, while an unbalanced transport map with a scaling factor is able to ameliorate the class imbalance (2nd panel). In this case, the scaling factor tells us explicitly how to downweigh or upweigh samples in the source distribution to balance the classes with the target distribution (3rd panel).
\end{example}

{\bf Relaxation.}
From an optimization standpoint, it is challenging to satisfy the constraint $T_{\#}(\xi\mu \times \lambda) = \nu$. We hence consider the following relaxation of (\ref{eq:ubOT_exact}) using a divergence penalty in place of the equality constraint:
\begin{equation}\label{eq:ubOT}
L_\psi(\mu, \nu) := \inf_{T, \xi} \int_\pazocal{X} \int_\pazocal{Z} c_1(x, T(x,z)) d\lambda(z) \xi(x) d\mu(x) + \int_\pazocal{X} c_2(\xi(x)) d\mu(x) + D_\psi(T_{\#}(\xi\mu \times \lambda) | \nu),
\end{equation}
using an appropriate choice of $\psi$ that satisfies the requirements of Lemma \ref{lem:div} in the Appendix%
\footnote{Note that $D_\psi(T_{\#}(\xi\mu \times \lambda) | \nu) = 0 \nRightarrow T_{\#}(\xi\mu \times \lambda) = \nu$ in general since the total mass of the transported measure is not constrained; for the relaxation to be valid, $\psi(s)$ should attain a unique minimum at $s=1$ (see Lemma \ref{lem:div}).}.%
This relaxation is the Monge-like version of the optimal-entropy transport problem (\ref{eq:transport}) by \citet{liero2018optimal}. Specifically, $(T, \xi)$ specifies a joint measure $\gamma \in \mathcal{M}_+(\pazocal{X} \times \pazocal{Y})$ given by
\[
\gamma(C) := \int_\pazocal{X} \int_\pazocal{Z} \mathbbm{1}_C(x,T(x,z)) d\lambda(z) \xi(x) d\mu(x), \quad \forall C \in \pazocal{B}(\pazocal{X}) \times \pazocal{B}(\pazocal{Y}),
\]
and by reformulating (\ref{eq:ubOT}) in terms of $\gamma$ instead of $(T, \xi)$, one obtains the objective function for optimal-entropy transport. The main difference between the formulations is their search space, since not all joint measures $\gamma \in \mathcal{M}_+(\pazocal{X} \times \pazocal{Y})$ can be specified by some choice of $(T, \xi)$. For example, if $T$ is a deterministic transport map, then $\gamma$ is necessarily restricted to the set of deterministic couplings. Even if $T$ is sufficiently random, it is generally not possible to specify all joint measures $\gamma \in \mathcal{M}_+(\pazocal{X} \times \pazocal{Y})$: in the asymmetric Monge formulation (\ref{eq:ubOT}), all the mass transported to $\pazocal{Y}$ must come from somewhere within the support of $\mu$, since the scaling factor $\xi$ allows mass to grow but not to materialize outside of its original support. Therefore equivalence can be established in general only when restricting the support of $\gamma$ to $supp(\mu) \times \pazocal{Y}$ as described in the following lemma, whose proof is given in the Appendix.

\begin{lemma} \label{lem:opt-entropy} Let $\pazocal{G}$ be the set of joint measures supported on $supp(\mu) \times \pazocal{Y}$%
, and define
\begin{equation*}\label{eq:OET}
\tilde W_{c, \psi_1, \psi_2}(\mu,\nu):= \inf_{\gamma \in \pazocal{G}} \int c ~d\gamma + D_{\psi_1}(\pi^\pazocal{X}_\#\gamma | \mu) + D_{\psi_2}(\pi^\pazocal{Y}_\#\gamma | \nu).
\end{equation*}
If $(\pazocal{Z}, \pazocal{B}(\pazocal{Z}), \lambda)$ is atomless and $c_2$ is an entropy function, then $L_\psi(\mu, \nu) = \tilde W_{c_1, c_2, \psi}(\mu, \nu)$.

\end{lemma}

Based on the relation between (\ref{eq:transport}) and (\ref{eq:ubOT}), several theoretical results for (\ref{eq:ubOT}) follow from the analysis of optimal entropy-transport by \citet{liero2018optimal}. Importantly, one can show the following theorem, namely that for an appropriate and sufficiently large choice of divergence penalty, solutions of the relaxed problem (\ref{eq:ubOT}) converge to solutions of the original problem (\ref{eq:ubOT_exact}). The proof is given in the Appendix.

\begin{theorem} \label{prop:relaxed}
Suppose $c_1, c_2, \psi$ satisfy the existence assumptions of Proposition \ref{prop:wellposed} in the Appendix, and let $(\pazocal{Z}, \pazocal{B}(\pazocal{Z}), \lambda)$ be an atomless probability space. Furthermore, let $\psi$ be uniquely minimized at $\psi(1)=0$. Then for a sequence $0 < \zeta^1< \cdots < \zeta^k < \cdots$ diverging to $\infty$ indexed by $k$, $\lim_{k \rightarrow \infty} L_{\zeta^k \psi}(\mu, \nu) = L(\mu,\nu)$. Additionally, let $\gamma^k$ be the joint measure specified by a minimizer of $L_{\zeta^k \psi}(\mu, \nu)$. If $L(\mu,\nu)<\infty$, then up to extraction of a subsequence, $\gamma^k$ converges weakly to $\gamma$, the joint measure specified by a minimizer of $L(\mu,\nu)$. 
\end{theorem}

{\bf Algorithm.} Using the relaxation of unbalanced Monge OT in (\ref{eq:ubOT}), we now show that the transport~map~and~scaling factor can be learned by stochastic gradient methods. While the divergence term cannot easily~be~minimized~using the definition in (\ref{eq:div}), we can write it as a penalty witnessed by an \emph{adversary} function $f: \pazocal{Y} \rightarrow (-\infty, \psi'_{\infty}]$ using the convex conjugate representation (see Lemma \ref{lem:vc2}):
\begin{equation}\label{eq:var-div}
D_\psi(T_{\#}(\xi\mu \times \lambda) | \nu) = \sup_f \int_\pazocal{X} \int_\pazocal{Z} f(T(x, z)) d\lambda(z) \xi(x) d\mu(x)  - \int_\pazocal{Y} \psi^*(f(y)) d\nu(y),
\end{equation}
where $\psi^*$ is the convex conjugate of $\psi$. The objective in (\ref{eq:ubOT}) can now be optimized using alternating stochastic gradient updates after parameterizing $T$, $\xi$, and $f$ with neural networks; see Algorithm \ref{alg:TAG} %
   \footnote{We assume that one has access to samples from $\mu, \nu$, and in the setting where $\mu, \nu$ are not normalized, then samples to the normalized measures $\tilde \mu, \tilde \nu$ as well as the normalization constants $c_\mu, c_\nu$. These are reasonable assumptions for practical applications: for example, in a biological assay, one might collect $c_\mu$ cells from time point 1 and $c_\nu$ cells from time point 2. In this case, the samples are the measurements of each cell and the normalization constants are $c_\mu, c_\nu$. }. %
The optimization procedure is similar to GAN training and can be interpreted as an adversarial game between $(T, \xi)$ and $f$:\\

\begin{itemize}
\itemsep0em 
\item $T$ takes a point $x \sim \mu$ and transports it from $\pazocal{X}$ to $\pazocal{Y}$ by generating $T(x, z)$ where $z \sim \lambda$.
\item $\xi$ determines the importance weight of each transported point.
\item Their shared objective is to minimize the divergence between transported samples and real samples from $\nu$ that is measured by the adversary $f$.
\item Additionally, cost functions $c_1$ and $c_2$ encourage $T, \xi$ to find the most cost-efficient strategy.
\end{itemize} 
 



\vspace{-0.2cm}

\begin{algorithm}[H]
\caption{Generative-Adversarial Framework for Unbalanced Monge OT}
\begin{algorithmic}
   \STATE {\bfseries Input:} Initial parameters $\theta$, $\phi$, $\omega$; step size $\eta$; normalized measures $\tilde \mu$, $\tilde \nu$, constants $c_\mu, c_\nu$.
   
   
   \STATE {\bfseries Output:} Updated parameters $\theta$, $\phi$, $\omega$.
   \WHILE{($\theta$, $\phi$, $\omega$) not converged}
   
   \STATE Sample $x_1, \cdots, x_n$ from $\tilde \mu$, $y_1, \cdots, y_n$ from $\tilde \nu$, $z_1, \cdots, z_n$ from $\lambda$;
   \vspace{-0.3cm}
   \begin{align}\label{eq:loss}
   \ell(\theta, \phi, \omega):= \frac{1}{n} \sum_{i=1}^n [ c_\mu c_1(x_i, T_\theta(x_i, z_i))\xi_\phi(x_i) &+ c_\mu c_2(\xi_\phi(x_i)) \\
   &+ c_\mu \xi_\phi(x_i) f_\omega(T_\theta(x_i, z_i)) - c_\nu \psi^*(f_\omega(y_i)). ] \nonumber
   \end{align}   
   \vspace{-0.5cm}
   \STATE Update $\omega$ by gradient descent on $-\ell(\theta, \phi, \omega)$.
   \STATE Update $\theta, \phi$ by gradient descent $\ell(\theta, \phi, \omega)$.

   \ENDWHILE
\end{algorithmic}
\label{alg:TAG}
\end{algorithm}

\vspace{-0.1cm}

Table \ref{tab:div} in the Appendix provides some examples of divergences with corresponding entropy functions and convex conjugates that can be plugged into (\ref{eq:var-div}). Further practical considerations for implementation and training are discussed in Appendix \ref{sec:prac}. 

{\bf Relation to other approaches.} The probabilistic Monge-like formulation (\ref{eq:ubOT}) is similar to the Kantorovich-like entropy-transport problem (\ref{eq:transport}) in theory, but they result in quite different numerical methods in practice. Algorithm \ref{alg:TAG} solves the non-convex formulation (\ref{eq:ubOT}) and learns a transport map $T$ and scaling factor $\xi$ parameterized by neural networks, enabling scalable optimization using stochastic gradient descent. The networks are immediately useful for many practical applications; for instance, it only requires a single forward pass to compute the transport and scaling of a point from the source domain to the target. Furthermore, the neural architectures of $T, \xi$ imbue their function classes with a particular structure, and when chosen appropriately, enable effective learning of these functions in high-dimensional settings. Due to the non-convexity of the optimization problem, however, Algorithm \ref{alg:TAG} is not guaranteed to find the global optimum. In contrast, the scaling algorithm of \citet{chizat2016scaling} based on (\ref{eq:transport}) solves a convex optimization problem and is proven to converge, but is currently only practical for discrete problems and has limited scalability. For completeness, in Section \ref{sec:met1} of the Appendix, we propose a \emph{new stochastic method} based on the same dual objective as \citet{chizat2016scaling} that can handle transport between continuous measures (Algorithm \ref{alg:1} in the Appendix). This method generalizes the approach of \citet{seguy2017large} for handling transport between continuous measures and overcomes the scalability limitations of \citet{chizat2016scaling}. However, the output is in the form of the dual solution, which is less interpretable for practical applications compared to the output of Algorithm~\ref{alg:TAG}. In particular, while one can compute a deterministic transport map known as a barycentric projection from the dual solution, it is unclear how best to obtain a scaling factor or a stochastic transport map that can generate samples outside of the target dataset. In the numerical experiments of Section \ref{sec:exp}, we show the advantage of directly learning a transport map and scaling factor using Algorithm \ref{alg:TAG}.

The problem of learning a scaling factor (or weighting factor) that ``balances" measures $\mu$ and $\nu$ also arises in causal inference. Generally, $\mu$ is the distribution of covariates from a control population and $\nu$ is the distribution from a treated population. The goal is to scale the importance of different members from the control population based on how likely they are to be present in the treated population, in order to eliminate selection biases in the inference of treatment effects. \cite{kallus2018deepmatch} proposed a generative-adversarial method for learning the scaling factor, but they do not consider transport.

\section{Numerical Experiments} \label{sec:exp}

In this section, we illustrate in practice how Algorithm \ref{alg:TAG} performs unbalanced OT, with applications geared towards population modeling.

{\bf MNIST-to-MNIST.} We first apply Algorithm \ref{alg:TAG} to perform unbalanced optimal transport between two modified MNIST datasets. The source dataset consists of regular MNIST digits with the class distribution shown in column 1 of Figure \ref{fig:mnist_data}a. The target dataset consists of either regular (for the experiment in Figure \ref{fig:mnist_data}b) or dimmed (for the experiment in Figure \ref{fig:mnist_data}c) MNIST digits with the class distribution shown in column 2 of Figure \ref{fig:mnist_data}a. The class imbalance between the source and target datasets imitates a scenerio in which certain classes (digits 0-3) become more popular and others (6-9) become less popular in the target population, while the change in brightness is meant to reflect population drift. We evaluated Algorithm \ref{alg:TAG} on the problem of transporting the source distribution to the target distribution, enforcing a high cost of transport (w.r.t. Euclidean distance). In both cases, we found that the scaling factor over each of the digit classes roughly reflects its ratio of imbalance between the source and target distributions (Figure \ref{fig:mnist_data}b-c). These experiments validate that the scaling factor learned by Algorithm \ref{alg:TAG} reflects the class imbalances and can be used to model growth or decline of different classes in a population. Figure \ref{fig:mnist_data}d is a schematic illustrating the reweighting that occurs during unbalanced OT.


\begin{figure}
\vspace{-0.2cm}
\subfigure[]{\includegraphics[scale=.55]{mnist_chart.png}} \quad
\subfigure[]{\includegraphics[scale=.55]{mnist_no_trans.png}}\quad
\subfigure[]{\includegraphics[scale=.55]{mnist_with_trans.png}}\quad\quad
\subfigure[]{\includegraphics[scale=.4]{mnist_schematic.png}}
\vspace{-0.2cm}
\caption{Learning weights on MNIST data using unbalanced OT.}
\vspace{-0.2cm}
\label{fig:mnist_data}
\end{figure}

{\bf MNIST-to-USPS.} Next, we apply unbalanced OT from the MNIST dataset to the USPS dataset. As before, these two datasets are meant to imitate a population sampled at two different time points, this time with a large degree of evolution. We use Algorithm \ref{alg:TAG} to model the evolution of the MNIST distribution to the USPS distribution, taking as transport cost the Euclidean distance between the original and transported images. A summary of the unbalanced transport is visualized in Figure \ref{fig:num_data}a. Each arrow originates from a real MNIST image and points towards the predicted appearance of this image in the USPS dataset. The size of the image reflects the scaling factor of the original MNIST image, i.e.~whether it is relatively increasing or decreasing in prominence in the USPS dataset compared to the MNIST dataset according to the unbalanced OT model. Even though the Euclidean distance is not an ideal measure of correspondence between MNIST and USPS digits, many MNIST digits were able to preserve their likeness during the transport (Figure \ref{fig:num_data}b). We analyzed which MNIST digits were considered as increasing or decreasing in prominence by the model. The MNIST digits with higher scaling factors were generally brighter (Figure \ref{fig:num_data}c) and covered a larger area of pixels (Figure \ref{fig:num_data}d) compared to the MNIST digits with lower scaling factors. These results are consistent with the observation that the target USPS digits are generally brighter and contain more pixels.

\begin{figure}[h]
\vspace{-0.2cm}
\centering
\subfigure[]{\includegraphics[scale=.4]{mnist_transport.png}} \quad
\subfigure[]{\includegraphics[scale=.4]{mnist_to_usps_examples.png}}\quad
\subfigure[]{\includegraphics[scale=.3]{mnist_to_usps_brightness.png}}\quad
\subfigure[]{\includegraphics[scale=.3]{mnist_to_usps_area.png}}
\vspace{-0.2cm}
\caption{Unbalanced Optimal Transport from MNIST to USPS digits.}
\vspace{-0.2cm}
\label{fig:num_data}
\end{figure}



{\bf CelebA-Young-to-CelebA-Aged.} We applied Algorithm \ref{alg:TAG} on the CelebA dataset to perform unbalanced OT from the population of young faces to the population of aged faces. This synthetic problem imitates a real application of interest, which is modeling the transformation of a population based on samples taken from two timepoints. Since the Euclidean distance between two faces is a poor measure of semantic similarity, we first train a variational autoencoder (VAE) \citep{kingma2013auto} on the CelebA dataset and encode all samples into the latent space. We then apply Algorithm \ref{alg:TAG} to perform unbalanced OT from the encoded young to the encoded aged faces, taking the transport cost to be the Euclidean distance in the latent space.
A summary of the unbalanced transport is visualized in Figure \ref{fig:celebA_data}a. Each arrow originates from a real face from the young population and points towards the predicted appearance of this face in the aged population. Generally, the transported faces retain the most salient features of the original faces  (Figure \ref{fig:celebA_data}b), although there are exceptions (e.g. gender swaps) which reflects that some features are not prominent components of the VAE encodings. Interestingly, the young faces with higher scaling factors were significantly enriched for males compared to young faces with lower scaling factors; 9.6\% (9,913/103,287) of young female faces had a high scaling factor as compared to 18.5\% (8,029/53,447) for young male faces  (Figure \ref{fig:celebA_data}c, top, $p = 0$). In other words, our model predicts growth in the prominence of male faces compared to female faces as the CelebA population evolves from young to aged. After observing this phenomenon, we confirmed based on checking the ground truth labels that there was indeed a strong gender imbalance between the young and aged populations: while the young population is predominantly female, the aged population is predominantly male (Figure \ref{fig:celebA_data}c, bottom). 

\begin{figure}[b]
\vspace{-0.2cm}
\subfigure[]{\includegraphics[scale=.30]{celebA_transport.png}} \quad
\subfigure[]{\includegraphics[scale=.20]{young_to_aged.png}} \quad
\subfigure[]{\includegraphics[scale=.15]{tables.png}}
\vspace{-0.2cm}
\caption{Unbalanced Optimal Transport from Young to Aged CelebA Faces.}
\label{fig:celebA_data}
\vspace{-0.2cm}
\end{figure}



{\bf Zebrafish embroygenesis.} A problem of great interest in biology is lineage tracing of cells between different developmental stages or during disease progression. This is a natural application of transport in which the source and target distributions are unbalanced: some cells in the earlier stage are more poised to develop into cells seen in the later stage. To showcase the relevance of learning the scaling factor, we apply Algorithm \ref{alg:TAG} to recent single-cell gene expression data from two stages of zebrafish embryogenesis~\citep{farrelleaar3131}. 
The source population is from a late stage of blastulation and the target population from an early stage of gastrulation (Figure \ref{fig:zebrafish}a). The results of the transport are plotted in Figure \ref{fig:zebrafish}b-c after dimensionality reduction by PCA and T-SNE \citep{maaten2008visualizing}. To assess the scaling factor, we extracted the cells from the blastula stage with higher scaling factors (i.e. over 90th percentile) and compared them to the remainder of the cells using differential gene expression analysis, producing a ranked list of upregulated genes. Using the GOrilla tool \citep{eden2009gorilla}, we found that the cells with higher scaling factors were significantly enriched for genes associated with differentiation and development of the mesoderm (Figure \ref{fig:zebrafish}d). This experiment shows that analysis of the scaling factor can be applied towards interesting and meaningful biological discovery.

\begin{figure}[h]
	\centering
	\includegraphics[scale=0.4]{zebrafish_figure.png}
	\caption{{\bf Unbalanced OT on Zebrafish Single-Cell Gene Expression Data.} (a) Illustration of the blastula and gastrula stages of zebrafish embryogenesis. (b) T-SNE plot \citep{maaten2008visualizing} of the unbalanced OT results for a subset of datapoints. The color of the transported points indicates the relative magnitude of the scaling factor (black = high, white = low). (c) Same plot as (b), where we have colored the source points instead of the transported points. (d) GOrilla output of significantly enriched processes \citep{eden2009gorilla} based on ranked list of enriched genes in cells with high scaling factors (black points from (c)) from a differential gene expression analysis. Processes in the graph are organized from more general (upstream) to more specific (downstream). }
	\label{fig:zebrafish}
\end{figure}





\pagebreak
\subsection*{Acknowledgements}
Karren Yang was supported by an NSF Graduate Fellowship and ONR (N00014-17-1-2147). Caroline Uhler was partially supported by NSF (DMS-1651995), ONR (N00014-17-1-2147 and N00014-18-1-2765), IBM, and J-WAFS. 

\bibliographystyle{iclr2019_conference}

