Learning Joint Wasserstein Auto-Encoders for Joint Distribution Matching

1 INTRODUCTION . The joint distribution matching problem has attracted extensive attention in computer vision , such as unsupervised image-to-image translation ( Zhu et al. , 2017 ; Liu et al. , 2017 ) and video-to-video synthesis ( Bashkirova et al. , 2018 ) . The goal of this problem is to learn the bidirectional mappings between unpaired data in two different domains . Unlike the marginal distribution in each domain , learning a joint distribution is often ignored and has the following two critical challenges . The first key challenge , from a probabilistic modeling perspective , is how to exploit the joint distribution of unpaired data by learning the bidirectional mappings between two different domains . In the unsupervised learning setting , there are two sets of samples drawn separately from two marginal distributions in two domains . Based on the coupling theory ( Lindvall , 2002 ) , there exist an infinite set of joint distributions given two marginal distributions , and hence infinite bidirectional mappings between two different domains . Therefore , directly learning the joint distribution without additional information between the marginal distributions is a highly ill-posed problem . Recently , many studies ( Zhu et al. , 2017 ; Yi et al. , 2017 ; Kim et al. , 2017 ) have been proposed to learn the mappings in two domains separately , which may incur the joint distribution mismatching issue . Therefore , how to exploit sufficient information from the joint distribution still remains an open question . Another important challenge is that the generalization ability w.r.t . the learned joint distribution of two different domains is still unknown . Existing theoretical results ( Pan et al. , 2018 ; Galanti et al. , 2018 ) ignore the joint distribution of different data and can not guarantee the generalization ability of such joint distribution . Moreover , it is also very hard to evaluate the generalization ability practically . Regarding this issue , according to ( Bojanowski et al. , 2018 ) , the generalization ability can be evaluated by the interpolation performance in the target domain . In this sense , we can extend image-to-image translation to video space by performing video interpolation in one domain and investigating the performance of the translated video in another domain . To achieve this , one may directly apply existing unsupervised image-to-image translation methods ( Zhu et al. , 2017 ; Kim et al. , 2017 ; Yi et al. , 2017 ) . However , these methods may result in significantly incoherent videos with low visual quality . Therefore , it is important to design an effective joint distribution learning method and provide necessary theoretical analysis . Regarding the above two challenges , in this paper , we propose a Joint Wasserstein Auto-Encoders ( JWAE ) to learn the bidirectional mappings between two domains by minimizing Wasserstein distance of joint distributions . Relying on the optimal transport theory , we are able to exploit sufficient information by matching latent distributions of images in two domains . The contributions of this paper are summarized as follows : • We propose a novel JWAE to solve the joint distribution matching problem . Based on Theorem 1 , an intractable primal problem of optimal transport can be reduced to a simple optimization problem . More critically , our method is a generalization of CycleGAN ( Liu et al. , 2017 ) and UNIT ( Liu et al. , 2017 ) . • We provide a generalization bound of JWAE ( see Theorem 4 ) . In particular , we theoretically prove that the generalization ability of our method w.r.t . the learned joint distribution can be guaranteed by minimizing Wasserstein distance of joint distributions . • To practically verify the generalization ability , we apply our method to unsupervised video-tovideo synthesis and obtain two visually smooth videos in two different domains . Experiments on real-world datasets show the superiority of the proposed method over several state-of-the-arts . 2 RELATED WORK . In this paper , we consider the joint distribution matching problem . Recently , this problem has attracted extensive attention in image-to-image translation and video-to-video synthesis . Image-to-image translation . Recently , Generative adversarial networks ( GAN ) ( Goodfellow et al. , 2014 ; Cao et al. , 2018 ; Salimans et al. , 2018 ) , Variational Auto-Encoders ( VAE ) ( Kingma & Welling , 2014 ) and Wasserstein Auto-Encoders ( WAE ) ( Tolstikhin et al. , 2017 ) have emerged as popular techniques for the image-to-image translation ( I2IT ) problem . For the unsupervised I2IT problem , CycleGAN ( Zhu et al. , 2017 ) , DiscoGAN ( Kim et al. , 2017 ) and DualGAN ( Yi et al. , 2017 ) aim at minimizing the adversarial loss and the cycle-consistent loss in different domains , which may induce a joint distribution mismatching issue . To address this , CoGAN ( Liu & Tuzel , 2016 ) learns a joint distribution by enforcing a weight-sharing constraint . Moreover , UNIT ( Liu et al. , 2017 ) builds upon CoGAN by using a shared-latent space assumption and the same weight-sharing constraint . However , these methods are not well-supported by any theoretical justifications . Video-to-video synthesis . In this paper , we consider unsupervised video-to-video synthesis ( V2VS ) problem . Existing image-to-image methods ( Zhu et al. , 2017 ; Kim et al. , 2017 ; Yi et al. , 2017 ) can not be directly used in the video-to-video synthesis problem , we further combine some video frame interpolation methods ( Zhou et al. , 2016 ; Ji et al. , 2017 ; Niklaus et al. , 2017 ; Liu et al. , 2018 ) to synthesize video . Although UNIT ( Liu et al. , 2017 ) can be applied to video synthesis by interpolating in the latent space , it often results in temporally incoherent videos of low visual quality . Recently , a video-to-video translation method ( Bashkirova et al. , 2018 ) is proposed to translate a video in one domain to a video in another domain , but this method can not conduct video frame interpolation . Moreover , Wang et al . ( 2018 ) propose a video-to-video synthesis method and synthesize video results , but it can not work for the unsupervised learning setting . 3 PROBLEM DEFINITION . Notations . We use calligraphic letters ( e.g. , X ) for space , capital letters ( e.g. , X ) for random variables , and bold lower case letter ( e.g. , x ) for their corresponding values . We denote probability distributions with capital letters ( i.e. , P ( X ) ) and corresponding densities with bold lower case letters ( i.e. , p ( x ) ) . Let ( X , PX ) be the domain , and P ( X ) be the set of all the probability measures over X , and PX be the marginal distribution over X . Sx= { xi } Mi=1 and Sy= { yi } Ni=1 are two sets of unpaired training data . For a set S and two functions F : S→R and G : S→R , we denote F ( s ) .G ( s ) , ∀s∈S if and only if ∃ C1 , C2 > 0 ( independent of s ) such that F ( s ) ≤C1·G ( s ) +C2 . 3.1 JOINT DISTRIBUTION MATCHING PROBLEM . This paper considers the Joint Distribution Matching Problem for the unsupervised image-to-image translation task . Given two domains ( X , PX ) and ( Y , PY ) , where PX and PY are distributions over X and Y , our goal is to learn cross-domain mappings , i.e. , f : X→Y and g : Y→X . Then we construct joint distributions in two domains , i.e. , PA ( X , f ( X ) ) , X∈X and PB ( g ( Y ) , Y ) , Y ∈Y , and we minimize the following distribution divergence d ( PA , PB ) between these two joint distributions , min f∈F , g∈G d ( PA , PB ) , ( 1 ) where F and G are the sets of cross-domain mappings . Note that d ( · , · ) is an arbitrary distribution divergence , e.g. , Wasserstein distance . In this paper , we study the joint distribution matching problem from the optimal transport ( OT ) point of view , which would bring some interesting results . 4 PROPOSED METHOD . 4.1 WASSERSTEIN DISTANCE OF JOINT DISTRIBUTION . To address the joint distribution matching problem , one can learn a shared latent space Z for two different domains ( Liu et al. , 2017 ) . In this sense , any pair of images in different domains can be mapped to the same latent representation . Inversely , there exist generative models PG1 ( X ′|Z ) and PG2 ( Y ′|Z ) that map a shared latent code Z to X ′=G1 ( Z ) and Y ′=G2 ( Z ) , respectively . Then PA ( X , Y ′ ) and PB ( X ′ , Y ) are two joint distributions between real and generated images , we can minimize Wasserstein distanceWc ( PA , PB ) between joint distributions PA and PB , i.e. , Wc ( PA , PB ) = min P∈P ( PA , PB ) E ( X , Y ′ ; X′ , Y ) ∼P [ c ( X , Y ′ ; X ′ , Y ) ] , ( 2 ) where P ( PA , PB ) is the set of couplings which is composed of joint probability distributions with the probability distributions ( PA , PB ) , and c is any measurable cost function . In practice , there are two important challenges on Wasserstein distance and the cost function . First , directly optimizing Problem ( 2 ) raises intractable computational difficulties ( Genevay et al. , 2018 ) . Second , how to choose a appropriate cost function is very challenging . In this paper , we set c ( X , Y ′ ; X ′ , Y ) =c1 ( X , X ′ ) +c2 ( Y ′ , Y ) ( Bhushan Damodaran et al. , 2018 ) , where c1 and c2 can be any metric to measure the distance in X×X and Y×Y , respectively . This helps to derive the following theorem so that the intractable Problem ( 2 ) can be reduced to a simple optimization problem . Theorem 1 Given two deterministic models PG1 ( X ′|Z ) and PG2 ( Y ′|Z ) as Dirac measures , i.e. , PG1 ( X ′|Z = z ) = δG1 ( z ) and PG2 ( Y ′|Z = z ) = δG2 ( z ) for all z ∈ Z , we have Wc ( PA , PB ) = inf Q∈Q1 EPXEQ ( Z1|X ) [ c1 ( X , G1 ( Z1 ) ) ] + inf Q∈Q2 EPY EQ ( Z2|Y ) [ c2 ( G2 ( Z2 ) , Y ) ] , ( 3 ) whereQ1= { Q ( Z1|X ) |QZ1=PZ=QZ2 , PY =PG2 } andQ2= { Q ( Z2|Y ) |QZ1=PZ=QZ2 , PX=PG1 } are the set of all probabilistic encoders , where QZ1 and QZ2 are the marginal distributions of Z1∼Q ( Z1|X ) and Z2∼Q ( Z2|Y ) , where X∼PX and Y∼PY , respectively . 1 As previously mentioned , finding an optimal couplings between joint distributions PA and PB is very challenging . Fortunately , according to Theorem 1 , we can instead optimize problem ( 3 ) for joint distribution matching . The details of objective functions and optimizations are given below . 1See supplementary materials for the proof .

Summary: This paper studies the joint distribution matching problem where given data samples in two different domains, one is interested in learning a bi-directional mapping between unpaired data elements in these domains. The paper proposes a joint Wasserstein auto-encoder (JWAE) to solve this problem. The paper shows that under the decomposable cost metric and deterministic decoding maps, the optimization problem associated with the JWAE formulation can be reduced to a tractable optimization problem. The paper also establishes a generalization bound for the JWAE formulation. Finally, the paper conducts an experimental evaluation of the proposed solution with the help of a video-to-video synthesis problem and show improved performance as compared to the existing results in the literature.