SlowGuess/NeurIPS_2023_OCR huggingface dataset:
where  $\iota$  is a convex indicator function. Note that we assume  $\Psi_1^*$  and  $\Psi_2^*$  are convex, non-decreasing, and differentiable. Since  $c\geq 0$ , by letting  $u\equiv -1$  and  $v\equiv -1$ , we can easily see that all three terms in Eq 25 are finite values. Thus, we can apply Fenchel-Rockafellar's theorem, which implies that the strong duality holds. Finally, combining the constraint  $u(x) + v(y)\leq c(x,y)$  with tightness condition, i.e.  $c(x,y) = u(x) + v(y)\pi$ -almost everywhere, we obtain the following semi-dual form:
$$
\sup  _ {v \in \mathcal {C}} \left[ \int_ {\mathcal {X}} - \Psi_ {1} ^ {*} \left(- \inf  _ {y \in \mathcal {Y}} [ c (x, y) - v (y) ]\right) \mu (x) + \int_ {\mathcal {Y}} - \Psi_ {2} ^ {*} (- v (y)) \nu (y) \right]. \tag {26}
$$
Now, we state precisely and prove Theorem 3.3.
Lemma A.2. For any  $(u,v)\in \mathcal{C}(\mathcal{X})\times \mathcal{C}(\mathcal{Y})$  , let
$$
I (u, v) := = \int_ {\mathcal {X}} - \Psi_ {1} ^ {*} (- u (x)) d \mu (x) + \int_ {\mathcal {Y}} - \Psi_ {2} ^ {*} (- v (y)) d \nu (y). \tag {27}
$$
Then, under the assumptions in Appendix A, the functional derivative of  $I(u,v)$  is
$$
\int_ {\mathcal {X}} \delta u (x) \Psi_ {1} ^ {* \prime} (- u (x)) d \mu (x) + \int_ {\mathcal {Y}} \delta v (y) \Psi_ {2} ^ {* \prime} (- v (y)) d \nu (y), \tag {28}
$$
where where  $(\delta u,\delta v)\in \mathcal{C}(\mathcal{X})\times \mathcal{C}(\mathcal{Y})$  denote the variations of  $(u,v)$ , respectively.
Proof. For any  $(f,g)\in \mathcal{C}(\mathcal{X})\times \mathcal{C}(\mathcal{Y})$  , let
$$
I (f, g) = \int_ {\mathcal {X}} - \Psi_ {1} ^ {*} (- f (x)) d \mu (x) + \int_ {\mathcal {Y}} - \Psi_ {2} ^ {*} (- g (y)) d \nu (y).
$$
Then, for arbitrarily given potentials  $(u,v)\in \mathcal{C}(\mathcal{X})\times \mathcal{C}(\mathcal{Y})$ , and perturbation  $(\eta ,\zeta)\in \mathcal{C}(\mathcal{X})\times \mathcal{C}(\mathcal{Y})$  which satisfies  $u(x) + \eta (x) + v(y) + \zeta (y)\leq c(x,y)$
$$
\mathcal {L} (\epsilon) := \frac {d}{d \epsilon} I (u + \epsilon \eta , v + \epsilon \zeta) = \int_ {\mathcal {X}} \eta \Psi_ {1} ^ {* \prime} (- ((u + \epsilon \eta) (x))) d \mu (x) + \int_ {\mathcal {Y}} \zeta \Psi_ {2} ^ {* \prime} (- ((v + \epsilon \zeta) (y))) d \nu (y).
$$
Note that  $\mathcal{L}(\epsilon)$  is a continuous function. Thus, the functional derivative of  $I(u,v)$ , i.e.  $\mathcal{L}(0)$ , is given by calculus of variation (Chapter 8, Evans [17])
$$
\int_ {\mathcal {X}} \delta u (x) \Psi_ {1} ^ {* \prime} (- u (x)) d \mu (x) + \int_ {\mathcal {Y}} \delta v (y) \Psi_ {2} ^ {* \prime} (- v (y)) d \nu (y), \tag {29}
$$
where  $(\delta u,\delta v)\in \mathcal{C}(\mathcal{X})\times \mathcal{C}(\mathcal{Y})$  denote the variations of  $(u,v)$ , respectively.
Theorem A.3. Suppose that  $\mu$  and  $\nu$  are probability densities defined on  $\mathcal{X}$  and  $\mathcal{Y}$ . Given the assumptions in Appendix A, suppose that  $\mu, \nu$  are absolutely continuous with respect to Lebesgue measure and  $\Psi^*$  is continuously differentiable. Assuming that the optimal potential  $v^{\star} = \inf_{v\in \mathcal{C}}J(v)$  exists,  $v^{\star}$  is a solution of the following objective
$$
\tilde {J} (v) = \int_ {\mathcal {X}} - v ^ {c} (x) d \tilde {\mu} (x) + \int_ {\mathcal {Y}} - v (y) d \tilde {\nu} (y), \tag {30}
$$
$$
T ^ {\star} (x) \in \operatorname {a r g i n f} _ {y \in \mathcal {Y}} [ c (x, y) - v ^ {\star} (y) ], \tag {31}
$$
$\mu$ -a.s..In particular,  $D_{\Psi_1}(\tilde{\mu}|\mu) + D_{\Psi_2}(\tilde{\nu}|\nu) \leq \tau \mathcal{W}_2^2(\mu,\nu)$  where  $\mathcal{W}_2(\mu,\nu)$  is a Wasserstein-2 distance between  $\mu$  and  $\nu$ .
Proof. By Lemma A, derivative of Eq 27 is
$$
\int_ {\mathcal {X}} \delta u (x) \Psi_ {1} ^ {* \prime} (- u (x)) d \mu (x) + \int_ {\mathcal {Y}} \delta v (y) \Psi_ {2} ^ {* \prime} (- v (y)) d \nu (y), \tag {32}
$$
Note that such linearization must satisfy the inequality constraint  $u(x) + \delta u(x) + v(y) + \delta v(y) \leq c(x, y)$  to give admissible variations. Then, the linearized dual form can be re-written as follows:
$$
\sup  _ {(u, v) \in \mathcal {C} (\mathcal {X}) \times \mathcal {C} (\mathcal {Y})} \int_ {X} \delta u (x) d \tilde {\mu} (x) + \int_ {\mathcal {Y}} \delta v (y) d \tilde {\nu} (y), \tag {33}
$$
where  $\delta u(x) + \delta v(y) \leq c(x,y) - u(x) - v(y)$ ,  $\tilde{\mu}(x) = \Psi_1^{*'}(-u(x))\mu(x)$ , and  $\tilde{\nu}(y) = \Psi_2^{*'}(-v(y))\nu(y)$ . Thus, whenever  $(u^{\star}, v^{\star})$  is optimal, the linearized dual problem is non-positive due to the inequality constraint and tightness (i.e.  $c(x,y) - u^{\star}(x) - v^{\star}(y) = 0$  -almost everywhere), which implies that Eq 33 achieves its optimum at  $(u^{\star}, v^{\star})$ . Note that Eq 33 is a linearized dual problem of the following dual problem:
$$
\sup  _ {u (x) + v (y) \leq c (x, y)} \left[ \int_ {\mathcal {X}} u (x) d \tilde {\mu} (x) + \int_ {\mathcal {Y}} v (y) d \tilde {\nu} (y) \right]. \tag {34}
$$
Hence, the UOT problem can be reduced into standard OT formulation with marginals  $\tilde{\mu}$  and  $\tilde{\nu}$ . Moreover, the continuity of  $\Psi^{*'}$  gives that  $\tilde{\mu},\tilde{\nu}$  are absolutely continuous with respect to Lebesgue measure and have a finite second moment (Note that  $\mathcal{X},\mathcal{Y}$  are compact. Hence,  $\Psi^{*'}$  is bounded on  $\mathcal{X},\mathcal{Y}$ , and  $\mathcal{X},\mathcal{Y}$  are bounded). By Theorem 2.12 in Villani [76], there exists a unique measurable OT map  $T^{\star}$  which solves Monge problem between  $\tilde{\mu}$  and  $\tilde{\nu}$ . Furthermore, as shown in Remark 5.13 in Villani et al. [77],  $T^{\star}(x)\in \operatorname{arginf}_{y\in \mathcal{Y}}[c(x,y) - v(y)]$ .
Finally, if we constrain  $\pi$  into  $\Pi (\mu ,\nu)$ , then the solution of Eq 4 is  $C_{ub} = \tau \mathcal{W}_2^2 (\mu ,\nu)$ . Thus, for the optimal  $\pi \in \mathcal{M}_{+}$  where  $\pi_0 = \tilde{\mu}$  and  $\pi_1 = \tilde{\nu}$ , it is trivial that
$$
\tau \mathcal {W} _ {2} ^ {2} (\mu , \nu) \geq \tau \mathcal {W} _ {2} ^ {2} (\tilde {\mu}, \tilde {\nu}) + D _ {\Psi_ {1}} (\tilde {\mu} | \mu) + D _ {\Psi_ {2}} (\tilde {\nu} | \nu), \tag {35}
$$
which proves the last statement of the theorem.
Theorem A.4 ([21]).  $J$  is convex. Suppose  $v$  is a  $\lambda$ -strongly convex,  $v^c$  and  $v$  are uniformly bounded on the support of  $\mu$  and  $\nu$ , respectively, and  $\Psi_1^*$ ,  $\Psi_2^*$  are strongly convex on every compact set. Then, it holds
$$
J (v) - J \left(v ^ {\star}\right) \geq \frac {1}{2 \lambda} \mathbb {E} _ {\tilde {\mu}} \left[ \| \nabla \left(v ^ {c} - v ^ {\star c}\right) \| _ {2} ^ {2} \right] + C _ {1} \mathbb {E} _ {\mu} \left[ \left(v ^ {c} - v ^ {\star c}\right) ^ {2} \right] + C _ {2} \mathbb {E} _ {\nu} \left[ \left(v - v ^ {\star}\right) ^ {2} \right], \tag {36}
$$
for some positive constant  $C_1$  and  $C_2$ . Furthermore,  $\tilde{J}(v) - \tilde{J}(v^{\star}) \geq \frac{1}{2\lambda} \mathbb{E}_{\tilde{\mu}} \left[ \| \nabla (v^c - v^{\star c}) \|_2^2 \right]$ .
Proof. See [21].
Unless otherwise stated, the source distribution  $\mu$  is a standard Gaussian distribution  $\mathcal{N}(0,\mathbf{I})$  with the same dimension as the target distribution  $\nu$ .
Toy Data We generated 4000 samples for each Toy dataset: Toy dataset for Outlier Robustness (Sec 5.1) and Toy dataset for Target Distribution Matching (Sec 5.2). In Fig 2 of Sec 5.1, the target dataset consists of  $99\%$  of samples from  $\mathcal{N}(1,0.5^2)$  and  $1\%$  of samples from  $\mathcal{N}(-1,0.5^2)$ . In Tab 1 and Fig 4 of Sec 5.2, the source data is composed of  $50\%$  of samples from  $\mathcal{N}(-1,0.5^2)$  and  $50\%$  of samples from  $\mathcal{N}(1,0.5)$ . The target data is sampled from the mixture of two Gaussians:  $\mathcal{N}(-1,0.5)$  with a probability of  $1/3$  and  $\mathcal{N}(2,0.5)$  with a probability of  $2/3$ .
CIFAR-10 We utilized all 50,000 samples. For each image  $x$ , we applied random horizontal flip augmentation of probability 0.5 and transformed it by  $2x - 1$  to scale the values within the  $[-1,1]$  range. In the outlier experiment (Fig 3 of Sec 5.1), we added additional 500 MNIST samples to the clean CIFAR-10 dataset. When adding additional MNIST samples, each sample was resized to  $3 \times 32 \times 32$  by duplicating its channel to create a 3-channeled image. Then, the samples were transformed by  $2x - 1$ . For the image generation task in Sec 5.2, we followed the source distribution embedding of Rout et al. [61] for the small model, i.e., the Gaussian distribution  $\mu$  with the size of  $3 \times 8 \times 8$  is bicubically upsampled to  $3 \times 32 \times 32$ .
CelebA-HQ We used all 120,000 samples. For each image  $x$ , we resized it to  $256 \times 256$  and applied random horizontal flip augmentation with a probability of 0.5. Then, we linearly transformed each image by  $2x - 1$  to scale the values within  $[-1, 1]$ .
Toy data For all the Toy dataset experiments, we used the same generator and discriminator architectures. The dimension of the auxiliary variable  $z$  is set to one. For a generator, we passed  $z$  through two fully connected (FC) layers with a hidden dimension of 128, resulting in 128-dimensional embedding. We also embedded data  $x$  into the 128-dimensional vector by passing it through three-layered ResidualBlock [68]. Then, we summed up the two vectors and fed it to the final output module. The output module consists of two FC layers. For the discriminator, we used three layers of ResidualBlock and two FC layers (for the output module). The hidden dimension is 128. Note that the SiLU activation function is used. We used a batch size of 256, a learning rate of  $10^{-4}$ , and 2000 epochs. Like WGAN [6], we used the number of iterations of the discriminator per generator iteration as five. For OTM experiments, since OTM does not converge without regularization, we set  $R_{1}$  regularization to  $\lambda = 0.01$ .
CIFAR-10 For the small model setting of UOTM, we employed the architecture of Balaji et al. [8]. Note that this is the same model architecture as in Rout et al. [61]. We set a batch size of 128, 200 epochs, a learning rate of  $2 \times 10^{-4}$  and  $10^{-4}$  for the generator and discriminator, respectively. Adam optimizer with  $\beta_{1} = 0.5$ ,  $\beta_{2} = 0.9$  is employed. Moreover, we use  $R_{1}$  regularization of  $\lambda = 0.2$ . For the large model setting of UOTM, we followed the implementation of Choi et al. [13] unless stated. We trained for 600 epochs with  $\lambda = 0.2$  for every iteration. Moreover, in OTM (Large) experiment, we trained the model for 1000 epochs. The other hyperparameters are the same as UOTM. Tt For the ablation studies in Sec 5.4, we reported the best FID score until 600 epoch for UOTM and 1000 epoch for OTM because the optimal epoch differs for different hyperparameters. Additionally, for the ablation study on  $\tau$  in the cost function (Fig 7), we reported the best FID score among three  $R_{1}$  regularizer parameters:  $\lambda = 0.2, 0.02, 1$ . To ensure the reliability of our experiments, we conducted three experiments with different random seeds and presented the mean and standard deviation in Tab 2 (UOTM (Large)).
CelebA-HQ  $(256\times 256)$  We set  $\tau = 10^{-5}$  and used  $R_{1}$  regularizer with  $\lambda = 5$  for every iteration. The model is trained for 450 epochs with an exponential moving average of 0.999. The other network architecture and optimizer settings are the same as Xiao et al. [80] and Choi et al. [13].
Evaluation Metric We used the implementation of Wang et al. [78] for the KL divergence in Tab 1 of Sec 5.2. We set  $k = 2$  (See Wang et al. [78]). For the evaluation of image datasets, we used 50,000 generated samples to measure IS and FID scores.
Ablation Study on Csiszar Divergence To ensure self-containedness, we include the definition of Softplus function employed in Table 4.
$$
\operatorname {S o f t p l u s} (x) = \log (1 + \exp (x)) \tag {37}
$$
Figure 8: Various types of  $\Psi^{*}$ .
C Additional Discussion on UOTM
C.1 Stability of UOTM
In this section, we provide an intuitive explanation for the stability and convergence of our model. Note that our model employs a non-decreasing, differentiable, and convex  $\Psi^{*}$  (Sec 3) as depicted in Fig 8. In contrast, OTM corresponds to the specific case of  $\Psi^{*}(x) = x$ . Now, recall our objective function for the potential  $v_{\phi}$  and the generator  $T_{\theta}$ :
$$
\mathcal {L} _ {v} = \frac {1}{| X |} \sum_ {x \in X} \Psi_ {1} ^ {*} (- c (x, T _ {\theta} (x, z)) + v _ {\phi} (T _ {\theta} (x, z))) + \frac {1}{| Y |} \sum_ {y \in Y} \Psi_ {2} ^ {*} (- v _ {\phi} (y)), \tag {38}
$$
$$
\mathcal {L} _ {T} = \frac {1}{| X |} \sum_ {x \in X} \left(c \left(x, T _ {\theta} (x, z)\right) - v _ {\phi} \left(T _ {\theta} (x, z)\right)\right). \tag {39}
$$
Then, the gradient descent step of potential  $v_{\phi}$  for a single real data  $y$  with a learning rate of  $\gamma$  can be expressed as follows:
$$
\phi - \gamma \nabla_ {\phi} \mathcal {L} _ {v _ {\phi}} = \phi - \gamma \nabla_ {\phi} \Psi_ {2} ^ {*} (- v _ {\phi} (y)) = \phi + \gamma \Psi_ {2} ^ {* \prime} (- v _ {\phi} (y)) \nabla_ {\phi} v _ {\phi} (y). \tag {40}
$$
Suppose that our potential network  $v_{\phi}$  assigns a low potential for  $y$ . Since the objective of the potential network is to assign high values to the real data, this indicates that  $v_{\phi}$  fails to allocate the appropriate potential to the real data point  $y$ . Then, because  $-v_{\phi}(y)$  is large,  $\Psi^{*'}(-v_{\phi}(y))$  in Eq 40 is large, as shown in Fig 8. In other words, the potential network  $v_{\phi}$  takes a stronger gradient step on the real data  $y$  where  $v_{\phi}$  fails. Note that this does not happen for OTM because  $\Psi^{*'}(x) = 1$  for all  $x$ . In this respect, UOTM enables adaptive updates for each sample, with weaker updates on well-performing data and stronger updates on poor-performing data. We believe that UOTM achieves significant performance improvements over OTM by this property. Furthermore, UOTM attains higher stability because the smaller updates on well-behaved data prevent the blow-up of the model output.
C.2 Qualitative Comparison on Varying  $\tau$
In the ablation study on  $\tau$ , UOTM exhibited a similar degradation in FID score  $(\geq 20)$  at  $\tau (\times 10^{-3}) = 0.1, 5$ . However, the generated samples are significantly different. Figure 9 illustrates the generated samples from trained UOTM for the cases when  $\tau$  is too small ( $\tau = 0.1$ ) and when  $\tau$  is too large ( $\tau = 5$ ). When  $\tau$  is too small (Fig 9 (left)), the generated images show repetitive patterns. This lack of diversity suggests mode collapse. We consider that the cost function  $c$  provides some regularization effect preventing mode collapse. For the source sample (noise)  $x$ , the cost function is defined as  $c(x,T(x)) = \tau \| x - T(x)\|_2^2$ . Hence, minimizing this cost function encourages the generated images  $T(x)$  to disperse by aligning each image to the corresponding noise  $x$ . Therefore, increasing  $\tau$  from 0.1 led to better FID results in Fig 7.
On the other hand, when  $\tau$  is too large Fig 9 (right), the generative samples tend to be noisy. Because the cost function is  $c(x,T(x)) = \tau \| x - T(x)\| _2^2$ , too large  $\tau$  induces the generated image  $T(x)$  to be similar to the noise  $x$ . This resulted in the degradation of the FID score for  $\tau = 5$  in Fig 7. As discussed in Sec 5.4, this phenomenon can also be understood through Thm 3.3.
Figure 9: Generated samples from UOTM trained on CIFAR-10 with Left:  $\tau = 0.1(\times 10^{-3})$  and Right:  $\tau = 5(\times 10^{-3})$ .
Figure 10: Generated samples from UOTM trained on CIFAR10 for  $(\Psi_{1},\Psi_{2}) = (KL,KL)$
Figure 11: Generated samples from UOTM trained on CIFAR10 for  $(\Psi_{1},\Psi_{2}) = (\chi^{2},\chi^{2})$
Figure 12: Generated samples from UOTM trained on CIFAR10 for  $(\Psi_{1},\Psi_{2}) = (KL,\chi^{2})$
Figure 13: Generated samples from UOTM trained on CIFAR10 for  $(\Psi_{1},\Psi_{2}) = (\chi^{2},KL)$
Figure 14: Generated samples from UOTM trained on CIFAR10 for  $(\Psi_{1},\Psi_{2}) =$  (Softplus, Softplus).
Figure 15: Generated samples from UOTM trained on CelebA-HQ  $(256\times 256)$