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 proposes a joint Wasserstein Auto-Encoder (JWAE) method to solve the problem of joint distribution matching. Instead of ﬁnding a coupling, the paper seeks a decoupling to make the primal problem of Wasserstein distance tractable. The decoupled version of joint Wasserstein distance is used for empirical reconstruction losses of within-domain Auto-Encodings and cyclic mappings. In addition, two GAN divergences are used to learn the cross-domain mappings such that the generated distributions are close to the real distribution, and another GAN divergence is imposed to align the latent distributions generated by two Auto-Encoders. Later, the paper applies the proposed model on the interpolation based video-to-video synthesis problem.