Wasserstein Geodesic Generator for Conditional Distributions

Most conditional generative models are extended from image generation methods. We first review three eminent image generative models: VAEs, GANs, and AAEs. To synthesize realistic data, all three methods aim to learn a generator that transforms latent factors, which follow a user-specified prior distribution. VAEs consist of encoder and decoder networks. They model the joint likelihood of latent factors and their transformations, feed-forwarded by decoder networks, and seek the maximum likelihood estimator for the marginal distribution of observations. Due to the intractability of the likelihood with nonlinear decoder networks, VAEs employ variational inference (Bishop, 2006) to maximize the evidence lower bound, using encoder networks to approximate the distribution of latent variables given observations. In contrast, GANs are composed of discriminator and generator networks. The generator networks in GANs serve the same function as decoders in VAEs—transforming latent factors to generate data. However, GANs introduce discriminator networks to form an adversarial loss for generator training. AAEs, similarly, employ discriminator and generator networks but also encoders. Unlike GANs, the discriminator of AAEs aims to classify encoded results and latent variables drawn from the prior distribution. Training AAEs can be interpreted as minimizing the $1$-Wasserstein distance between distributions of real data and generation results, a special case of Wasserstein autoencoders (WAEs) (Tolstikhin et al., 2018).

Statistical distances employed in forming training objectives in generative models are pivotal components for the quality of generation results. Both VAEs and GANs utilize the $f$-divergence (Csiszár, 1964) to learn distributions of images. As an alternative statistical distance, Arjovsky et al. (2017) showed that the Wasserstein distance has advantages over the $f$-divergence when the supports of data distributions are on a low-dimensional manifold as in image data. The Wasserstein distance yields differentiable losses, while $f$-divergences do not define losses well or yield non-differentiable losses, possibly due to these advantages, Wasserstein distance-based approaches including AAEs, WAEs, and Wasserstein GANs (Arjovsky et al., 2017) often have outperformed $f$-divergence-based approaches.

The main extension from image generative models to conditional models is to concatenate domain labels into the latent variable. In cVAE, cGAN, and cAAE, domain labels are incorporated into encoder networks of VAE, discriminator and generators of GAN, and decoder of AAE, respectively. To enhance the visual quality and diversity of generation results, Kameoka et al. (2018), Odena et al. (2017), and Zhao et al. (2018) introduce auxiliary classifier to match the observed and predicted domain labels for conditional generation results. Conditional generative models can be used to generate data for unobserved domains. This is accomplished by inserting unobserved domain label values into trained models, a technique demonstrated in zero-shot learning approaches (Xian et al., 2018; Chao et al., 2016), aimed at boosting classification performance for unseen classes. However, this approach faces limitations. Firstly, it assumes that latent variables for unobserved domains follow similar patterns to those for observed domains. Secondly, the generated data distribution for unobserved domains has not been justified with a metric space w.r.t. distributions. Intuitively, the unobserved intermediate distribution serves as the centroid of observed distributions, but this property has not been discussed. In contrast, our approach constructs geodesics in the Wasserstein space for unobserved domains, and the proposed distribution is the Wasserstein barycenter when distributions of representations are identical across observed domains. This result is achieved without assuming the homogeneity of representations across both observed and unobserved domains, distinguishing our approach from existing works.

We employ conditional generative models for learning data distributions from observed domain labels, observed vertices in the Wasserstein space. In Section 3, we present theoretical results to provide a tractable objective for minimizing the formulated Wasserstein distances between conditional distributions.

Conditional generative models find a transformation from latent variables and domain labels to data, effectively learning conditional distributions. In contrast, another line of work, which we refer to as Data-to-Data Translation (Kim et al., 2017; Choi et al., 2018; Zhu et al., 2017), operates under a different paradigm. In Data-to-Data Translation, the source and target domains are predefined, and the goal is to find a transport map from the source data to the target data. Typical examples include unpaired translations, such as converting daytime scenes to nighttime scenes, where corresponding pairs of daytime and nighttime scenes for the same location are not required during training. Prominent methods within this field include multi-modal translations across different data domains (Xu et al., 2018; Isola et al., 2017), such as Text-to-Image Translation, as demonstrated by DALL-E (Ramesh et al., 2021), and multi-domain Image-to-Image Translation, as exemplified by StarGAN (Choi et al., 2018).

CycleGAN (Zhu et al., 2017) is a pioneering work in unpaired Image-to-Image Translation. It minimizes the adversarial loss with target data and translated source data, with the cycle consistency loss encouraging the inverse relation between translation maps between two domains. In this case, the transformation is encouraged to match distributions of the target and converted source data, but it does not learn conditional distributions given domain labels since the real data from the source domains are mandatory to construct target data. Liu et al. (2017) introduce a latent variable model to learn conditional distributions. However, the properties of the conditional distributions given unobserved intermediate domains have not been discussed in the Data-to-Data Translation literature.

Our proposed method leverages Data-to-Data Translation techniques to find optimal transport maps between observed conditional distributions, thereby defining edges between observed vertices in the Wasserstein space. In Section 4, we propose to generate intermediate data from edges, Wasserstein geodesics, and extend it to approximate the centroid, which is the Wasserstein barycenter of observed distributions.

The Wasserstein space, a metric space of distributions endowed with the Wasserstein distance, has found widespread applications in various fields of generative models. We review related works in image processing, domain adaptation, and data augmentation. Other important applications of Wasserstein space include density matching (Cisneros-Velarde and Bullo, 2020), distribution alignment (Zhou et al., 2022), online learning (Korotin et al., 2021), and Bayesian inference (Srivastava et al., 2018).

In image processing, most applications focus on transporting point clouds (Cuturi, 2013) or specific features—such as texture (Rabin et al., 2012), colors (Rabin et al., 2014), and shapes (Solomon et al., 2015)—from a source image a target image. While these approaches have shown remarkable results, they typically require training models or solving an optimization problem for each pair of source and target images. Taking a different route, Mroueh (2020) introduces universal style transfer that employs autoencoders. The method exploits the Wasserstein geodesic of Gaussian measures, using features extracted by encoder networks. However, in Mroueh’s framework, distributions of generated samples for unobserved intermediate domains cannot be characterized by Wasserstein spaces. On a related note, Korotin et al. (2019) put forth a generative model to solve the dual form of the $2$-Wasserstein distance between two distributions of images, which they apply to image style transfer tasks.

In domain adaptation, Xie et al. (2019) proposes latent variable models using single representations to generate multiple images from each domain while minimizing the transportation costs between them. However, this method requires modality-specific generators and can not generate intermediate distributions. The Wasserstein Barycenter Transport (WBT) (Montesuma and Mboula, 2021) is a closely related work to ours. The WBT targets the Wasserstein barycenter of multiple observed source distributions to generate unobserved intermediate domains, but it requires pairs of observations from all the source domains and solves optimization problems for every generation.

In data augmentation, several recent works utilize the Wasserstein space. Bespalov et al. (2022) propose to augment landmark coordinates of facial images with the Wasserstein barycenter, but their method requires computing Wasserstein distances between all pairs of images to oversample landmark data. Zhu et al. (2023) augment data from the Wasserstein barycenter of distributions of images to learn robust classifiers, but this method assumes the Gaussianity of conditional distributions. The work by Fan and Alvarez-Melis (2023) is closely related to this work. Their approach synthesizes data for unobserved domains by applying linear combinations of optimal transport maps between datasets, essentially generating data from the generalized Wasserstein geodesic of observed data distributions. Despite the merits of generalized geodesics, such as convexity and impressive performance in transfer learning tasks, the method employs the optimal transport dataset distance (Alvarez-Melis and Fusi, 2020), dependent on classification labels from each domain. Additionally, the generalized geodesic differs from the Wasserstein barycenter, and the method uses an alternate transportation cost, the $(2,\nu)$-transport metric (Craig, 2016) in optimization.

Existing methods typically rely on strong assumptions, such as the data following Gaussian distributions and the need to solve optimization problems each time data is generated. Other assumptions include the existence of the Wasserstein distance, optimal transport map, Wasserstein geodesics, and Wasserstein barycenters with Euclidean distances on the data space, implying that the distribution of high-dimensional data is continuous. In contrast, our work fills this gap by generating and justifying intermediate, unobserved distributions without the aforementioned assumptions.

We provide basic notations as follows. Random variables, their realizations, and their supports are denoted by capital, small, and calligraphy capital letters, respectively. The real data, generated samples, and domain labels are denoted by $X$, $\tilde{X}$, and $C$, respectively. We denote the set of distributions defined on a given support $\mathcal{X}$ by $\mathcal{P}(\mathcal{X})$ and the conditional distribution of $X$ given $Y$ by $\mathbb{P}_{X|Y}$.

For any metric $d$ on $\mathcal{X}$ and probability measures $\mathbb{P}_{X}$ and $\mathbb{P}_{Y}$ in $\mathcal{P}(\mathcal{X})$, the $p$-Wasserstein distance between $\mathbb{P}_{X}$ and $\mathbb{P}_{Y}$ w.r.t. $d$ is denoted by $W_{p}(\mathbb{P}_{X},\mathbb{P}_{Y};d):=\big{(}\underset{\pi\in\Pi(\mathbb{P}_{X},\mathbb{P}_{Y})}{\inf}\int d^{p}(x,y)d\pi(x,y)\big{)}^{1/p}$ where $p\in[1,\infty)$ and $\Pi(\mathbb{P}_{X},\mathbb{P}_{Y})$ is the set of all couplings of $\mathbb{P}_{X}$ and $\mathbb{P}_{Y}$. For brevity, we omit $d$, the metric on $\mathcal{X}$, in the Wasserstein distance if there is no confusion. We assume compactness and convexity of $\mathcal{X}$ to ensure that the Wasserstein space $(\mathcal{P}(\mathcal{X}),W_{p})$ is a geodesic space where every two points can be connected by the constant-speed geodesic (Santambrogio, 2015).

Generating samples given a specific label requires to learn conditional distributions given domain labels. In this section, we formulate distances between target and model conditional distributions and derive a tractable upper bound of the Wasserstein distance between conditional distributions.

Denoting the latent variable independent of domain labels by $Z\sim\mathbb{P}_{Z}$, the distribution of the observed domain labels by $\mathbb{P}_{C}$, and the conditional generator by $\text{Gen}:\mathcal{Z}\times\mathcal{C}\to\mathcal{X}$, the model conditional distribution can be expressed as $\mathbb{P}_{\text{Gen}(Z,C)|C}(\cdot|c)$ where $c\in\mathcal{C}$. To learn $\mathbb{P}_{X|C}(\cdot|c)$, we formulate a class of distances between conditional distributions as

where $\mathcal{D}$ is a measure between distributions. Various statistical distances can be considered for $\mathcal{D}$, and when we choose the Kullback-Leibler divergence, Equation (1) can be minimized by maximizing the expectation of the variational lower bound of the conditional log-likelihood over the distribution of domain labels. For the case of the Jensen-Shannon divergence, adversarial learning with discriminator and generator incorporating domain labels minimizes Equation (1).

To bring advantages over $f$-divergences, we focus on Equation (1) equipped with Wasserstein distances to learn the Wasserstein geodesic. Since there is no previous work formulating or deriving properties of the Wasserstein distance between conditional distributions given domain labels, in the next section, we derive an upper bound of Equation (1) that has a tractable representation.

In this section, we lay a theoretical groundwork by deriving a tractable upper bound of the expected Wasserstein distance between conditional distributions.

We first propose a new set of couplings for conditional generation, conditional sub-coupling.

The conditional sub-coupling is the set of all probability measures induced by couplings of conditional distributions. It is nonempty and equals to $\Pi(\mathbb{P}_{X},\mathbb{P}_{Y})$ if $(\mathbb{P}_{X|C}(\cdot|c),\mathbb{P}_{Y|C}(\cdot|c))=(\mathbb{P}_{X},\mathbb{P}_{Y})$ for all $c\in\mathcal{C}$. The following example provides cases where the conditional sub-coupling is a proper subset of $\Pi(\mathbb{P}_{X},\mathbb{P}_{Y})$. Let $N(\mu_{1},\mu_{2},\sigma_{1},\sigma_{2},\rho)$ denote bivariate Gaussian distribution with mean $(\mu_{1},\mu_{2})^{T}$ and covariance $\begin{pmatrix}\sigma_{1}^{2}&\rho\sigma_{1}\sigma_{2}\\ \rho\sigma_{1}\sigma_{2}&\sigma_{2}^{2}\end{pmatrix}$.

Further discussions and proofs about the conditional sub-coupling are provided in Appendix A. With the conditional sub-coupling, we derive an upper bound of the expected $p$-Wasserstein distance in the following theorem.

That is, the minimum transport cost over conditional sub-coupling is an upper bound of expected Wasserstein distance between conditional distributions.

We show a tractable representation of the upper bound in the following theorem. We denote the set of all $\mathbb{Q}_{Z|X,C}$ satisfying $\mathbb{P}_{(Z,C)}(z,c)=\big{(}\int\mathbb{Q}_{Z|X,C}(z|x,c)d\mathbb{P}_{X|C}(x|c)\big{)}\mathbb{P}_{C}(c)$ by $\mathcal{Q}$; the RHS can be considered as an aggregate posterior (Makhzani et al., 2015).

That is, the upper bound of the Wasserstein distance between conditional distributions, the RHS of Equation (2), can be expressed as the infimum of the reconstruction error over encoders $\mathbb{Q}_{Z|X,C}\in\mathcal{Q}$. Note that the integrand in the LHS of Equation (2) depends on the conditioning data $c$ and requires to evaluate the Wasserstein distance for every realization $c$, which is infeasible. In contrast, the derived representation can be computed by solving a stochastic optimization problem. When the terms related to domain label $C$ are removed, Theorem 3 reduces to the representation of the Wasserstein distance between marginal distributions provided by Tolstikhin et al. (2018).

For a motivating example, suppose data come from one of two observed domains whose label values are $c_{0}$ and $c_{1}$. Existing methods in conditional generation literature have considered $\text{Gen}(Z,(1-t)c_{0}+tc_{1})$ as intermediate samples (Zhang et al., 2017b). However, without a strong assumption such as a linear structure, the interplay between $\mathbb{P}_{X|C}(\cdot|c_{0})$, $\mathbb{P}_{X|C}(\cdot|c_{1})$, and $\mathbb{P}_{\text{Gen}(Z,C)|C}(\cdot|(1-t)c_{0}+tc_{1})$ is difficult to formalize.

A desirable property of generated samples for unobserved intermediate domains would be their conditional distributions change smoothly from one observed domain to another. The next section proposes a new conditional generator that constructs samples from distributions on the constant-speed geodesic in the Wasserstein space.

That is, a constant-speed geodesic in a Wasserstein space is a parameterized curve whose speed equals to the Wasserstein distance. Our method yields the conditional distribution given an unobserved intermediate domain label as an interpolation point between the conditional distributions given observed domain labels in the Wasserstein space. Unlike existing methods, the generated distributions are fully characterized by the Wasserstein space $(\mathcal{P}(\mathcal{X}),W_{p})$.

This section proposes the Wasserstein geodesic generator, a novel conditional generator for learning Wasserstein geodesics. Our method learns the conditional distributions given observed domains and the optimal transport maps between them to construct the Wasserstein geodesic. The proposed method consists of three networks: encoder $\text{Enc}:\mathcal{X}\times\mathcal{C}\to\mathcal{Z}$, generator $\text{Gen}:\mathcal{Z}\times\mathcal{C}\to\mathcal{X}$, and transport map $T:\mathcal{X}\times\mathcal{C}^{2}\to\mathcal{X}$.

We first define the optimal transport map, and then provide the proposed method.

The optimal transport map refers maps yielding the minimum transportation cost. The optimal transport map uniquely exists if the cost function is the $p$-th power of $l_{2}$-distance denoted by $\|\cdot\|$ where $p>1$ and measures are absolutely continuous on compact domains. The minimum transportation cost by the optimal transport map is known as $W^{p}_{p}(\mathbb{P}_{X},\mathbb{P}_{Y})$.

We now present the Wasserstein geodesic generator. The proposed method postulates the encoder, generator, and transport map satisfying the following conditions $(A1)$ and $(A3)$ to generate intermediate samples. Here, $d_{\text{Enc}}((x_{1},c_{1}),(x_{2},c_{2})):=\lVert(\text{Enc}(x_{1},c_{1}),c_{1})-(\text{Enc}(x_{2},c_{2}),c_{2})\rVert$ is a metric defined on $\mathcal{Z}\times\mathcal{C}$.

1. (A1)

   (One-to-one mapping between $\mathcal{X}\times\mathcal{C}$ and $\mathcal{Z}\times\mathcal{C}$) For any $(x,c)\in\mathcal{X}\times\mathcal{C}$ and $(z,c)\in\mathcal{Z}\times\mathcal{C}$, $\text{Gen}(\text{Enc}(x,c),c)=x$ and $\text{Enc}(\text{Gen}(z,c),c)=z$.

2. (A2)

   (Absolutely continuous representations) For any $c\in\mathcal{C}$, $\mathbb{P}_{\text{Enc}(X,C)|C}(\cdot|c)$ is absolutely continuous, is defined on a compact set, and has finite second moments.

3. (A3)

   (Optimal transportation) For any observed domain labels $c_{m}$ and $c_{m^{\prime}}$, $(T(\cdot,\cdot,c_{m^{\prime}}),c_{m^{\prime}}):\mathcal{X}\times\{c_{m}\}\to\mathcal{X}\times\{c_{m^{\prime}}\}$ is the optimal transport map from $\mathbb{P}_{(X,c_{m})|C}(\cdot|c_{m})$ to $\mathbb{P}_{(X,c_{m^{\prime}})|C}(\cdot|c_{m^{\prime}})$ w.r.t. $d^{p}_{\text{Enc}}$.

Note that condition $(A1)$ is about inverse relations between the encoder and generator for fixed domain labels, $(A2)$ is to guarantee the existence and uniqueness of optimal transport maps between observed conditional distributions w.r.t. $d^{p}_{\text{Enc}}$, and $(A3)$ is to build paths between observed conditional distributions with optimal transport maps.

The encoder and generator can define a distance-preserving mapping (called isometric mapping) between two Wasserstein spaces, and we can connect their geometric structures, including geodesics and barycenters. With Lemma 6 and optimal transport maps, we can first generate geodesics on $(\mathcal{P}(\mathcal{Z}\times\mathcal{C}),W_{p}(\cdot,\cdot;\lVert\cdot\rVert))$, and then project them to geodesics on $(\mathcal{P}(\mathcal{X}\times\mathcal{C}),W_{p}(\cdot,\cdot;d_{\text{Enc}}))$ to generate intermediate unobserved conditional distributions.

The conditional distributions of the samples generated by the proposed method constitute the Wasserstein geodesic, yielding the minimum transportation cost quantified by $d_{\text{Enc}}$ between the conditional distributions of observed domains.

This section extends our Wasserstein geodesic generator to accommodate scenarios involving multiple observed conditional distributions. We explain how the proposed distribution approximates the centroid of observed conditional distributions with an interpretable upper bound of the approximation error. Furthermore, we derive that the proposed distribution is the Wasserstein barycenter under some conditions.

We first define the Wasserstein barycenter, the centroid of distributions within the Wasserstein space.

Specifically, when $M=2$, the constant-speed geodesic $w(t)$ in Definition 4 serves as the Wasserstein barycenter of two distributions with weights $(1-t,t)$. In this context of defining a centroid between two distributions, the Wasserstein geodesic is referred to as McCann’s interpolant (McCann, 1995). The Wasserstein barycenters have been acknowledged as an effective solution for aggregating high-dimensional distributions (Korotin et al., 2022), across various applications including data augmentation (Huguet et al., 2022; Bespalov et al., 2021; Zhu et al., 2023) and domain adaptation (Montesuma and Mboula, 2021).

The Wasserstein barycenter has several advantages in generating unobserved conditional distributions. First, it provides smooth and stable transitions between observed distributions, which is essential for synthesizing data for new, unobserved domains. This allows the generated data to inherit characteristics from observed conditional distributions without abrupt changes. Second, it minimizes the average optimal transportation costs to observed distributions, thus minimally altering observed data to synthesize unobserved data. Last, the Wasserstein barycenter can be employed to infer the characteristics of unobserved domains, e.g., estimating the conditional average treatment effect in clinical trials (Huguet et al., 2022). Despite these advantages, the computational complexity of estimating Wasserstein barycenter remains a significant bottleneck (Cuturi, 2013; Lin et al., 2020).

We extend the Wasserstein geodesic generator to generate unobserved intermediate distributions with multiple observed distributions. In Theorem 10, we establish an interpretable bound on the error incurred when approximating the Wasserstein barycenter using out proposed distribution. We further demonstrate that, under condition $(A4)$ relating to the condition on the homogeneity of representations across observed domains, the proposed distribution is the Wasserstein barycenter. We begin by introducing a lemma that elucidates how the homogeneity affects on the average squared Wasserstein distances.

In Equation (3), the first term represents the Wasserstein variance (Martinet et al., 2022) of $(\mathbb{P}_{\text{Enc}(X,C)|C}(\cdot|c_{m}))_{m=1}^{M}$, while the second term denotes the variance of $(c_{m})_{m=1}^{M}$ w.r.t. weights $(\alpha_{m})_{m=1}^{M}$. When multiple convex combinations are possible, we select the optimal combination that minimizes the variance, $\sum_{m=1}^{M}\alpha_{m}\lVert\bar{c}-c_{m}\rVert^{2}$. In the special case where $M=2$ and $c$ is univariate, this approach is equivalent to using the two nearest observed distributions to generate intermediate distributions. The Wasserstein variance quantifies the homogeneity of representations across observed conditional distributions and is zero if and only if the distributions of representations from all observed domains are identical, which is to say that the following condition holds.

• (A4)

  $(\text{Homogeneous representations across observed domains})$ For any two observed domain labels $c_{m}$ and $c_{m^{\prime}}$, $\mathbb{P}_{\text{Enc}(X,C)|C}(\cdot|c_{m})=\mathbb{P}_{\text{Enc}(X,C)|C}(\cdot|c_{m^{\prime}})$.

For any $M$ observed domain labels $(c_{m})_{m=1}^{M}$, their convex combination $\bar{c}:=\sum_{m=1}^{M}\alpha_{m}c_{m}$, and $X_{1}\sim\mathbb{P}_{X|C}(\cdot|c_{1})$, the latent interpolation result of $X_{1}$ and its transported results $(T(X_{1},c_{1},c_{m}))_{m=2}^{M}$ can be expressed as $\tilde{X}(\alpha_{1},\dots,\alpha_{M}):=\text{Gen}(\sum_{m=1}^{M}\alpha_{m}\text{Enc}(T(X_{1},c_{1},c_{m}),c_{m})),\bar{c})$. In the subsequent theorem, we derive an upper bound for the difference between the average squared Wasserstein distances with our proposed distribution and with the Wasserstein barycenter. Additionally, when we further suppose the homogeneity of representations, condition $(A4)$, the proposed method generates the Wasserstein barycenter of observed conditional distributions.

To the best of authors’ knowledge, this is the first result that justifies latent interpolation from an optimal transport point of view, without resorting to Gaussian or univariate assumptions. The upper bound represents half of the average squared distances between $T(X_{1},c_{1},c_{m^{\prime}})$ and $T(T(X_{1},c_{1},c_{m}),c_{m},c_{m^{\prime}})$. They are identical in the univariate scenario, but not necessarily in other cases, which makes the non-negative upper bound. Although we have detailed results for the $2$-Wasserstein distance where both existence and uniqueness of barycenters are extensively examined, similar results hold for general $p$-Wasserstein distances. If $(A1)$ through $(A4)$ hold and $\underset{\mathbb{P}\in\cup_{c\in\mathcal{C}}\mathcal{P}(\mathcal{X}\times\{c\})}{\inf}\sum_{m=1}^{M}\alpha_{m}W^{2}_{p}(\mathbb{P},\mathbb{P}_{(X,C)|C}(\cdot|c_{m});d_{\text{Enc}})$ has solutions (called Fréchet mean w.r.t. $W_{p}$), then the proposed distribution is a solution. Note that condition $(A4)$ is weaker than the following condition.

• (A5)

  $(\text{Homogeneous representations})$ $\mathbb{P}_{\text{Enc}(X,C)|C}(\cdot|c)$ is constant w.r.t. $c\in\mathcal{C}$, equivalently,
  $\text{Enc}(X,C)\perp\!\!\!\perp C$.

Most existing conditional generative models, including cVAE, cGAN, and cAAE, assume $(A5)$. When the unobserved domains possess patterns distinctly different from the observed ones, this condition may not satisfied. For example, in face frontalization (Huang et al., 2017), if we have facial images captured from the right and left angles and aim to synthesize frontal facial images, the frontal images might exhibit unique features, such as symmetrical facial structures, clear visibility of both eyes, and a full view of facial landmarks like the nose bridge and forehead. In the following theorem, under $(A5)$, we derive that the generation result follows the true conditional distribution.

In summary, when $(A1)$ through $(A3)$ hold, we can construct distributions that change smoothly in-between observed conditional distributions by generating geodesics. When we further suppose $(A4)$, we can generate data from the Wasserstein barycenter without observations as described in Algorithm 1. Furthermore, the proposed distribution is the true conditional distribution when $(A5)$ holds.

The training of the proposed method is to learn networks satisfying conditions $(A1)$ through $(A4)$ and consists of two steps. The first step is to learn the encoder and generator pair and the second step is to learn the transport map with the learned encoder. These two steps learn the vertices and edges of the geodesic, respectively.

For the first step, motivated by Theorem 3 and condition $(A1)$, we minimize the reconstruction error with two penalty terms

where the first term is the reconstruction error, the second term is to enforce the constraint $\mathcal{Q}$ on the encoder network in the derived representation of the upper bound in Theorem 3, and the last term is to enforce condition $(A1)$. We substitute the deterministic encoder Enc with $\mathbb{Q}_{Z|X,C}$ and $\mathcal{D}$ with $\mathcal{D}_{\text{JS}}$, and set $\mathbb{P}_{(Z,C)}=\mathbb{P}_{Z}\mathbb{P}_{C}$ to learn the Wasserstein geodesic on which the information independent of domain labels is minimally changed. In implementation, we apply GAN for the second term and interpolation results of encoded values for the last term.

For the second step, with learned encoder from the first step, we solve Monge–Kantorovich transportation problems w.r.t. $d^{p}_{\text{Enc}}$,

for all observed domain label values $c_{0}$ and $c_{1}$. The objective is

Here, $\mathbb{P}_{(T(X_{0},C_{0},C_{1}),C_{1},C_{0})}$ is the distribution of $(T(X_{0},C_{0},C_{1}),C_{1},C_{0})$ where $C_{0}$ and $C_{1}$ are independent samples from $\mathbb{P}_{C}$ and $X_{0}$ is sampled from $\mathbb{P}_{X|C}(\cdot|C_{0})$. The first term is the transportation cost measured by $d_{\text{Enc}}^{p}$ and the second term is to enforce constraint on conditional distributions of transported data. Note that the second term is zero if and only if $\mathbb{P}_{(X,c_{1})|C}(\cdot|c_{1})=\mathbb{P}_{(T(X,c_{0},c_{1}),c_{1})|C}(\cdot|c_{0})$ for all $(c_{0},c_{1})$ sampled from $\mathbb{P}_{C}\mathbb{P}_{C}$. The last term is the cycle consistency loss that encourages the inverse relation of the transport map from one domain to another and its vice versa. The cycle consistency loss has been proposed in the Data-to-Data Translation literature (Zhu et al., 2017; Choi et al., 2018) from a heuristic point of view to avoid mode collapse, but in our method, it enforces the transport map to satisfy the inverse relation of optimal transport maps between the two domains. We derive that the minimizer of objective in Equation (6) is unique and is the optimal transport map between observed domains in the following theorem.

In implementation, we apply WGAN with gradient penalty (Gulrajani et al., 2017) for the second term, employ an auxiliary regressor to enhance the visual quality and diversity of generated samples and add a reconstruction error, $\int d(x,T(x,c,c))d\mathbb{P}_{(X,C)}(x,c)$, for a regularization purpose.

With the learned encoder, generator, and transport map, we can generate samples whose conditional distributions are on the Wasserstein geodesic on which the information independent of domain labels is minimally changed.

If either of the two steps of our training algorithm is dropped and an appropriate modification is made, our algorithm reduces to either cAAE or CycleGAN.

Suppose the second step is dropped. The first step alone learns the encoder and generator pair through minimizing objective in Equation (5). If we remove $\mathbb{P}_{C}$ in the second term and drop the third term, objective in Equation (5) reduces to that of cAAE. Thus it is the $\mathbb{P}_{C}$ in the second term that enables the proposed algorithm to learn the features independent of domain labels.

Now suppose the first step is removed. Since the second step optimizes objective in Equation (6) that includes $d_{\text{Enc}}$, it inherently depends on the first step. However, if we redefine the distance not to depend on the encoder, the second step can be operated independently. By not inheriting the $d_{\text{Enc}}$, one no longer learns a transport map while minimally changing contents independent of conditioning variables. If we replace $d_{\text{Enc}}$ in the first term with the $l_{2}$ norm and $W_{p}$ in the second term with the Jensen-Shannon divergence, and if there are only two observed domain labels $c_{0}$ and $c_{1}$, objective in Equation (6) reduces to that of CycleGAN. Inheriting the results from Theorem 12, we can establish the following corollary.

Lu et al. (2019) specifically point out that theoretically, there is no claim on the detailed properties of the mapping established by CycleGAN. On the other hand, Korotin et al. (2019) used the cycle consistency term to promote inverse relations between optimal transport maps, albeit with an objective based on the dual form of $2$-Wasserstein distance, which distinguishes it from CycleGAN’s objective. Because CycleGAN can be viewed a special case of our proposed method, now it can be interpreted in terms of optimal transport theory, which has not been established.

We conduct experiments on the Extended Yale Face Database B (Extended Yale-B) (Georghiades et al., 2001) dataset. The Extended Yale-B dataset consists of face images from $29$ subjects, $10$ poses, and $64$ light conditions. We consider light conditions as domain labels. For light conditions, the azimuth and elevation of the light source are provided. The total number of images is about $15,000$.

We split the Extended Yale-B dataset by subjects to construct training, validation, and test sets. The number of subjects is $21$, $4$, and $4$ for training, validation, and test, respectively. For a data pre-processing, we apply a face detection algorithm proposed by Viola and Jones (2001) to crop the face part.²²2We use the official code in OpenCV (Bradski, 2000). The range of images is scaled to $[0,1]$ and the horizontal flip with probability $0.5$ is applied during training.

For baselines, we consider cAAE, CycleGAN, and StarGAN (Choi et al., 2018). As described in Section 4.5, our encoder and generator pair can reduce to cAAE and our transport map can reduce to CycleGAN. StarGAN is a state-of-the-art Data-to-Data Translation method and we add it as a baseline for the transport map. Algorithms of CycleGAN and StarGAN in their published works are inapplicable to continuous domain labels, so we add source and target light conditions as a input of CycleGAN and change the auxiliary classifier of StarGAN to regressor for experiments.

In all methods, architectures of the encoder and generator networks are adopted from DCGAN (Radford et al., 2015), and the architecture of the transport map is adopted from StarGAN. Architectures are modified to concatenate light conditions to latent variables, and we control the size of the networks for a fair comparison. For both the proposed method and baselines, we train the encoder and generator pair for $100,000$ iterations with batch size of $32$, and train the transport map for $50,000$ iterations with batch size of $16$. We use the Adam (Kingma and Ba, 2014) optimizer and set the initial learning rate to $0.0002$ and to linearly decrease to $0$ for the encoder and generator pair and to $0.0001$ for the transport map. Implementation details including architectures are provided in Appendix B. ³³3The implementation code is provided at the following link:
http://github.com/kyg0910/Wasserstein-Geodesic-Generator-for-Conditional-Distributions.

For evaluation, we use the Fréchet inception distance (FID) (Heusel et al., 2017), a dissimilarity measure between distributions utilizing features of Inception-v3 (Szegedy et al., 2016). We evaluate FID between distributions of the real and generated face images.

We compare the proposed method and baselines with three tasks: (i) Conditional generation for unobserved intermediate domains, (ii) Data-to-Data Translation, and (iii) Latent interpolation with the real data and their translation results. Figures 1, 2, and 3 present results for tasks (i), (ii), and (iii), respectively. The (ii) and (iii) evaluate transport maps when an image is given. We reiterate CycleGAN and StarGAN cannot generate samples while the proposed method can.

Figure 1 presents the conditional generation results for unobserved intermediate domains. We compare the proposed method and cAAE. The proposed method produces face images with clearer eyes, noses, and mouths than baselines. For each method, the leftmost and rightmost columns show generation results for observed domains whose values of (azimuth, elevation) are $(0,0)$ and $(70,45)$, respectively, and intermediate columns show results for unobserved intermediate domains. For the proposed method, as described in Section 4.2, we first generate samples in the leftmost column, then transport the samples to the domain of the rightmost column, and finally apply latent interpolation where $t$ increases from $0$ to $1$ at equal intervals. For cAAE, we fix latent variable $Z$ for each row and interpolate domain label $C$ with equal spacing.

Figure 2 compares the transportation results. We compare the proposed method, CycleGAN, and StarGAN. The leftmost panel visualizes the transportation results. The proposed method gradually casts a shadow reflecting the three-dimensional structure of the nose and mouth in the face, which makes the outcomes visually sharper and more plausible than baselines. The bottom row shows face images from various azimuth for a fixed subject, pose, and elevation. The elevation is $0$ and azimuth values are shown at the bottom. For the first three rows, the middle column shows the real data in the fourth row and other columns show transportation results by various methods for observed domains corresponding to each column. Each method translates the real data in the middle column to other observed domains. The middle panel shows the box plots of FID scores from various transport maps. The lower values are better and the proposed method outperforms baselines. The means of the FID scores of CycleGAN, StarGAN, and the proposed method with standard deviation are 174.3 (28.9), 109.7 (9.7), and 74.2 (10.9), respectively. To calculate the FID scores, we transport real face images from (azimuth, elevation) of $(0,0)$ to other observed domains and evaluate FID between transportation results and real images for every domains. The absolute value of azimuth and elevation is less than or equal to $50$ and $35$, respectively, to remove outliers from extreme light conditions. The rightmost panel presents the surfaces of the FID scores from various transport maps as functions of azimuth and elevation. In all combinations of light conditions, the proposed method outperforms baselines. We estimate scores for unobserved domains by interpolating scores from adjacent observed domains to draw the plot.

Figure 3 presents the results of latent interpolation on unobserved intermediate domains with a real image and its translation result. We compare the proposed method, CycleGAN, and StarGAN. We fix the encoder and generator pair for all methods, and only change the transport map. The bottom row shows the ground-truth images of fixed subject and pose from two observed domains with (azimuth, elevation) of $(0,0)$, for the leftmost, and $(70,0)$, for the rightmost. In the top three rows, the leftmost column shows the ground-truth, the rightmost column shows the transportation results of the ground-truth by various transport maps, and the intermediate columns show latent interpolation results for unobserved intermediate domains. As in Figure 2, the outputs of the proposed method are visually sharper and more plausible than baselines.