Exploiting Chain Rule and Bayes' Theorem to Compare Probability Distributions

1 Introduction . 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 . 2 Chain rule and Bayes ’ theorem based conditional transport . 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 .

Summary: 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.