KALE Flow: A Relaxed KL Gradient Flow for Probabilities with Disjoint Support

1 Introduction . We consider the problem of transporting probability mass from a source distribution P to a target distribution Q using a Wasserstein gradient flow in probability space . When the density of the target is well-defined and available , the Wasserstein gradient flow of the Kullback-Leibler ( KL ) divergence provides a simple way to transport mass towards the target through the Fokker-Planck equation as established in the seminal work of [ 31 ] . Its time discretization yields a practical algorithm , the Unadjusted Langevin Algorithm ( ULA ) , which comes with strong convergence guarantees [ 22 , 19 ] . A more recent gradient flow approach , Stein Variational Gradient Descent ( SVGD ) [ 36 ] , also leverages the analytic expression of the density and constructs a gradient flow of the KL , albeit using a metric different from the Wasserstein metric . The KL divergence is of particular interest due to its information theoretical interpretation [ 56 ] and its use in Bayesian Inference [ 13 ] . The KL defines a strong notion of convergence between probability distributions , and as such is often widely used for learning generative models , through Maximum Likelihood Estimation [ 20 ] . Using the KL as a loss requires knowledge of the density of the target , however ; moreover , this loss is well-defined only when the distributions share the same support . Consequently , we can not use the KL in settings where the probability distributions are mutually singular , or when they are only accessible through samples . In particular , the Wasserstein gradient flow of the KL in these settings is ill-defined . ⇤Work mostly completed at the Gatsby Unit . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Recent works have considered the gradient flow of Integral Probability Metrics ( IPM ) [ 46 ] instead of the KL , in settings where only samples ( and not the density ) of the target are known . This includes the Maximum Mean Discrepancy ( MMD ) [ 4 ] and the Kernelized Sobolev Discrepancy ( KSD ) [ 45 , 44 ] . One motivation for considering these particle flows is their connection with the training of Generative Adversarial Networks ( GANs ) [ 28 ] using IPMs such as the Wasserstein distance [ 7 , 29 , 27 ] , the MMD [ 24 , 34 , 33 , 9 , 10 , 5 ] or the Sobolev discrepancy [ 43 ] . As discussed in [ 45 , Section 3.3 ] , these flows define update equations that are similar to those of a generator in a GAN . Thus , studying the convergence flows can provide helpful insight into conditions for GAN convergence , and ultimately , improvements to GAN training algorithms . A second motivation lies in the connection between the training dynamics of infinitely wide 2-layer neural networks and the Wasserstein gradient flow of particular functionals [ 52 ] . Thus , analyzing the asymptotic behavior of such flows [ 40 , 58 , 18 ] can ultimately provide convergence guarantees for the training dynamics of neural networks . Establishing such results remains challenging for some classes of IPMs , however , such as the MMD [ 4 ] . In this paper , we construct the gradient flow of a relaxed approximation of the KL , termed the KALE ( KL Approximate Lower bound Estimator ) . Unlike the KL , the KALE is well-defined given any source and target , regardless of their relative absolute continuity . The KALE is obtained by solving a regularized version of the Fenchel dual problem defining the KL , defined over a restricted function class [ 48 , 6 ] , and can be estimated solely from samples from the data . The version of the KALE we consider in this work benefits from two important features that are crucial for defining and analyzing a relaxed gradient flow of the KL . ( 1 ) We define the function class to be a Reproducing Kernel Hilbert Space ( RKHS ) . This makes the optimization problem defining the KALE convex and allows for practical algorithms computing it . ( 2 ) We consider a regularized version of the problem defining the KALE , thus providing a simpler expression for the gradient flow by virtue of the envelope theorem [ 42 ] . In Section 2 , we review the KALE , and show that it is a divergence that metrizes the weak convergence of probability measures , while interpolating between the KL and the MMD depending on the amount of regularization . We then construct in Section 3 the Wasserstein Gradient Flow of the KALE , and we show global convergence of the KALE flow provided that the trajectories are sufficiently regular . In Section 4 , we introduce the KALE particle descent algorithm as well as a practical way to implement it . In Section 5 , we present the results obtained by running the KALE particle descent algorithm on a set of problems with different geometrical properties . We show empirically that the sensitivity to support mismatch of the KALE inherited from the KL leads to well-behaved trajectories compared to the MMD flow , making the KALE flow a desirable alternative when a KL flow can not be defined . Related work . The Fenchel dual formulation of the KL , and more generally f -divergences , has a rich history in Machine Learning : [ 48 ] relied on this dual formulation to estimate the KL between two probability distributions when their density ratios belong to an approximating class . They derived a plug-in estimator for the KL which comes with convergence guarantees . In the context of GANs , [ 49 ] used the Fenchel dual representation of f -divergences , of which the KL is a particular instance , as a GAN critic . Later , [ 41 ] used Fenchel duality to estimate the KL in the context of Variational Inference ( VI ) when the variational distribution is chosen to be an implicit model , thus allowing more flexible models at the expense of tractability of a KL term appearing in the expression of the ELBO . In both the GAN and VI settings , the function class defining the f -divergence was restricted to neural networks . Recently , [ 6 ] showed that controlling the smoothness of such a function class results in a divergence , the KL Approximate Lower bound Estimator ( KALE ) , that metrizes the weak convergence of distributions [ 21 ] , unlike the KL which defines a stronger topology [ 62 ] . The KALE is therefore well-suited for learning Implicit Generative Models which are only accessible through sampling , as advocated in [ 8 ] . When neural network classes are used , however , the method has no optimization guarantees , as the dual problem becomes non-convex due to the choice of the function class . This is unlike our setting [ and that of 48 ] , since our dual problem is strongly convex and comes with guarantees . [ 12 ] presented a general framework incorporating both f -divergences and Integral Probability Metrics , but rely on another variational formulation than the Fenchel dual one . In parallel to work related to f -divergences , [ 30 , 16 , 53 , 1 ] have investigated the task of sampling in the case where the source and the target have disjoint supports . Again , unlike our setting , these works assume that the log-density of the target distribution is known . 2 Interpolating between KL and MMD using KALE . In this section , we introduce the KALE , a relaxed approximation of the KL divergence . Although we will use the KALE to define a relaxed KL gradient flow , we show in this section that the KALE is an object of independent interest outside the gradient flow setting : indeed , it is a valid probability divergence that metrizes the weak convergence of probability distributions , and interpolates between the KL and the Maximum Mean Discrepancy . Mathematical details and notation We start by introducing some notation . We denote by P ( Rd ) the set of probability measures defined on Rd endowed with its Borelian -algebra , and by P2 ( Rd ) ⇢ P ( Rd ) the set of elements of P ( Rd ) with finite second moment . Weak convergence of a sequence of probability measures ( Pn ) n 0 towards P is written Pn * P. A positive definite kernel on the set Rd will be denoted k : Rd⇥Rd 7 ! R , with RKHS H , and we will use kx to refer to the RKHS function y 7 ! k ( x , y ) obtained by fixing the first argument of k. The Dirac delta measure for x 2 Rd will be written x . We denote by C1c ( Rd ⇥ ( 0 , +1 ) ) the set of infinitely differentiable functions with compact support on Rd ⇥ ( 0 , +1 ) , and by C0 b ( Rd ) the set of continuous bounded functions from Rd to R. Sets of N points in Rd will be indexed using a superscript { x ( i ) } N i=1 , while a sequence of points in Rd will use a subscript : ( xn ) n2N . If random , elements of such sets X ( i ) or iterates of such sequences Xn will be capitalized . If not , they will be kept in lower-case . For the sake of notational lightness , the choice of the norm used for a specific object ( vectors , functions , operators ) will be specified with a subscript ( e.g . khkH for the RKHS norm ) only if the said choice is not obvious from the context . This remark also holds when referring to the null element of a vector space ( 0H , 0Rd , ... ) . 2.1 The KL Approximate Lower bound Estimator ( KALE ) . The central equation to derive the KALE is the ( Fenchel ) dual formulation of the KL [ 3 , Lemma 9.4.4 ] : KL ( P || Q ) = sup h2C 0 b ( Rd ) ⇢ 1 + Z hdP Z e hdQ . ( 1 ) KALE is obtained from Eq . ( 1 ) by restricting the variational set to an RKHS H with reproducing kernel k , and by adding a penalty to the objective that controls the RKHS norm of the test function h. This regularization ensures that the KALE is well-defined for a broader class of probabilities compared to the KL , even when P and Q are mutually singular . Its complete definition is stated below : Definition 1 ( KALE ) . Let > 0 , and H be an RKHS with kernel k. The Kullback-Leibler Approximate Lower bound Estimator ( KALE ) is given by : KALE ( P || Q ) = ( 1 + ) max h2H ⇢ 1 + Z hdP Z e hdQ 2 khk2 H . ( 2 ) The ( 1+ ) scaling will prevent a degenerate decay to 0 in the large regime ( see Proposition 1 ) . The definition we consider here also differs from the one in [ 6 ] , which first finds the optimal function h ? solving Eq . ( 2 ) , and then defines KALE by evaluating the KL objective in Eq . ( 1 ) , thereby discarding the regularization term when evaluating the divergence . Mathematical Assumptions To prove the theoretical results stated in this work , we will make the following basic assumptions on the kernel k : Assumption 1 ( Boundedness ) . There exists K > 0 such that k ( x , x ) K , for all x 2 Rd . Assumption 2 ( Smoothness ) . The kernel is 2-times differentiable in the sense of [ 60 , Definition 4.35 ] : for all i , j 2 { 1 , . . . , d } @ i @ i+dk and @ i @ j @ i+d @ j+dk exist . Moreover , we have : kr1kxk2 = P d i=1 k @ ikxk2 K1d and kH1kxk2 = P d i , j=1 k @ i @ jkxk2 K2d , where d indicates an expected scaling with dimension . Assumption 1 guarantees the integrability of the objects intervening in KALE , and implies boundedness of the RKHS functions . Assumption 2 guarantees first and second order smoothness of the RKHS functions , a property invoked to control the KALE flow trajectories . Indeed , both the differential and the hessian of any f 2 H can now be bounded in operator norm : using the CauchySchwarz inequality and the kernel reproducing derivative property [ 60 , Corollary 4.36 ] , we have : | @ if ( x ) | k @ ikxk kfk and | @ i @ jf ( x ) | k @ i @ jkxk kfk , implying krf ( x ) k p K1d kfk , and kH ( f ( x ) ) kOp kH ( f ( x ) ) kF p K2d kfk . KALE is a probability divergence We first show that KALE is a probability divergence , and presents topological properties compatible with its use in generative models , such as GANs and Adversarial VAEs : weak continuity , and metrizing the weak convergence of probability distributions . We recall that a functional D ( · || · ) is a probability divergence if both D ( P || Q ) 0 and D ( P || Q ) = 0 ( ) P = Q , for any P , Q 2 P ( Rd ) . Theorem 1 ( Topological properties of KALE ) . Let P , Q 2 P ( Rd ) . Let ( Pn ) n 0 be a sequence of probability measures . Then , under Assumption 1 : ( i ) KALE is weakly continuous : Pn * P = ) lim n ! 1 KALE ( Pn || Q ) = KALE ( P || Q ) ( ii ) If k is universal [ 57 ] , then for any > 0 , KALE is a probability divergence . Moreover , KALE metrizes the weak topology between probability measures with finite first order moments . Central to the proof of all points in this theorem is a link between KALE and the MMD witness function fP , Q , which we report in the next lemma . We recall that given an RKHS H associated to a kernel k , and two probability distributions P and Q , the MMD is defined as the RKHS norm of the difference of mean embeddings of P and Q : MMD ( P || Q ) = kfP , Qk ( fP , Q = Z kxdP ( x ) Z kxdQ ( x ) = µP µQ ) . ( 3 ) Lemma 1 . Let P , Q 2 P ( Rd ) , and K : H 7 ! R be the objective maximized by KALE , e.g . K ( h ) = 1 + R hdP R e hdQ 2 khk 2 . Then , under Assumption 1 , K is Fréchet differentiable . Moreover , the following relationship holds : rK ( 0 ) = fP , Q Intuitively , noting that K ( 0 ) = 0 , Lemma 1 ensures that KALE presents “ equivalent ” regularity and discriminative properties to those of MMD ( a divergence which is itself , under the assumptions of this theorem , weakly continuous and that metrizes the weak convergence of probability distributions ) . The proof of the second point of Theorem 1 is inspired by [ 6 ] , which in turn derives from [ 66 , 37 ] , and is adapted to account for the extra norm penalty term in the version of the KALE in this paper . Interpolating between the MMD and the KL using the KALE The KALE includes a positive regularization parameter , inducing two asymptotic regimes : ! 0 and ! 1 . In these regimes , the KALE asymptotically recovers on the one hand the KL divergence , and on the other hand the MMD . Proposition 1 ( Asymptotic properties of KALE ) . Let P , Q 2 P ( Rd ) . Then , under Assumption 1 , the following result holds : lim ! +1 KALE ( P || Q ) = 1 2 MMD2 ( P || Q ) . ( 4 ) Suppose additionally that log dPdQ 2 H. Then , lim ! 0 KALE ( P || Q ) = KL ( P || Q ) . ( 5 ) Proposition 1 shows that the MMD can be seen as solving a degenerate version of the KL objective . Eq . ( 5 ) is natural given the original definition of the KALE , and highlights the continuity of the KALE objective w.r.t the regularization parameter . Both the MMD and the KL exhibit limitations when used for defining gradient flows , however : as discussed in [ 4 , 25 , 14 ] , the MMD induces a “ flat ” geometry , making its use in generative models tricky [ 5 ] . On the other hand , the KL comes with stronger convergence guarantees [ 3 ] , but its use in sampling algorithms is limited to cases where the target distribution has a density , discarding cases satisfying the widely known manifold hypothesis [ 47 , 14 , 17 ] , stating that typical high dimensional data used in machine learning are distributed on a lower-dimensional manifold . For this reason , we argue that the true interest of the KALE does not lie in its interpolation properties , but rather in the geometry it generates at intermediate values of . The KALE ’ s dual objective Interestingly , the KALE itself admits a dual formulation , with a strong connection to the original KL expression : ( KALE ( P || Q ) / minf > 0 nR ( f ( log f 1 ) + 1 ) dQ+ 12 R f ( x ) kxdQ ( x ) µP 2 H o h ? = R f ? ( x ) kxdQ ( x ) µP ( 6 ) The solution f ? of Eq . ( 6 ) can be seen as an entropically-regularized density ratio estimate on the support of Q ( additional details on the KALE dual objective are given in the appendix ) . Eq . ( 6 ) also yields an elegant estimation procedure , as discussed below . Computing KALE ( P || Q ) in practice As for other IPMs , computing KALE ( P || Q ) for arbitrary P and Q is intractable , and is therefore approximated using a discretization procedure . A common procedure is to assume access to samples { Y ( i ) } N i=1 and { X ( i ) } Ni=1 from P and Q and to solve the empirical equivalent of Eq . ( 6 ) ( e.g . Eq . ( 6 ) , but where P and Q are replaced by their plug-in estimators bPN = 1 N P N i=1 Y ( i ) and bQN = 1N P N i=1 X ( i ) ) . This empirical equivalent is written min f > 0 1 N NX i=1 f ( X ( i ) ) log ( f ( X ( i ) ) ) f ( X ( i ) ) + 1 + 1 2 1 N NX i=1 f ( X ( i ) ) k ( X ( i ) , · ) µbP 2 H ( 7 ) which is a strongly convex N -dimensional problem , and can be solved using standard euclidean optimization methods . By adapting arguments of [ 6 ] , it can be shown that the discrepancy between the KALE ’ s empirical and population value , |KALE ( bPN || bQN ) KALE ( P || Q ) | ( often called “ sample complexity ” ) , is at most O ( 1p N ) . This rate is identical that of Sinkhorn divergences [ 26 ] , another family of entropically-regularized divergences .

Summary: Gradient flows defined on the Wasserstein space have received particular attention from the statistical and machine learning community, since they allow to fit a target distribution starting from another source distribution. They consist in minimizing a certain functional, and the choice of this functional strongly affects its theoretical properties and empirical performance. Recent work have studied gradient flows based on specific functionals, including the maximum mean discrepancy (MMD) and the Kullback-Leibler divergence (KL). This paper focuses on these two cases and proposes a novel gradient flow method based on an approximation of KL, in order to alleviate the limitations induced by MMD and KL.   First, the authors introduce a novel divergence on the space of probability measures, which corresponds to a relaxed version of the KL and is called the "KL approximate lower-bound estimator" (KALE, Definition 1). They study some of its topological properties (Theorem 1), show that KALE interpolates between KL and a specific instance of MMD if the associated kernel is bounded (Proposition 1). Besides, KALE admits a dual formulation (equation (6)), which can be approximated using samples from the two distributions of interest (eq. (7)).   Then, the authors develop the gradient flow in the Wasserstein space based on KALE, by solving a specific instance of the continuity equation in eq. (8): see Proposition 2. The resulting gradient flow is shown to ensure global convergence under certain assumptions (Proposition 3). A practical algorithm is developed to approximate the solution of the KALE gradient flow, and imply discretizing the time of the continuity equation and performing gradient descent on finitely many samples (Algorithm 1). This "particle descent" scheme induce an approximation error (due to the use of samples instead of their underlying densities), which is upper bounded in Proposition 4. The authors also present a variant of this algorithm, largely inspired by prior work on MMD flows, which incorporate Gaussian noise (Section 4.2). The goal of this approach is to relax the convergence guarantees associated to time discretization, which are too strict (Proposition 5). Experiments on different synthetic settings are finally conducted to illustrate the advantages yielded by the KALE gradient flow over the MMD and KL counterparts (Section 5).