Exploiting Chain Rule and Bayes’ Theorem to Compare Probability Distributions
Year: 2022

Abstract:
To measure the difference between two probability distributions, referred to as the source and target, respectively, we exploit both the chain rule and Bayes’ theorem to construct conditional transport (CT), which is constituted by both a forward component and a backward one. The forward CT is the expected cost of moving a source data point to a target one, with their joint distribution defined by the product of the source probability density function (PDF) and a sourcedependent conditional distribution, which is related to the target PDF via Bayes’ theorem. The backward CT is defined by reversing the direction. The CT cost can be approximated by replacing the source and target PDFs with their discrete empirical distributions supported on mini-batches, making it amenable to implicit distributions and stochastic gradient descent-based optimization. When applied to train a generative model, CT is shown to strike a good balance between modecovering and mode-seeking behaviors and strongly resist mode collapse. On a wide variety of benchmark datasets for generative modeling, substituting the default statistical distance of an existing generative adversarial network with CT is shown to consistently improve the performance. PyTorch code is provided.

heading: 1 Introduction
text: Measuring the difference between two probability distributions is a fundamental problem in statistics and machine learning [1–3]. A variety of statistical distances, such as the Kullback–Leibler (KL) divergence [4], Jensen–Shannon (JS) divergence [5], maximum mean discrepancy (MMD) [6], and Wasserstein distance [7], have been proposed to quantify the difference. They have been widely used for generative modeling with different mode covering/seeking behaviors [8–13]. The KL divergence, directly related to both maximum likelihood estimation and variational inference [14–16], requires the two probability distributions to share the same support and is often inapplicable if either is an implicit distribution whose probability density function (PDF) is unknown [17–20]. Variational auto-encoders (VAEs) [8], the KL divergence based deep generative models, are stable to train, but often exhibit mode-covering behaviors in its generated data, producing blurred images. The JS divergence is directly related to the min-max loss of a generative adversarial net (GAN) when the discriminator is optimal [9], while the Wasserstein-1 distance is directly related to the min-max loss of a Wasserstein GAN [11], whose critic is optimized under the 1-Lipschitz constraint. However, it is difficult to maintain a good balance between the updates of the generator and discriminator/critic, making (Wasserstein) GANs notoriously brittle to train. MMD [6] is an RKHS-based statistical distance behind MMD-GANs [10, 21, 22], which have also shown promising results in generative modeling when trained with a min-max loss. Different from VAEs, these GAN-based models often exhibit mode dropping and face the danger of mode collapse if not well tuned during the training.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
In this paper, we introduce conditional transport (CT) as a new method to quantify the difference between two probability distributions, which will be referred to as the source distribution pX(x) and target distribution pY (y), respectively. The construction of CT is motivated by the following observation: the difference between pX(x) and pY (y) can be reflected by the expected difference of two dependent random variables x and y, whose joint distribution π(x,y) is constrained by both pX(x) and pY (y) in a certain way. Denoting c(x,y) ≥ 0 as a cost function to measure the difference between points x and y, such as c(x,y) = ‖x−y‖22, the expected difference is expressed as Eπ(x,y)[c(x,y)]. A basic way to constrain π(x,y) with both pX(x) and pY (y) is to let π(x,y) = pX(x)pY (y), which means drawing x and y independently from pX(x) and pY (y), respectively; this expected difference EpX(x)pY (y)[c(x,y)] is closely related to the energy distance [23]. Another constraining method is to require both ∫ π(x,y)dy = pX(x) and ∫ π(x,y)dx = pY (y), under which minπ{Eπ(x,y)[c(x,y)]} becomes the Wasserstein distance [7, 24–26]. A key insight of this paper is that by exploiting the chain rule and Bayes’ theorem, there exist two additional ways to constrain π(x,y) with both pX(x) and pY (y): 1) A forward CT that can be viewed as moving the source to target distribution; 2) A backward CT that reverses the direction. Our intuition is that given a source (target) point, it is more likely to be moved to a target (source) point closer to it. More specifically, if the target distribution does not provide good coverage of the source density, then there will exist source data points that lie in low-density regions of the target, making the expected cost of the forward CT high. Therefore, we expect that minimizing the forward CT will encourage the target distribution to exhibit a mode-covering behavior with respect to (w.r.t.) the source PDF. Reversing the direction, we expect that minimizing the backward CT will encourage the target distribution to exhibit a mode-seeking behavior w.r.t. the source PDF. Minimizing the combination of both is expected to strike a good balance between these two distinct behaviors.
To demonstrate the use of CT, we apply it to train implicit (or explicit) distributions to model both 1D and 2D toy data, MNIST digits, and natural images. The implicit distribution is defined by a deep generative model (DGM) that is simple to sample from. We provide empirical evidence to show how to control the mode-covering versus mode-seeking behaviors by adjusting the ratio of the forward CT versus backward CT. To train a DGM for natural images, we focus on adapting existing GANs, with minimal changes to their settings except for substituting the statistical distances in their loss functions with CT. We leave tailoring the network architectures and settings to CT for future study. Modifying the loss functions of various existing DGMs with CT, our experiments show consistent improvements in not only quantitative performance and generation quality, but also learning stability. Our code is available at https://github.com/JegZheng/CT-pytorch.

heading: 2 Chain rule and Bayes’ theorem based conditional transport
text: Exploiting the chain rule and Bayes’ theorem, we can constrain π(x,y) with both pX(x) and pY (y) in two different ways, leading to the forward CT and backward CT, respectively. To define the forward CT, we use the chain rule to factorize the joint distribution as
π(x,y) = pX(x)πY (y |x), where πY (y |x) is a conditional distribution of y given x. This construction ensures ∫ π(x,y)dy =
pX(x) but not ∫ π(x,y)dx = pY (y). Denote dφ(h1,h2) ∈ R as a function parameterized by φ,
which measures the difference between two vectors h1,h2 ∈ RH of dimension H . While allowing∫ π(x,y)dx 6= pY (y), to appropriately constraint π(x,y) by pY (y), we treat pY (y) as the prior distribution, view e−dφ(x,y) as an unnormalized likelihood term, and follow Bayes’ theorem to define
πY (y |x) = e−dφ(x,y)pY (y)/Q(x), Q(x) := ∫ e−dφ(x,y)pY (y)dy, (1)
where Q(x) is a normalization term that ensures ∫ πY (y |x)dy = 1. We refer to πY (y |x) as the forward “navigator,” which specifies how likely a given x will be mapped to a target point y ∼ pY (y). We now define the cost of the forward CT as
C(X → Y ) = Ex∼pX(x)Ey∼πY (· |x)[c(x,y)]. (2)
In the forward CT, we expect large c(x,y) to typically co-occur with small πY (y |x) as long as pY (y) provides a good coverage of the density of x. Thus minimizing the forward CT cost is expected to encourage pY (y) to exhibit a mode-covering behavior w.r.t. pX(x). Such kind of behavior is also
expected when minimizing the forward KL divergence as KL(pX ||pY ) = Ex∼pX [ ln pX(x)pY (x) ] , which calls for pY (x) > 0 whenever pX(x) > 0.
Reversing the direction, we construct the backward CT, where the joint is factorized as π(x,y) = pY (y)πX(x |y) and the backward navigator is defined as
πX(x |y) = e−dφ(x,y)pX(x)/Q(y), Q(y) := ∫ e−dφ(x,y)pX(x)dx. (3)
This ensures ∫ π(x,y)dx = pY (y); while allowing ∫ π(x,y)dy 6= pX(x), it constrains π(x,y) by treating pX(x) as the prior to construct πX(x |y). The backward CT cost is now defined as
C(X ← Y ) = Ey∼pY (y)Ex∼πX(· |y)[c(x,y)]. (4)
In the backward CT, we expect large c(x,y) to typically co-occur with small πX(x |y) as long as pX(x) has good coverage of the density of y. Thus minimizing the backward CT cost is expected to encourage pY (y) to exhibit a mode-seeking behavior w.r.t. pX(x). Such kind of behavior is also expected when minimizing the reverse KL divergence as KL(pY ||pX) = Ex∼pY [ ln pY (x)pX(x) ] , which allows pY (x) = 0 when pX(x) > 0 and it is fine for pY to just fit some portion of pX .
In comparison to the forward and revers KLs, the proposed forward and backward CT are more broadly applicable as they don’t require pX and pY to share the same distribution support and have analytic PDFs. For the cases where the KLs can be evaluated, we introduce
D(X,Y ) = KL(pX ||pY )− KL(pY ||pX)
as a formal way to quantify the mode-seeking and mode-covering behavior of pY w.r.t. pX , with D(X,Y ) > 0 implying mode seeking and with D(X,Y ) < 0 implying mode covering.
Combining both the forward and backward CTs, we now define the CT cost as
Cρ(X,Y ) := ρC(X → Y ) + (1− ρ)C(X ← Y ), (5)
where ρ ∈ [0, 1] is a parameter that can be adjusted to encourage pY (y) to exhibit w.r.t. pX(x) mode-seeking (ρ = 0), mode-covering (ρ = 1), or a balance of two distinct behaviors (ρ ∈ (0, 1)). By definition we have Cρ(X,Y ) ≥ 0, where the equality can be achieved when pX = pY and the navigator parameter φ is optimized such that e−dφ(x,y) is equal to one if and only if x = y and zero otherwise. We also have Cρ=0.5(X,Y ) = Cρ=0.5(Y,X). We fix ρ = 0.5 unless specified otherwise.

heading: 2.1 Conjugacy based analytic conditional distributions
text: Estimating the forward and backward CTs involves πY (y |x) and πX(x |y), respectively. Both conditional distributions, however, are generally intractable to evaluate and sample from, unless pX(x) and pY (y) are conjugate priors for likelihoods proportional to e−d(x,y), i.e., πX(x |y) and πY (y |x) are in the same probability distribution family as pX(x) and pY (y), respectively. For example, if d(x,y) = ‖x− y‖22 and both pX(x) and pY (y) are multivariate normal distributions, then both πX(x |y) and πY (y |x) will follow multivariate normal distributions. To be more specific, we provide a univariate normal based example, with x, y, φ, θ ∈ R and
pX(x) = N (0, 1), pY (y) = N (0, eθ), dφ(x, y) = (x− y)2/(2eφ), c(x, y) = (x− y)2. (6)
Here we have D(X,Y ) = KL[N (0, 1)||N (0, eθ)]− KL[N (0, eθ)||N (0, 1)] = θ − sinh(θ), which is positive when θ < 0, implying mode-seeking, and negative when θ > 0, implying mode-covering. As shown in Appendix C, we have analytic forms of the forward and backward navigators as
πY (y |x) = N (σ(θ − φ)x, σ(θ − φ)eφ), πX(x | y) = N (σ(−φ)y, σ(φ)),
where σ(a) = 1/(1 + e−a) denotes the sigmoid function, and forward and backward CT costs as
C(X → Y ) = σ(φ− θ)(eθ + σ(φ− θ)), C(X ← Y ) = σ(φ)(1 + σ(φ)eθ).
As a proof of concept, we illustrate the optimization under CT using the above example, for which θ = 0 is the optimal solution that makes pX = pY . Thus when applying gradient descent to minimize the CT cost Cρ=0.5(X,Y ), we expect the generator parameter θ → 0 with proper learning dynamic, as long as the learning of the navigator parameter φ is appropriately controlled. This is confirmed by
Fig. 1, which shows that as the navigator φ gets optimized by minimizing CT cost, it is more obvious that θ will minimize the CT cost at zero. This suggests that the navigator parameter φ mainly plays the role in assisting the learning of θ. The right four subplots describe the log-scale curves of forward cost, backward cost and bi-directional CT costs w.r.t. θ as φ gets optimized to four different values. It is worth noting that the forward cost is minimized at eθ > 1, which implies a mode-covering behavior, and the backward cost is minimized at eθ → 0, which implies a mode-seeking behavior, while the bi-directional cost is minimized at around the optimal solution eθ = 1; the forward CT cost exhibits a flattened curve on the right hand side of its minimum, adding to which the backward CT cost not only moves that minimum left, making it closer to θ = 0, but also raises the whole curve on the right hand side, making the optimum of θ become easier to reach via gradient descent.
To apply CT in a general setting where the analytical forms of the distributions are unknown, there is no conjugacy, or we only have access to random samples from the distributions, below we show we can approximate the CT cost by replacing both pX(x) and pY (y) with their corresponding discrete empirical distributions supported on mini-batches. Minimizing this approximate CT cost, amenable to mini-batch SGD based optimization, is found to be effective in driving the target (model) distribution pY towards the source (data) distribution pX , with the ability to control the mode-seeking and mode-covering behaviors of pY w.r.t. pX .

heading: 2.2 Approximate CT given empirical samples
text: Below we use generative modeling as an example to show how to apply the CT cost in a general setting that only requires access to random samples of both x and y. Denote x as a data taking its value in RV . In practice, we observe a finite set X = {xi}|X |i=1, consisting of |X | data samples assumed to be iid drawn from pX(x). GivenX , the usual task is to learn a distribution to approximate pX(x), for which we consider a deep generative model (DGM) defined as y = Gθ( ), ∼ p( ), where Gθ is a generator that transforms noise ∼ p( ) via a deep neural network parameterized by θ to generate random sample y ∈ RV . While the PDF of the generator, denoted as pY (y;θ), is often intractable to evaluate, it is straightforward to draw y ∼ pY (y;θ) with Gθ. While knowing neither pX(x) nor pY (y;θ), we can obtain discrete empirical distributions pX̂N and pŶM supported on mini-batches x1:N and y1:M , as defined below, to guide the optimization of Gθ in an iterative manner. With N random observations sampled without replacement from X , we define
pX̂N (x) = 1 N ∑N i=1 δ(x− xi), {x1, . . . ,xN} ⊆ X (7)
as an empirical distribution for x. Similarly, with M random samples of the generator, we define
pŶM (y) = 1 M ∑M j=1 δ(y − yj), yj = Gθ( j), j iid∼ p( ) . (8)
Substituting pY (y;θ) in (2) with pŶM (y), the continuous forward navigator becomes a discrete one as
π̂Y (y |x) = ∑M j=1 π̂M (yj |x,φ)δyj , π̂M (yj |x,φ) :=
e−dφ(x,yj)∑M j′=1 e −dφ(x,yj′ ) . (9)
Thus given pŶM , the cost of a forward CT can be approximated as
Cφ,θ(X → ŶM ) = E y1:M iid∼ pY (y;θ) Ex∼pX(x)
[∑M j=1 c(x,yj)π̂M (yj |x,φ) ] , (10)
which can be interpreted as the expected cost of following the forward navigator to stochastically transport a random source point x to one of the M randomly instantiated “anchors” of the target distribution. Similar to previous analysis, we expect this approximate forward CT to stay small as long as pY (y;θ) exhibits a mode covering behavior w.r.t. pX(x).
Similarly, we can approximate the backward navigator and CT cost as π̂X(x |y) = ∑N i=1 π̂N (xi |y,φ)δxi , π̂N (xi |y,φ) :=
e−dφ(xi,y)∑N i′=1 e −dφ(xi′ ,y) ,
Cφ,θ(X̂N ← Y ) = E x1:M iid∼ pX(x) Ey∼pY (y;θ)
[∑N i=1 c(xi,y)π̂N (xi |y,φ) ] . (11)
Similar to previous analysis, we expect this approximate backward CT to stay small as long as pY (y;θ) exhibits a mode-seeking behavior w.r.t. pX(x).
Combining (10) and (11), we define the approximate CT cost as
Cφ,θ,ρ(X̂N , ŶM ) = ρCφ,θ(X → ŶM ) + (1− ρ)Cφ,θ(X̂N ← Y ), (12) an unbiased sample estimate of which, given mini-batches x1:N and y1:M , can be expressed as
Lφ,θ,ρ(x1:N ,y1:M ) = ∑N i=1 ∑M j=1 c(xi,yj) ( ρ N π̂M (yj |xi,φ) + 1−ρ M π̂N (xi |yj ,φ) ) = ∑N i=1 ∑M j=1 c(xi,yj) ( ρ N
e−dφ(xi,yj)∑M j′=1 e −dφ(xi,yj′ ) + 1−ρM e−dφ(xi,yj)∑N i′=1 e −dφ(xi′ ,yj)
) . (13)
Lemma 1. Approximate CT in (12) is asymptotic as limN,M→∞ Cφ,θ,ρ(X̂N , ŶM ) = Cφ,θ,ρ(X,Y ).

heading: 2.3 Cooperatively-trained or adversarially-trained feature encoder
text: To apply CT for generative modeling of high-dimensional data, such as natural images, we need to define an appropriate cost function c(x,y) to measure the difference between two random points. A naive choice is some distance between their raw feature vectors, such as c(x,y) = ‖x − y‖22, which, however, is known to often poorly reflect the difference between high-dimensional data residing on low-dimensional manifolds. For this reason, with cosine similarity [27] as cos(h1,h2) :=
hT1 h2√ hT1 h1 √ hT2 h2
, we further introduce a feature encoder Tη(·), parameterized by η, to help redefine the point-to-point cost and both navigators as
cη(x,y) = 1− cos(Tη(x), Tη(y)), dφ ( Tη(x) ‖Tη(x)‖ , Tη(y) ‖Tη(y)‖ ) . (14)
To apply the CT cost to train a DGM, we find that the feature encoder Tη(·) can be learned in two different ways: 1) Cooperatively-trained: Training them cooperatively by alternating between two different losses: training the generator under a fixed Tη(·) with the CT loss, and training Tη(·) under a fixed generator with a different loss, such as the GAN discriminator loss, WGAN critic loss, and MMD-GAN [10] critic loss. 2) Adversarially-trained: Viewing the feature encoder as a critic and training it to maximize the CT cost, by not only inflating the point-to-point cost, but also distorting the feature space used to construct the forward and backward navigators’ conditional distributions.
To be more specific, below we present the details for the adversarial way to train Tη . Given training data X , to train the generator Gθ, forward navigator πφ(y |x), backward navigator πφ(x |y), and encoder Tη , we view the encoder as a critic and propose to solve a min-max problem as
min φ,θ max η E x1:N⊆X , 1:M iid∼ p( ) [Lφ,θ,ρ,η(x1:N , {Gθ( j)}Mj=1)], (15)
where Lφ,θ,ρ,η is defined the same as in (13), except that we replace c(xi,yj) and dφ(·, ·) with their corresponding ones shown in (14) and use reparameterization in (8) to draw y1:M := {Gθ( j)}Mj=1. With SGD, we update φ and θ using∇φ,θLφ,θ,ρ,η(x1:N , {Gθ( j)}Mj=1)) and, if the feature encoder is adversarially-trained, update η using −∇ηLφ,θ,ρ,η(x1:N , {Gθ( j)}Mj=1)). We find by experiments that both ways to learn the encoder work well, with the adversarial one generally providing better performance. It is worth noting that in (Wasserstein) GANs, while the adversarially-trained discriminator/critic plays a similar role as a feature encoder, the learning dynamics between the discriminator/critic and generator need to be carefully tuned to maintain training stability and prevent trivial solutions (e.g., mode collapse). By contrast, the feature encoder of the CT cost based DGM can be stably trained in two different ways. Its update does not need to be well synchronized with the generator and can be stopped at any time of the training.

heading: 3 Related work
text: In practice, variational auto-encoders [8], the KL divergence based deep generative models, are stable to train, but often exhibit mode-covering behaviors and generate blurred images [28–32]. By contrast, both GANs and Wasserstein GANs can generate photo-realistic images, but they often suffer from stability and mode collapse issues, requiring the update of the discriminator/critic to be well synchronized with that of the generator. This paper introduces conditional transport (CT) as a new method to quantify the difference between two probability distributions. Deep generative models trained under CT not only allow the balance between mode-covering and mode-seeking behaviors to be adjusted, but also allow the encoder to be pretrained or frozen at any time during cooperative/adversarial training.
As the JS divergence requires the two distributions to have the same support, the Wasserstein distance is often considered as more appealing for generative modeling as it allows the two distributions to have non-overlapping support [24–26]. However, while GANs and Wasserstein GANs in theory are connected to the JS divergence and Wasserstein distance, respectively, several recent works show that they should not be naively understood as the minimizers of their corresponding statistical distances, and the role played by their min-max training dynamics should not be overlooked [33–35]. In particular, Fedus et al. [34] show that even when the gradient of the JS divergence does not exist and hence GANs are predicted to fail from the perspective of divergence minimization, the discriminator is able to provide useful learning signal. Stanczuk et al. [35] show that the dual form based Wasserstein GAN loss does not provide a meaningful approximation of the Wasserstein distance; while primal form based methods could better approximate the true Wasserstein distance, they in general clearly underperform Wasserstein GANs in terms of the generation quality for high-dimensional data, such as natural images, and require an inner loop to compute the transport plan for each mini-batch, leading to high computational cost [12, 35–38]. See previous works for discussions on the approximation error and gradient bias when estimating the Wasserstein distance with mini-batches [10, 23, 39, 40].
MMD-GAN [10, 21, 22] that calculates the MMD statistics in the latent space of a feature encoder is the most similar to the CT cost in terms of the actual loss function used for optimization. In particular, both the MMD-GAN loss and CT loss, given mini-batches x1:N and y1:M , involve computing the differences of all NM pairs (xi,yj). Different from MMD-GAN, there is no need in CT to choose a kernel and tune its parameters. We provide below an ablation study to evaluate both 1) MMD generator + CT encoder and 2) MMD encoder + CT generator, which shows 1) performs on par with MMD, while 2) performs clearly better than MMD and on par with CT.

heading: 4 Experimental results
text: Forward and backward analysis: To empirically verify our previous analysis of the mode covering (seeking) behavior of the forward (backward) CT, we train a DGM with (12) and show the corresponding interpolation weight from the forward CT cost to the backward one, which means CTρ reduces from forward CT (ρ = 1), to the CT in (12) (ρ ∈ (0, 1)), and to backward CT (ρ = 0). We consider the squared Euclidean (i.e. L22) distance to define both cost c(x,y) = ‖x − y‖22 and dφ(x,y) = ‖Tφ(x) − Tφ(y)‖22, where Tφ denotes a neural network parameterized by φ. We consider a 1D example of a bimodal Gaussian mixture pX(x) = 14N (x;−5, 1) + 3 4N (x; 2, 1) and a 2D example of 8-modal Gaussian mixture with equal component weight as in Gulrajani et al. [41]. We use an empirical sample set X , consisting of |X | = 5, 000 samples from both 1D and 2D cases, and illustrate in Fig. 2 the KDE of 5000 generated samples yj = Gθ( j) after 5000 training epochs. For the 1D case, we take 200 grids in [−10, 10] to approximate the empirical distribution of p̂X and p̂Y , and report the corresponding forward KL (KL[p̂X ||p̂Y ]), reverse KL (KL[p̂Y ||p̂X ]), and their difference D(X,Y ) = KL[p̂X ||p̂Y ]− KL[p̂X ||p̂Y ] below each corresponding sub-figure in Fig. 2. Comparing the results of different ρ in Fig. 2, it suggests that minimizing the forward CT cost only encourages the generator to exhibit mode-covering behaviors, while minimizing the backward CT cost only encourages mode-seeking behaviors. Combining both costs provides a user-controllable balance between mode covering and seeking, leading to satisfactory fitting performance, as shown in Columns 2-4. Note that for a fair comparison, we stop the fitting at the same iteration; in practice, we find if training with more iterations, both ρ = 0.75 and ρ = 0.25 can achieve comparable results as ρ = 0.5 in this example. Allowing the mode covering and seeking behaviors to be controlled by adjusting ρ is an attractive property of CTρ.
Resistance to mode collapse: We continue to use a 8-Gaussian mixture to empirically evaluate how well a DGM resists mode collapse. Unlike the data in Fig. 2, where 8 modes are equally weighted, here the mode at the left lower corner is set to have weight γ, while the other modes are set to have the same weight of 1−γ7 . We set X with 5000 samples and the mini-batch size as N = 100. When γ is lowered to 0.05, its corresponding mode is shown to be missed by GAN, WGAN, and SWD-based DGM, while well kept by the CT-based DGM. As an explanation, GANs are known to be susceptible to mode collapse; WGAN and SWD-based DGMs are sensitive to the mini-batch size, as when γ equals to a small value, the samples from this mode will appear in the mini-batches less frequently than those from any other mode, amplifying their missing mode problem. Similarly, when γ is increased to 0.5, the other modes are likely to be missed by the baseline DGMs, while the CT-based DGM does not miss any modes. The resistance of CT to mode dropping can be attributed to its forward component’s mode-covering property. The backward’s mode-seeking property further helps distinguish the density of each mode component to avoid making components of equal weight.
CT for 2D toy data and robustness in adversarial feature extraction: To test CT with more general cases, we further conduct experiments on 4 representative 2D datasets for generative modeling evaluation [41]: 8-Gaussian mixture, Swiss Roll, Half Moons, and 25-Gaussian mixture. We apply the vanilla GAN [9] and Wasserstein GAN with gradient penalty (WGAN-GP) [41] as two representatives of min-max DGMs that require solving a min-max loss. We then apply the generators trained under the sliced Wasserstein distance (SWD) [42] and CT cost as two representatives of min-max-free DGMs. Moreover, we include CT with an adversarial feature encoder trained with (14) to test the robustness of adversary and compare with the baselines in solving the min-max loss.
On each 2D data, we train these DGMs as one would normally do during the first 5k epochs. We then only train the generator and freeze all the other learnable model parameters, which means we freeze the discriminator in GAN, critic in WGAN, the navigator parameter φ of the CT cost, and both
Table 1: FID comparison with different cooperative training on CIFAR-10 (lower FID is preferred).
Critic space FID ↓ Discriminator 29.7
Slicing 32.4 Adversarial CT 22.1
Table 2: FID Comparison with using MMD (Rational quadratic kernel/distance kernel) and CT loss in training critic/generator on CIFAR10 (lower FID is preferred).
MMD-rq Generator loss MMD-dist Generator loss
MMD CT MMD CT Critic MMD 39.9 24.1 Critic MMD 40.3 28.8 loss CT 41.4 23.9 loss CT 30.9 29.4
(φ,η) of CT with an adversarial feature encoder, for another 5k epochs. Figs. 7-10 in Appendix E.1 illustrate this training process on each dataset, where for both min-max baseline DGMs, the models collapse after the first 5k epochs, while the training for SWD remains stable and that for CT continues to improve. Compared to SWD, our method covers all data density modes and moves the generator much closer to the true data density. Notably, for CT with an adversarially trained feature encoder, although it has switched from solving a min-max loss to freezing the feature encoder after 5k epochs, the frozen feature encoder continues to guide the DGM to finish the training in the last 5k epochs, which shows the robustness of the CT cost.
Ablation of cooperatively-trained and adversarially-trained CT: As previous experiments show the adversarially-trained feature encoder could provide a valid feature space for CT cost, we further study the performance of the encoders cooperatively trained with other losses. Here we leverage, as two alternatives, the space of an encoder trained with the discriminator loss in GANs and the empirical Wasserstein distance in sliced 1D spaces [43]. We test these settings on both 8-Gaussian, as shown Fig. 4, and CIFAR-10 data, as shown in Table 1. It is confirmed these encoders are able to cooperatively work with CT, in general producing less appealing results with those trained by maximizing CT. From this view, although CT is able to provide guidance for the generators in the feature space learned with various options, maximizing CT is still preferred to ensure the efficiency. Moreover, as observed in Figs. 4b-4e, CT clearly improves the fitting with sliced Wasserstein distance. To explain why CT helps improve in the sliced space, we further provide a toy example in 1D to study the properties of CT and empirical Wasserstein distance in Appendix E.3.
Ablation of MMD and CT: As MMD also compares the pair-wise sample relations in a mini-batch, we study if MMD and CT can benefit each other. The feature space of MMD-GAN can be considered as Tη◦k, where k is the rational quadratic or distance kernel in Bińkowski et al. [10]. Here we evaluate the combinations of MMD/CT as the generator/encoder criterion to train DGMs. On CIFAR-10, shown in Table 2, combining MMD and CT generally has improvement over MMD alone in FID. It is interesting to notice that for MMD-GAN, learning its generator with the CT cost shows more obvious improvement than learning its feature encoder with the CT cost. We speculate the estimation of MMD relies on a supremum of its witness function, which needs to be maximized w.r.t Tη ◦ k and cannot be guaranteed by maximizing CT w.r.t Tη. In the case of MMD-dist, using CT for witness function updates shows a more clear improvement, probably because CT has a similar form as MMD when using the distance kernel. From this view, CT and MMD are naturally able to be combined to compare the distributional difference with pair-wise sample relations. Different from MMD, CT does not involve the choice of kernel and its navigators assist to improve the comparison efficiency. Below we show on more image datasets, CT is compatible with many existing models, and achieve good results to show improvements on a variety of data with different scale.
Adversarially-trained CT for natural images: We conduct a variety of experiments on natural images to evaluate the performance and reveal the properties of DGMs optimized under the CT cost. We consider three widely-used image datasets, including CIFAR-10 [44], CelebA [45], and
LSUN-bedroom [46] for general evaluation, as well as CelebA-HQ [47], FFHQ [48] for evaluation in high-resolution. We compare the results of DGMs optimized with the CT cost against DGMs trained with their original criterion including DCGAN [49], Sliced Wasserstein Generative model (SWG) [42], MMD-GAN [10], SNGAN [50], and StyleGAN2 [51]. For fair comparison, we leverage the best configurations reported in their corresponding paper or Github page. The detailed setups can be found in Appendix D. For evaluation metric, we consider the commonly used Fréchet inception distance (FID, lower is preferred) [52] on all datasets and Inception Score (IS, higher is preferred) [53] on CIFAR-10. Both FID and IS are calculated using a pre-trained inception model [54].
The summary of FID and IS on previously mentioned model is reported in Table 3. We observe that trained with CT cost, all the models have improvements with different margin in most cases, suggesting that CT is compatible with standard GANs, SWG, MMD-GANs, WGANs and generally helps improve generation quality, especially for data with richer modalities like CIFAR-10. CT is also compatible with advanced model architecture like StyleGAN2, confirming that a better feature space could make CT more efficient to guide the generator and produce better results.
The qualitative results shown in Fig. 5 are consistent with quantitative results in Table 3. To additionally show how CT works for more complex generation tasks, we show in Fig. 6 example higher-resolution images generated by CT-SNGAN on LSUN bedroom (128x128) and CelebA-HQ (256x256), as well as images generated by CT-StyleGAN2 on LSUN bedroom (256x256), FFHQ (256x256), and FFHQ (1024x1024).
On the choice of ρ for natural images: In previous experiments, we fix ρ = 0.5 by default when we prefer neither mode-covering nor mode-seeking. We further tune ρ as an additional ablation study on CIFAR-10 dataset with both the CT + DCGAN backbone and CT + SNGAN backbone to see its affects in terms of certain metrics, such as the FID score. The results
shown in Table 4 suggest that CT is not sensitive to the choice of ρ as long as 0 < ρ < 1, and the FID score could be further improved if we choose a smaller ρ to bias towards mode-seeking.

heading: 5 Conclusion
text: We propose conditional transport (CT) as a new criterion to quantify the difference between two probability distributions, via the use of both forward and backward conditional distributions. The forward and backward expected cost are respectively with respect to a source-dependent and targetdependent conditional distribution defined via Bayes’ theorem. The CT cost can be approximated with discrete samples and optimized with existing stochastic gradient descent-based methods. Moreover, the forward and backward CT possess mode-covering and mode-seeking properties, respectively. By combining them, CT nicely incorporates and balances these two properties, showing robustness in resisting mode collapse. On complex and high-dimensional data, CT is able to be calculated and stably guide the generative models in a valid feature space, which can be learned by adversarially maximizing CT or cooperatively deploying existing methods. On various benchmark datasets for deep generative modeling, we successfully train advanced models with CT. Our results consistently show improvement over the original ones, justifying the effectiveness of the proposed CT loss.
Discussion: Note CT brings consistent improvement to these DGMs by neither improving their network architectures nor gradient regularization. Thus it has great potential to work in conjunction with other state-of-the-art architectures or methods, such as BigGAN [55], self-attention GANs [56], partition-guided GANs [57], multimodal-DGMs [58], BigBiGAN [59], self-supervised learning [60], and data augmentation [61–63], which we leave for future study. As the paper is primarily focused on constructing and validating a new approach to quantify the difference between two probability distributions, we have focused on demonstrating the efficacy and interesting properties of the proposed CT on toy data and benchmark image data. We have focused on these previously mentioned models as the representatives in GAN, MMD-GAN, WGAN under CT, and we leave to future work using the CT to optimize more choices of DGMs, such as VAE-based models [8] and neural-SDE [64].

heading: Acknowledgments
text: The authors acknowledge the support of NSF IIS-1812699, the APX 2019 project sponsored by the Office of the Vice President for Research at The University of Texas at Austin, the support of a gift fund from ByteDance Inc., and the Texas Advanced Computing Center (TACC) for providing HPC resources that have contributed to the research results reported within this paper.

Review summaries:
This paper presents a new cost function to measure the difference between two probability distributions. The proposed cost function called the conditional transport (CT) is a mixture of two cost functions. The two cost functions of the mixture are derived by factorizing the joint probability density function in two reverse manners. A sample analogue of the CT is given, and it is shown that this sample analogue converges to its population version as the number of samples tends to infinity. Experiments are given to assess the performance of the CT and compare the performance of CT-based methods with that of existing methods.

This paper proposes a new distance measure between two probability distributions by exploiting the chain rule and Bayes' theorem. The new distance measure is obtained by minimizing out an additional parameter from the conditional transport (CT) cost. The CT costs is a convex combination of a forward component and a backward one. The forward CT is the expected cost of moving a source data point to a target one, with their joint distribution defined by the product of the source probability density function (PDF) and a surrogate source-dependent conditional distribution, which is related to the target PDF via Bayes' theorem. The backward CT is defined by reversing the direction. The CT cost can be approximated with discrete samples and optimized with existing gradient descent-based methods. The new measure nicely balances the mode-covering property of forward CT and mode-seeking property of backward CT, showing strong robustness against mode collapse. Experiments on a wide variety of benchmark datasets show that substituting the default statistical distance of an existing generative adversarial network with CT consistently improves the performance.

The paper defines conditional transport (CT), a new divergence between probability distributions (Equation 5), and reports experimental findings where this divergence serve as a good loss function for training generative models over natural images (Table 3, Figure 4 and so on). Among the tunable parameters of the divergence, there is one parameter that modulates between the extremes of mode seeking and mode covering. In real-data experiments, generative models fitted using CT as training loss have low Frechet inception distance (FID) and high inception score, which overall indicate high-quality generated images.