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

Abstract:
We study the gradient flow for a relaxed approximation to the Kullback-Leibler (KL) divergence between a moving source and a fixed target distribution. This approximation, termed the KALE (KL Approximate Lower bound Estimator), solves a regularized version of the Fenchel dual problem defining the KL over a restricted class of functions. When using a Reproducing Kernel Hilbert Space (RKHS) to define the function class, we show that the KALE continuously interpolates between the KL and the Maximum Mean Discrepancy (MMD). Like the MMD and other Integral Probability Metrics, the KALE remains well-defined for mutually singular distributions. Nonetheless, the KALE inherits from the limiting KL a greater sensitivity to mismatch in the support of the distributions, compared with the MMD. These two properties make the KALE gradient flow particularly well suited when the target distribution is supported on a low-dimensional manifold. Under an assumption of sufficient smoothness of the trajectories, we show the global convergence of the KALE flow. We propose a particle implementation of the flow given initial samples from the source and the target distribution, which we use to empirically confirm the KALE’s properties.

heading: 1 Introduction
text: 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 cannot 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 cannot 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.

heading: 2 Interpolating between KL and MMD using KALE
text: 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 , ...).

heading: 2.1 The KL Approximate Lower bound Estimator (KALE)
text: 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.

heading: 3 KALE Gradient Flow
text: Having introduced KALE as a relaxed approximation of the KL, we now construct the KALE gradient flow, and assert its well-posedness. We provide conditions for global convergence of the flow, and discuss its relationship with the MMD flow and the KL flow. All proofs are given in the appendix.

heading: 3.1 Wasserstein Gradient Flow of the KALE
text: Wasserstein Gradient Flows of divergence functionals F(P || Q) aim at transporting mass from an initial probability distribution P0 towards a target distribution Q by following a path Pt in probability space. The path is required to dissipate energy, meaning that t 7! F(Pt || Q) is a decreasing function of time. Additionally, it is constrained to satisfy a continuity equation that allows only local movements of mass without jumping from a location to another. This equation involves a time dependent vector field Vt which serves as a force that drives the movement of mass at any time t:
@tPt + div(PtVt) = 0. (8)
Eq. (8) holds in the sense of distributions, meaning that for any test function ' 2 C1 c (Rd⇥(0,+1)), we have: Z @t'(x, t)dPtdt+ Z hrx'(x, t), VtiRd dPtdt = 0.
The Wasserstein gradient flow of a (sufficiently regular) functional F is then obtained by choosing Vt as the gradient of first variation of F , defined as the Gâteaux derivative of P along the direction ,
DPF(P; ) = lim ✏!0 ✏ 1 (F(P+ ✏ ) F(P)) ,
where R d = 0, and provided that such a limit exists. This choice recovers a particle Euclidean gradient flow when P0 is a finite sum of Dirac distributions, and can thus be seen as a natural extension of gradient flows to the space of probability distributions [3, 63, 64]. In the next proposition, we show that the functional P 7 ! KALE(P || Q) admits a well-defined gradient flow of this form.
Proposition 2 (KALE Gradient Flow). Let > 0, and P0,Q 2 P2(Rd). Under Assumptions 1 and 2, the Cauchy problem
@tPt div(Pt(1 + )rh?t ) = 0, Pt=0 = P0, (9)
where h? t is the unique solution of
h ? t = argmax
h2H
⇢ 1 + Z hdPt Z e hdQ
2 khk2 , (10)
admits a unique solution (Pt)t 0, which is the Wasserstein Gradient Flow of the KALE.
The proof, given in Appendix D, requires studying the regularity properties of KALE in the metric space P2(Rd) endowed with the Wasserstein-2 Distance [64], using the framework developed in [3].

heading: 3.2 Convergence properties of the KALE flow
text: Proposition 2 hints at a connection between the KALE flow and the MMD flow, which solves:
@tPt div (PtrfPt,Q) = 0, Pt=0 = P0 (11)
The MMD flow and the KALE flow thus differ in the choice of witness function characterizing their velocity field. A convergence analysis of the MMD flow was proposed for a wide range of kernels in [4] using inequalities of Lojasiewicz type; in particular, the MMD flow is guaranteed to converge provided that the quantity Pt Q remains bounded in the negative Sobolev distance kPt QkḢ 1(Pt) [50]. We recall that the negative weighted negative Sobolev distance [4] between µ and ⌫ is defined as:
kµ ⌫k Ḣ 1(P) = sup
kfk Ḣ(P)1
Z
fd(µ ⌫) ,
which is obtained by duality with the weighted Sobolev semi-norm kfk Ḣ(P) = (
R krfk2dP) 12 . Note
the important role of the latter quantity in the energy dissipation formula of the KALE gradient flow:
dKALE(Pt || Q) dt
= Z
(1 + )2 krh?k2 dP = (1 + )2 kh?k2 Ḣ(P) . (12)
In the next proposition, we extend the condition ensuring the global convergence of the MMD flow [4] to the KALE flow:
Proposition 3. Under Assumptions 1 and 2, if kPt QkḢ 1(Pt)  C for some C > 0, then:
KALE(Pt || Q)  C
CKALE(P0 || Q) + t .
Proposition 3 ensures a convergence rate in O(1/t) provided that kPt QkḢ 1(Pt) remains bounded. This convergence rate is slower than the linear rate of the KL along its gradient flow [38] and could be an effect of RKHS smoothing.

heading: 4 KALE Particle Descent
text: We now derive a practical algorithm that computes the solution of a KALE gradient flow, given an initial source-target pair P0 and Q. Because of the continuous-time dynamics, and the possibly continuous nature of P0 and Q, solutions of Eq. (9) are intractable to compute and manipulate. To address this issue, we first introduce the KALE Particle Descent Algorithm that returns a sequence (bPN
n )n 0 of discrete probability measures able to approximate the forward Euler discretization of Pt with arbitrary precision. Additionally, we show that the KALE particle descent algorithm can be regularized using noise injection [4], which guarantees global convergence of the flow under a suitable noise schedule. All proofs are given in the appendix.

heading: 4.1 The KALE Particle Descent Algorithm
text: Time-discretized KALE Gradient Flow As a first step towards deriving the KALE particle descent algorithm, let us first consider a time-discretized version of the KALE gradient flow (Eq. (9) and Eq. (10)), obtained by applying a forward-Euler scheme to Eq. (9) with step size . This time-discretized equation is given by
Pn+1 = (I (1 + )rh?n)#Pn, Pn=0 = P0. (13) The function h?
n is a discrete time analogue of Eq. (10), in that it is solution to the following
optimization problem:
h ? n = argmax
h2H
⇢ 1 + Z hdPn Z e hdQ
2 khk2 . (14)
The solution Pn of Eq. (13) is a sensible approximation of Pt: indeed, it can be shown under suitable smoothness assumptions [54, 4] that the piecewise-constant trajectory (t 7 ! Pn if t 2 [n , (n+ 1) )) obtained from the time-discretized gradient flow of a functional F will recover the true gradient flow solution Pt of F as ! 0.
Approximation using finitely many samples: the KALE particle descent algorithm Despite its discrete-time nature, the sequence (Pn)n 0 may still be intractable to compute: for generic P0 and Q, Eq. (14) will contain intractable expectations and have an infinite dimensional search space. To address this issue, we propose the KALE particle descent algorithm: this algorithm approximates the true time-discrete iterates Pn given N samples {X(i)}Ni=1 and {Y (i) 0 }Ni=1 of Q and P0, by computing the probabilities bPN n
solving the time-discrete KALE gradient flow arising from the empirical sourcetarget pair bQN = 1
N
P N
i=1 X (i) and bPN0 = 1N
P N
i=1 Y (i) 0 . As opposed to Pn, it is possible to exactly
compute bPN n : indeed, the recursion equation Eq. (13) implies that bPN n
remains discrete for all n. More precisely, we have bPN
n = 1
N
P N
i=1 Y (i) n , where
Y (i) n+1 = Y (i) n (1 + )rbh? n (Y (i) n ), (15)
and bh? is defined as
bh? n = argmax
h2H
⇢Z hdbPN
n
Z e h dbQN
2 khk2 H . (16)
As in the sample-based setting of Eq. (7), bPN n and bQN are discrete, meaning that Eq. (16) reduces to an N - dimensional convex problem, and bh?
n can be tractably computed. The alternate execution of
Eq. (15) and Eq. (16) for a finite number of time steps defines the KALE Particle Descent Algorithm, that we lay out in Algorithm 1.
Consistency of the KALE Particle Descent Algorithm Note that the source of error in the KALE particle descent algorithm lies in the use of an approximate witness function bh?
n instead of the true, but
intractable, h? n . Indeed, one can show, using the theory of McKean-Vlasov representative processes [39], that the n-th iterates of the sequence defined by:
Ȳ (i) n+1 = Ȳ (i) n (1 + )rh? n (Ȳ (i) n ), Ȳ (i)0 ⇠ P0, 1  i  N (17)
are distributed according to the nth iterate Pn of the true discrete-time KALE gradient flow solution defined in Eq. (13). As such, the discrete probability P̄N
n = 1
N
P N
i=1 Ȳ (i)n may be considered as
an unbiased space-discretization of Eq. (13). In the next proposition, we show that the iterates bPN n returned by the KALE particle descent algorithm can approximate the unbiased P̄N n
with arbitrarily low error. Proposition 4 (Consistency of the KALE particle descent). Let {Y (i)0 }Ni=1 ⇠ P0. Let (P̄Nn )n 0 be the sequence of discrete probabilities arising from Eq. (17) with initial conditions {Y (i)0 }Ni=1, and let (bPN
n )n 0 be the sequence arising from Eq. (15) with the same initial conditions {Y (i)0 }Ni=1. Let
nmax 0. Then, under Assumptions 1 and 2, for all n  nmax , the following bound holds:
EW2(bPNn , P̄Nn )  A
B p N (e Bnmax 1)
with A = q 2KK1d(1 + e 8K ) ⇥ 1
4 p KK1d+K2d
B = (1+ )(4 p KK1d+ p K2d) , and K,K1d,K2d are
the constants defined in Assumptions 1 and 2.
Proposition 4 shows that given a finite time horizon nmax , and given sufficiently many samples of P0 and Q, one can approximate an exact discrete KALE flow between n = 0 and n = nmax with arbitrary precision. The proof of Proposition 4 (given in Appendix F) relies on the regularity of the KALE witness function x 7 ! bh?
n (x), but also on the regularity of the mapping bPN n 7 ! bh? n (using
the 2-Wasserstein distance as the metric on P2(Rd).
Algorithm 1 KALE Particle Descent Algorithm
Input: {Y (i)0 }Ni=1 ⇠ P0, {X(i)}Ni=1 ⇠ Q, max_iter, , k, Output {Y (i)max_iter}Ni=1 for n = 0 to max_iter 1 do f_star dual_solve(X(1), Y (1)
i , . . . , X
(N) , Y (N) i
, k, ) # See Eq.6 h_star compute_log_ratio(f_star, X(1), Y (1)
i , . . . , X
(N) , Y (N) i
, k, ) # Ditto for j = 1 to N do v (1 + )grad(h_star(Y (j)
i ))
Y (j) i+1 Y (j) i ⇥v
end for end for

heading: 4.2 Regularization of KALE particle descent using Noise Injection
text: In practice, guaranteeing the convergence of the KALE gradient flow (and its corresponding KALE particle descent) by relying on the condition given in Proposition 3 is cumbersome for two reasons: first, this condition is hard to check, and second, it does not tell us what to do when the condition is not met. Noise injection [4, 11] is a practical regularization technique originally introduced for the MMD flow, that trades off some of the “steepest descent” property of gradient flow trajectories with some additional smoothness (in negative Sobolev norm) in order to improve convergence to the target trajectory. We recall that the solution of a (discrete time) noise injected gradient flow with velocity field (1 + )rh?
n and noise schedule n is defined as the sequence (Pn)n 0 whose iterates verify: Pn+1 = ((x, u) 7 ! x (1 + )rh?n(x+ nu))# (Pn ⌦ g), (18)
where g is a standard unit Gaussian distribution. As we show in the next proposition, under a suitable noise schedule, noise injection can also be applied to ensure global convergence of the KALE flow. Proposition 5 (Global Convergence under noise injection dynamics). Let Pn be defined as Eq. (18). Let ( n)n 0 be a sequence of noise levels, and define D n,Pn = Ey⇠Pn,u⇠g krh?n(x+ nu)k
2 with g the density of a standard Gaussian distribution. Then, under Assumptions 1 and 2, and for a choice of n such that:
8K2d 2n 2 KALE(Pn || Q)  D n,Pn(Pn),
the following holds: KALE(Pn+1 || Q) KALE(Pn || Q)  2 (1 3 p KK2d)D n,Pn(Pn). Moreover, if 1P i=1 i = +1, then limn!1 KALE(Pn || Q) = 0.
As in [4], convergence of the regularized KALE flow is guaranteed when the noise schedule satisfies an inequality for all n, which is hard to check in practice. Nonetheless, we empirically observe that in all our problems a small, constant noise schedule can help the KALE flow reach a lower KALE value at convergence.
Let us stress that the noise injection scheme given in Proposition 5 is a population scheme that includes an intractable convolution. To use noise injection in the KALE particle descent algorithm, we approximate this convolution using a single sample U (i)n for each particle update. Eq. (15) becomes:
Y (i) n+1 = Y (i) n (1 + )rbh? n (Y (i) n + nU (i) n ), U (i) n ⇠ N (0, 1). (19)
Implementation The particle descent algorithm can be implemented using automatic differentiation software such as the pytorch library in python. This allows us to easily compute the gradient of the log-density ratio estimate bh?
n appearing in the particle update rule Eq. (15).
Computing bh? n
can be achieved using methods such as gradient descent, coordinate descent or higher order optimization methods such as Newton’s method and L-BFGS [35].

heading: 5 Experiments
text: In this section, we empirically study the behavior of the KALE particle descent algorithm in three settings reflecting different topological properties for the source-target pair: a pair with a target supported on a hypersurface (zero volume support), a pair with disjoint supports of positive volume, and a pair of distributions with a positive density supported on Rd. KALE flow for targets defined on hypersurfaces Our first example consists in a target supported (and uniformly distributed) on a lower-dimensional surface that defines three non-overlapping rings. The initial source is a Gaussian distribution with a mean in the vicinity of the target Q. This setting is a perfect candidate to illustrate the failure modes of both the KL and the MMD when used in particle descent algorithms: on the one hand, the measures P0 and Q are mutually singular, and thus the KL gradient flow from P0 to Q does not exist. By contrast, the KALE is well-defined in this case, and inherits from the KL an increased sensitivity to support discrepancy. For that reason, we hypothesize that the trajectories of the KALE flow will converge towards a better limit compared to its MMD flow counterpart. We sample N = 300 points from the target and the initial source distribution and run an implementation of Algorithm 1 for n = 50000 iterations. The complete set of parameters is given in the appendix. Results are plotted in Fig. 1. We indeed notice that the KALE flow trajectory remains
close to the target support and recovers the target almost perfectly. This illustrates the ability of the KALE flow to relax the hard support-sharing constraints of the KL flow into soft support closeness constraints. These soft constraints are not present in the MMD flow, where particles of the source can remain scattered around the plane.
KALE flow between probabilities with disjoint support In our second example, we consider a source/target pair that are supported on disjoint subsets each with a finite, positive volume (unlike the previous example). The support of the source and the target consist respectively of a heart and a spiral, and the two distributions have a uniform density on their support. Again, because the supports of the source and the target are disjoint, the KL flow cannot be defined, nor simulated for this pair. We run a KALE particle descent algorithm, and compare it as before with an MMD flow, as well as with a “Sinkhorn descent algorithm” [25]. Results are in Fig. 2.
As we can see, the soft support-sharing constraint informing the KALE flow forces the source to quickly recover the spiral shape, much before the Sinkhorn and MMD flow trajectories. However, compared to Sinkhorn, the two KALE-generated spirals have a harder time recovering outliers, disconnected from the main support of the spiral.
KALE flow for probabilities with densities We consider the setting where the target admits a positive density on Rd. Hence, unlike in the two previous examples, the KL gradient flow is well-defined, and can be simulated using the Unadjusted Langevin Algorithm (ULA). Echoing the interpolation property of the KALE between the MMD and the KL shown in Proposition 1, we propose to investigate whether this property is preserved in a gradient flow setting. We consider a
balanced mixture of 4 Gaussians with means located on the 4 corners of the unit square for the target and a source distribution given by a unit Gaussian in the vicinity of the unit square. We then run KL, MMD, and KALE flows with different values of , and compute the Wasserstein distance between reference particles at iteration n from either the MMD or KL flow and particles obtained from the KALE flow at the same iteration n. The choice of the Wasserstein distance is natural for Wasserstein Gradient Flows. As shown in Fig. 3a, for “small” values of , particles from a KALE flow remain close to the ULA particles, while for “large” ones they remain close to the MMD particles (Fig. 3b).
Impact of noise injection On all three examples, using a regularized KALE flow with an appropriately tuned n schedule always improves the proximity to the global minimum P1 = Q. Its effect is particularly impactful in the mixture of Gaussians example, where a small, constant noise schedule n allows for faster mixing times for Pn, as opposed to its unregularized counterpart, see Fig. 3c. We provide further details on the impact of noise injection in the appendix.

heading: 6 Discussion and further work
text: We have constructed the KALE flow, a gradient flow between probability distributions that relaxes the KL gradient flow for probabilities with disjoint support. Using the KALE Particle Descent Algorithm, we have shown on several examples that in cases where a KL gradient flow cannot be defined, trajectories of the KALE flow empirically exhibit better convergence properties when compared to the MMD flow, a flow that the KALE is also able to interpolate. In cases where the KL flow can be defined, we notice empirically that the KALE flow can approximate the trajectories of the KL flow, but using only information from samples of the target. This latter property is in sharp contrast with KL Gradient Flow discretizations like the Unadjusted Langevin Algorithm: in this regard, we could use the KALE flow as a sample-based approximation of the KL flow, which is to our knowledge a novel concept. Future work would analyze when the KALE flow is a consistent estimator of the KL flow in the large sample limit.

Review summaries:
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). 

This paper studies the Wasserstein gradient flow of the KALE, a lower bound estimator of the KL divergence. By time discretization, this flow leads to an algorithm to sample from a target distribution for which samples are known (in contrast to the setting where the log density is known). Moreover, this KALE flow performs well in the case when the initial and the target distribution have disjoint supports, a setting in which the KL flow is not defined. 

This paper proposes a new probability divergence called KALE (KL approximate lower-bound estimator) that interpolates between the KL divergence and MMD through a scaling parameter. The motivation comes from regularizing the dual formulation of KL, and KALE has been shown to exhibit desirable topological properties. In addition, the Wasserstein gradient flow of KALE is derived and a particle-based descent algorithm is introduced to compute the KALE flow in practice. Consistency and global convergence under noise injection of the proposed algorithm have been shown. Experiments on toy examples have been conducted to demonstrate empirically that the KALE flow indeed interpolates between those of KL and MMD, while having a "non-flat" geometry (compared to MMD) and requiring only sample access to compute (compared to KL).