KALE Flow: A Relaxed KL Gradient Flow for Probabilities with Disjoint Support
Authors: Pierre Glaser
Michael Arbel
Arthur Gretton
Venue: NeurIPS 2021 Poster (2021)
Keywords: Generative Models
Kernel Methods
Optimal Transportation
Probability Divergences

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.

KALE Flow: A Relaxed KL Gradient Flow for Probabilities with Disjoint Support
Pierre Glaser
Michael Arbel
Gatsby Computational Neuroscience Unit University College London
Université Grenoble Alpes, Inria, CNRS, Grenoble INP, LJK,38000 Grenoble, France" mi michael n. arber lobeeemmmmll... com
pierroglaser@gmal. com
Arthur Gretton
Gatshy Computational Neuroscience Unit University College London
ar arthur . wwwtteongmaill.cooommm com
A Abstract
We study the gradient flow for a relaxed approximation to the Kullback-Leite (KL) divergence between a moving source and fixed target distribution. This approximation, termed the KALE (KL Approximate Lower bound Estimator). solves regularized version of the Fenchel dual problem defining the KI. 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 inter- polates between the KL and the Maximum Mcan Discrepancy (MMD). Like the MMD and other Integral Probability Metries, the KALE remains well-defined for mutually singular distributions. Nonetheless, the KALE inherits from the limiting KL greater sensitivity 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 -iw-eennnional 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.
1 Introduction
We consider the problem of transporting probability mass from source distribution P to a target distribution Q using Wasserstein gradient flow in probability space. When the density of the target is well-defined and available, the Wasserstein gradient flow the Kullback-Leitn provides a simple way to transport mass towards the target through the Fokker-Planck equation established in the seminal work of [31]. Its time discretization yields 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 the density and constructs gradient flow of KL, albeit using a metric different from the Wasserstein metric.
(KL) divergence
The KL divergence of particular interest due information theoretical interpretation [56] and 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 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, cannot the KL in settings where the probability distributions mutually singular, or when they 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 (NeurTPS 2021).

---

Recent works have considered the gradient flow Integral Probability Metrics (IPM) [46] instead of the KI., 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, 271. the MMD [24, 34, 33, 9, 10, S] the Sobolev discrepancy [43]. As discussed in [45, Section 3.3], these flows define update equations that are similar to those of generator 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, construct the gradient flow of 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 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 K.I.EE we consider in this work benefits from two important features that are crucial for defining and analyzing relaxed gradient flow of the KL. (1) We define the function class to be Reproducing Kernel Hilbert Space (RKHS). This makes the optimization problem defining the KALE convex and allows for practical algorithms computing it. (2) We consider 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 the amount of regularization. We then construct in Section 3 the Wasserstein Gradient Flow of the KALE, and show global convergence of the KALE flow provided that the trajectories are sufficiently regular. Section 4, we introduce the KALE particle descent algorithm as well as a practical way to implement it. 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 desirable alternative when KL flow cannot be defined.
Related work. The Fenchel dual formulation of the KL, and more generally f-divergences, has rich history in Machine Learning: [48] relied on this dual formulation 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 (VT) when the variational distribution is chosen to be implicit model, thus allowing more flexible models the expense of tractability of KL term appearing in the expression of the ELBO. In both the GAN and VI settings, the function class defining the -ddivergence 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], J, 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, advocated in [8]. When neural network classes are used, however, the method has optimization guarantees, as the dual problem becomes non-conve convec due 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 onc. In parallel to work related to /-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- of the target distribution known.
2

---

2 Interpolating between KL and MMD using KALE
In this section, we introduce the KALF,, relaxed approximation of the KI. divergence. Although we will use the KALE to define relaxed KL gradient flow, show in this section that the KALE is an object of independent interest outside the gradient flow setting: indeed, it 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 ((R.)) the set of probability measures defined on R" endowed with its Borelian o-algebra, and by P2 C P(R.) the of elements of ((R))) with finite second moment. Weak convergence a sequence of probability measures (P 1200 towards P is written P., P. A positive definite kernel on the set I will be denoted A Rd x R + R. with RKHS H. and we will use k₂ refer the RKHS function y m k(a,g) obtained by fixing the first argument of k. The Dirac delta measure for E Re will be written 6z. We denote by C2 (IRd x (0))) the set of infinitely differentiable functions with compact support on Rd (0, +00), and by CR*) the set of continuous bounded functions from Rd to R. Sets of N points in RRd will be indexed using superscript {
Fi=1 while sequence of points in Hd will use a subscript: (kn))eR: If random, elements of such sets ( or iterates of such sequences X,, will he 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 subscript (e.g. ||А]| 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 (0x, 0ge,.... ...).
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]:
{ 1 (
f
hdP f
KL Q
sup hECf(Rd))
(1)
KALE is obtained from Eq. (1) by restricting the variational set to RKHS 71 with reproducing kernel k and by adding a penalty the objective that controls the RKHS norm of the test function
h. This regularization ensures that the KALE well-defined for a broader class of probabilities compared to the KL, even when I' and Q are mutually singular. Its complete definition is stated helow:
Definition 1 (KALE), Let X > O, and H. be an RKHS with kernel k. The Kullbuck-Leibee Approxi- mate Lower bound Estimator (KALE) is given by:
11A -
{
|+ InhP hndd-- 20000-22
KALE((P || Q) = (1 X) max hEW
(2)
The (1 ++ A) 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, S K. for all E R Assumption 2 (Smoothness).. The kernel is 2-times differentiable in the sense of [60, Definition 4.35|: for all i, E fl,...,,d d) 0,0,+ak and 0,0,0.+a0 exist. Moreover, we have: 1|Vikall Ei ||0,k 1| < Kia and || ,t, 2 Ed (|,,0, & K2d, where d indicates an expected scaling with dimension.
Assumption guarantees the integrability of the objects intervening in KALE, and implies bound- edness of the RKHS functions. Assumption 2 guarantees first and second order smoothness of the RKHS functions, property invoked to control the KALE flow trajectories. Indeed, both the
3

---

differential and the hessian of any f E can now be bounded in operator norm: using the Cauchy- Schwarz incquality and kernel reproducing derivative property [60, Corollary 4.36], we have: [a,((a) & [ and [0,0 v ||,,,,k [[] implying || [V()l & // [|/ and [H))))))) [f())))) 22 Il
KALE is a probability divergence We first show that KALE is 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) 2 0 and D(P II
0 P = Q. for any P, P(
Theorem (Topological properties KALE), P, Q E P(R2). Let (Pr) sequence of probability measures. Then, under Assumption 1:
(i) KALE is weakly continuous: P., P
lim KALE(P, 11 = KALE(P Q) n-+00
(ii) If k is universal [57], then for any À > O, KALE is 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 fF.O. which report in the next lemma. We recall that given an RKHS H. associated to kernel k and two probability distributions IP and Q, the MMD is defined the RKHS norm of the difference of mean embeddings of P and Q:
(M Q) 1
pp pQ)-
(3)
Lemma 1. Ler P. Q E )(IR)) and K. + R. be the objective maximized by KALE, e.g. <C(h) = 1+ f hdP - I dQ - A |lAlie. Then, under Assumption 1, K. is Fréchet differentiable. 2
+
Moreover, the following relationship holds:
VA(O) fe.
Intuitively, noting that KC(0) O. Lemma ensures that KALE presents "cquivalent" regularity and discriminative properties to those of MMD (a divergence which itself, under the assumptions of this theorem, weakly continuous and metrizes the weak convergence probability distributions). The proof of the second point of Theorem 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 positive regularization parameter X, inducing two asymptotic regimes: À - and À - oc. In these regimes, the KALE asymptotically recovers on the one hand the KL divergence, and on the other hand the MMD.
(Asymptotic properties of KALE).. Ler P. Q E P(RR)). Then, under Assumption 1, the following result holds:
Proposition
lim KALE(P
Q 2 MMD (TP Q).
(4)
A
Suppose additionally that log
H. Then,
x
lim KALE(P
Q KL || Q).
(5)
Proposition 1 shows that the MMD can be seen as solving 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 A. 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 use in generative models tricky [5]. On the other hand, the KI. 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 distributed on lower-dimensional manifold. For this reason, argue that the true interest of the KALE does not lie in its interpolation properties, but rather in the geometry generates intermediate values of X.
4

---

The KALE's dual objective Interestingly, the KALE itself admits a dual formulation, with a strong connection the original KI. expression:
KALE(PP () minsso {s (f(log]
1) + Ex II ((=)k, mellic)
(6)
h- ( (a)), -
The solution f Eq. (6) can be seen as an entropically-regulaiin 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 KA Q) in practice As for other IPMs, computing KALE(() Q) for arbitrary P and Q is intractable. and is therefore approximated using discretization procedure. A common procedure is to assume access to samples { L.
Fi=l and (XO) from and Q and solve the empirical equivalent of Eq. (6) (e.g. Eq, (6), but where and Q are replaced by their plug-in estimators I N Ein1 8yt and Q
N Ei (x()). This empirical equivalent is written
1 A 2A N Σ I(X)) - -MB (7)
min /20
Σ (o0")))((((0000) ((x))) +1+
which is strongly convex N 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, IKALE(PN || QM)- KALE(P || Q) (often called "sample complexity",, is at most O (TV ). This rate is identical that of Sinkhorn divergences [26),
another family of entropically-regulazzee divergences.
3 KALE Gradient Flow
Having introduced KALE as relaxed approximation of KL, we now construct the KALE gradient flow; and assert its well pell-poss,ness. We provide conditions for global convergence of flow, and discuss its relationship with the MMD flow and the KL flow. All proof' are given in the appendix.
3.1 Wasserstein Gradient Flow of the KALE
Wasserstein Gradient Flows of divergence functionals F(P Q) aim transporting mass from initial probability distribution Po towards target distribution by following a path P. probability space. The path is required to dissipate energy, meaning that / + F(P, Q) is decreasing function of time. Additionally, it constrained to satisfy a continuity equation that allows only local movements of mass without jumping from location to another. This equation involves time dependent vector field Vi which serves force that drives the movement of mass at any time f:
O,P, + (iv(
(8)
Eq. (8) holds the sense of distributions, meaning that for any test function yp E C. (Rd (0,+)))); we have:
J @ep(a, (dd + (Vos(e,,), (), Vilga d.Pdtt
The Wasserstein gradient flow of (sufficiently regular) funetional F is then obtained by choosing Vi the gradient of first variation of F. defined the Gâteaux derivative of P along the direction X.
Dr DrF((P: (P:x)) 2 lim
-0
where f dx = O, and provided that such limit exists. This choice recovers a particle Euclidean grudient flow when Po is a finite sum of Dirac distributions. and can thus be seen as natural extension of gradient flows to the space of probability distributions [3. 63, 64]. In the next proposition, show that the functional
KALE(PP || Q) admits dell-defined gradient flow of this form.
S

---

Proposition (KALE Gradient Flow), Let O, Po. Q E P2(R)). Under Assumptions and 2, the Caucky problem
a ( = 0, Pr-o Po.
(9)
where Hig the unique solution of
hdP:: - - dd |1m]
mAx REK
1
(10)
admits a unique solution (220 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 (2(RR) endowed with the Wasserstein-2 Distance [64], using the framework developed [3].
3.2 Convergence properties of the KALE flow
Proposition 2 hints connection between the KALE flow and the MMD flow, which solves:
a,Pr div (P, V fe. = 0, I Po
(11)
The MMD flow and the KALEE 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 incqualities of Lojasiewicz type; in particular, the MMD flow is guaranteed to converge provided that the quantity P. Q remains bounded in the negative Sobolev distance PP, Ql (-M()) [50]. recall that the negative weighted negative Sobolev distance [4] between μ v defined as:
|| vl (
sup M
/f
sd(f - -1)p
which is obtained by duality with the weighted Sobolev semi-norm | (J |V /V// Note the important role of the latter quantity in the energy dissipation formula of the KALE gradient flow:
dKALE(P: 11 Q) (1+
-
-(1 x)2 |Milie)
(12)
df
In the next proposition, extend the condition ensuring the global convergence of the MMD flow [4] to the KALE flow:
Proposition 3. Under Assumptions and 2, Il, Q S C. for some > 0, then:
KALE(F', || Q) < CF ++++
Proposition ensures convergence rate in O(1/t) provided that ||P, i 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.
4 KALE Particle Descent
We now derive practical algorithm that computes the solution of KALE gradient flow, given initial source-target pair Fo and Q. Because of continuous-time dynamics, and possibly continuous nature of Po and @ solutions of Eq. (9) are intractable compute and manipulate. To address this issue, we first introduce the KALE Particle Dexcent Algorithm that returns a sequence (PN) haso of discrete probability measures able to approximate the forward Fuler discretization of P, 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 suitable noise schedule. All proofs are given in the appendix.

---

4.1 The KALE Particle Descent Algorithm
Time-diseretized KALE Gradient Flow As 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 7. This time-discretized equation is given by
Fo. The function h; is discrete time analogue of Eq. (10), in that it is solution the following optimization problem:
Pot1 = (1 7(1 + y Pa-0
(13)
h5 arg max
1 +
(14)
The solution PA of Eq. (13) is a sensible approximation of Pe: indeed, it can be shown under suitable smoothness assumptions [54, 4] that the piecewise-constant trajectory (! m P, if e [ny,, (n + 1)r)) obtained from the time-diseretized gradient flow of functional will recover the true gradient flow solution P. of F 7 ++ O.
Approximation using finitely many samples: the KALE particle descent algorithm Despite its discrete-time nature, the sequence (IP 2)220 may still be intractable to compute: for generic Po and Q Eq. (14) will contain intractable expectations and have infinite dimensional search space. To address this issue, we propose the KALE particle descent algoritom; this algorithu approximates the true time-discrete iterates P. given N samples {X(1 L and (Yo ( ) of Q and Po. by computing the probabilities F. solving the time-discrete KALE gradient flow arising from the empirical source- target pair @N = N >m ( and PP = Min Y0 As opposed to P... it is possible to exactly compute F.N. indeed, the recursion equation Eq. (13) implies that foN remains discrete for all n. More precisely, we have T..- = E y-(() where
v = Y 7(1 + A)V
(15)
and A defined as
/ hoTP. / f /CO00""
A 11
= arg max
mag
(16)
NEH
As in the sample-hased setting of Eq. (7), P. and @ are discrete, meaning that Eq. (16) reduces to an N- dimensional convex problem, and R; can be tractably computed. The alternate exccution of Eq. (15) and Eq. (16) for 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
partiele descent algorithm lies in the use of an approximate witness function h; instead of the truc, but intractable, 14- Indeed, one can show, using the theory of McKean-Vlason representative processes [39], that the n-th iterates of the sequence defined by:
( = 7 - 711 + (F D00
Po. ISISN
(17)
are distributed according to the thth iterate Pn, of the true discrete-timg KALE gradient flow solution defined in Eq. (13). As such, the discrete probability F- = N Σi 6pc0 may be considered as an unbiased space-discretization of Eq. (13). In the next proposition, we show that the iterates f.vv returned by the KALE partiele descent algorithm can approximate the unbiased PN with arbitrarily low error.
Proposition 4 (Consistency of the KALE partiele descent). Let (1) y Po. Les )e20 be
the sequence of discrete probabilities arising from Eq. (17) with initial conditions ( Je-, and
(L 1n20 the sequence arising from Eq. (15) with the same initial conditions (Y, N Let
Tonax 2. 0. Then, under Assumptions and 2, for all n S Tmaxx the following bound holds:
A
EW PPN P.N
(6 - 1)
BVN
7

---

with A 2K2 + x 1/KKia-Ka the constants defined in Assumptions 1 and
B (L+A)(4/KK /Kas) and K, Kud Kza are
Proposition 4 shows that given finite time horizon Tenax. and given sufficiently many samples of Po and Q. one can approximate exact discrete KALE flow between 0 and Rmaa with arbitrary precision. The proof of Proposition 4 (given in Appendix F) relies the regularity of the KALE witness function + An (z), but also on the regularity of the mapping P.V - hi (using the Wasserstein distance as the metric on P2(()).
Algorithm KALE Partiele Descent Algorithm
Input: (00))
[==1 Po. {X li=1 Q, max max_iterr X, h,
Output { p
1.N
for n to wtamma__itt max_
f_star + dual_solve y, ...., ,...,,i,,, XiM, y(()) k,
# See Eq. 6
h_star +- rattuttf. staar, x(1,, x(1),,11 y,() ...., X(M, y!M, A,
# Ditto
for j 1 to N do
v (1 + ((grad(h_star(((()))
Y y, yxv
end for end for
4.2 Regularization of KALE particle descent using Noise Injection
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 practical regularization technique originally introduced for the MMD flow, that trades off some of the "stcepest descent" property of gradient flow trajectories with some additional smoothness (in negative Sobolev norm) in order to improve convergence the target trajectory. We recall that the solution of a (discrete time) noise injected gradient flow with velocity field (1 + X) AJV,,, l,, and noise schedule 3, is defined as the sequence (In)noo whose iterates verify:
Fa+1 = ((x, u) ب x 7(1 + A)V
(Ir + B,,u)) # (Pn @ g).
(18)
where 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). Ler P. be defined as Eg. (18). Let (Bu)220 be a sequence of noise levels, and define D3. Ey-FF V ht,( (x + 3,,a)l with the density of a standard Gaussian distribution. Then under Assumptions and 2, and for a choice of In such that:
8K2a82 KALE n ( S (P.....P.).
X2
the following holds: KALE(P
KALE((P, || Q) S -1111-330000
K Do. (PP)).
Moreover, if Σ B +00, then lim KALE Q) = 0.
As in [4]. convergence of the regularized KALE flow is guaranteed when the noise schedule satisties an incquality for all n, which is hard to check in practice. Nonetheless, we empirically observe that in all our problems small, constant noise schedule can help the KALEE flow reach lower KALE value at convergence.
Let us stress that the noise injection scheme given in Proposition 5 is population scheme that includes an intractable convolution. To use noise injection the KALE partiele descent algorithm, we approximate this convolution using single sample U.( for each particle update. Eq. (15) becomes:
y( r-11
Y 7(1 + A)VR, B, U ~N10, 1).
(19)
8

---

Implementation The particle descent algorithm can be implemented using automatic differentiation software such as the pytorch library python. This allows us to easily compute the gradient the log-density ratio estimate i, appearing in the particle update rule Eq. (15).
Computing ht can achieved using methods such as gradient descent, coordinate descent or higher order optimization methods such as Newton's 's method and L-BFGS [35].
5 Experiments
In this section, empirically study the bchavior of the KALE partiele descent algorithm in three settings reflecting different topological properties for the source-targel pair: pair with target supported on hypersurface (zero volume support), a pair with disjoint supports of positive volume, and pair of distributions with a positive density supported on R"
KALE flow for targets defined on hypersurfaces Our first example consists in 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 perfect candidate to illustrate the failure modes of both the KI. and the MMD when used in particle descent algorithms: on the one hand, the measures Po and Q are mutually singular, and thus the KL gradient flow from Po to Q does not exist. By contrast, the KALE is well-defined in this ease, and inherits from the KL inercased sensitivity support discrepancy. For that reason, we hypothesize that the trajectories the KALE flow will converge towards 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 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
T=
T=2
T-30
T-99
MMD OO
KALE
Figure 1: MMD and KALE flow trajectories for "three rings" target
close to target support and recovers the target almost perfectly. This illustrates ability 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 spiral, and the two distributions have uniform density their support. Again, because the supports the source and the target are disjoint, the flow cannot be defined, nor simulated for this pair. We run KALE particle descent algorithm, and compare it as before with an MMD flow, as well as with "Sinkhorn descent algorithm" [25]. Results are in Fig. 2.
As we can see, the soft support-sharinn constraint informing the KALE flow forces 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 Rdd. Hence, unlike in the two previous examples, the KI, 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 preserved in a gradient flow setting. We consider a
9

---

T= 0
T= 10
T=
T= 60
Target
KALE 0000
Sillon
Figure Shape Transfer using KALE flow
balanced mixture of 4 Gaussians with means located on the 4 corners of the unit square for the target and source distribution given by unit Gaussian in the vicinity of the unit square. We then run KL, MMD, and KALE flows with different values of X, and compute the Wasserstein distance between reference particles iteration n from either the MMD or KL flow and particles obtained from the KALE flow at same iteration n. The choice the Wasserstein distance is natural for Wasserstein Gradient Flows. As shown in Fig. 3a, for "small" values of A. particles from a KALE flow remain close to the ULA particles, while for "large" ones they remain close to the MMD particles (Fig. 3h).
distance to KL
distance to MMD
KALE
A
10
1,00
0.75
10
A 0.001
-
0.001
0C1
B 2--0 0.3
0.25
A 10000
10000
8
0
20000
40000
60000
20000
40000
60000
20000
40000
60000
Iteration
Iteration
iteration
(a)
(b)
(c)
Figure 3: (a): Evolution of the Wasserstein distance between reference particles from the ULA algorithm and the KALE particle descent algorithm with various values of X. (b) Left: same (a), but taking particles from the MMD flow as reference. (c) Evolution the KALE along the traeectories a KAIE. descent algorithm with same mixture of Gaussians as target, using = 0.1. Orange: without noise injection. Blue: with noise injection using a constant noise schedule.
Impact of noise injection On all three examples, using regularized KALE flow with appropri- ately tuned 3. schedule always improves the proximity to the global minimum IFx Q. effect is particularly impactful in the mixture of Gaussians example, where small, constant noise schedule 3n allows for faster mixing times for P., as opposed to its unregularized counterpart, see Fig. 3c. We provide further details on the impact of noise injection in the appendix.
6 Discussion and further work
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 KL gradient flow cannot be defined, trajectories of the KALE flow empirically exhibit better convergence properties when compared the MMD flow, flow that the KALE is also able to interpolate. In cases where KL flow can defined, we notice empirically that the KALE flow can approximate the trajectories but using only information from samples of the target. This latter property is in sharp contrast with KI. Gradient Flow discretizations like the Unadjusted I Langevin Algorithm: in this regard, we could use the KALE flow as sample-based approximation of the KL flow, which is to our knowledge a novel concept. Future work would analyze when the KALE flow is consistent estimator of the KL flow in the large sample limit.
KI. flow,
10

---

References
[1] K. Ahn and S. Chewi. Efficient Constrained Sampling via the Langev-Langeein Algorithm. arXiv preprint 20Ii:20000.12222, 2020.
[2] S. M. Ali and S. D. Silvey, A General Class of Coefficients of Divergence of One Distribution from Another. Journal of the Royal Statistical Society: Series (Methodological), 1966.
[3] L. Ambrosio, N. Gigli, and G. Savaré. Gradient Flows: in Metric Spaces and in the Space of Probability Measures. Springer Science & Business Media, 2008.
[4] M. Arbel, Korba, A. Salim, and Gretton. Maximum Mean Discrepancy Gradient Flow, Advances in Neural Information Processing Systems, 2019.
[5] M. Arbel, D. Sutherland, M. Binkowski, and A. Gretton. On Gradient Regularizers for MMD GANs. In Advances Neural Information Processing Systemx, 2018.
[6] M. Arbel, L. Zhou, and A. Gretton. Generalized Energy Based Models. In International Conference on Learning Representations, 2021.
[7] M. Arjovsky and I.. Botton. Towards Principled Methods for Training Generative Adversarial Networks. International Conference on Learning Reprexentations, 2017.
[8] M. Arjovsky. S. Chintala, and L. Bottou. Wasserstein Generative Adversarial Networks. In International Conference Machine Learning, 2017.
[9] M. G. Bellemare, I. Danihelka, W. Dabney, S. Mohamed, Lakshminarayanan, S. Hoyer, and R. Munos. The Cramer Distance as a Solution to Biased Wasserstein Gradients. CoRR, abs/1 705,10743, 2017.
[10] M. Binkowski, D. J. Sutherland, M. Arbel, and Gretton. Demystifying MMD GANs. International Conference on Learning Representations, 2018.
[11] T. Birdal, M. Arbel, U. Simsekli, and L. Guibas. Synchronizing Probability Measures on Rotations via Optimal Transport. In Conference on Computer Vision and Pattern Recognition, 2020.
[12] J. Birrell, P. Dupuis, M. A. Katsoulakis, Y. Pantazis, and L. Rey-Bellet, -dive-divessssces:: Interpolating between f-divergences and integral probability metrics. arXiv preprint wxXv:2011.05953. 2020.
[13] M. Blei, A. Kucukelbir, and D. McAuliffe, Variational Inferenee: A Review for Statisticians. Journal of the American statistical Association, 2017.
[14] L. Bottou, M. Arjovsky. D. Lopez-Paz. and M. Oquab. Geometrical Insights for Implicit Generative Modeling. In Braverman Readings Machine Learning. Key Ideas from Inception to Current State. 2018.
[15] H. Brezis. Functional Anolysis, Sobolev Spaces and Partial Differential Equations. Springer New York, 2010.
[16] S. Bubeck, R. Eldan, and J. Lehec. Sampling from a Lon-ccn Distribution with Projected Langevin Monte Carlo. Discrete & Computational Geometry, 2018.
[17] L. Cayton. Algorithms for Manifold Learning. Univ. of California at San Diego Tech. Rep, 2005.
[18] L. Chizat and F. Bach. On the Global Convergence of Gradient Descent for Over parameterized Models using Optimal Transport. In Advances in Neural Information Processing Systems, 2018. 8.
[19] A. S. Dalalyan and A. Karagulyan. User-friendly guarantees for the langevin monte carlo with inaccurate gradient. Stochastic Processes and their Applications, 2019.
[20] H. Daniels. The asymptotic efficiency of maximum likelihood estimator. In Berkeley Symposiun Mathematical Statistics and Probability, 1961.
11

---

[21] R. M. Dudley, Real Analysis and Probability. Cambridge University Press, 2002.
[22] A. Durmus, S. Majewski, and B. Miasojedow. Analysis of langevin monte carlo via convex optimization. arXiv preprint arXiv: /802.09188, 88, 2018.
[23] A. Durmus and E. Moulines. Nonasymptotic Convergence Analysis for Unadjusted Langevin Algorithm. The Annals of Applied Probability, 2017.
[24] G. K. Dziugaite, D. M. Roy, and Z. Ghahramani. Training Generative Neural Networks via Maximum Mean Discrepancy Optimization. In Uncertainty in Artificial Intelligence, 2015.
[25] J. Feydy, I. Séjourné, F.-X. Vialard, S.-i. Amari, A. Trouve, and G. Peyré. Interpolating between Optimal Transport and MMD using Sinkhorn Divergences. In International Conference on Artificial Intelligence and Statistics, 2019.
[26] A. Genevay, L. Chizat, Bach. M. Cuturi, and G. Peyré. Sample Complexity of Sinkhorn Divergences. International Conference on Artificial Intelligence and Statistics, 2019.
[27] A. Genevay, G. Peyré, and M. Cuturi. Learning Generative Models with Sinkhom Divergences. In International Conference on Artificial Intelligence and Statistics, 2018
[28] I. Goodfellow, J. Pouget-Abadie, M. Mirza, B. Xu, D. Warde-Farley, S. Ozair, A. Courville, and Y. Bengio. Generative Adversarial Nets. In Advances in Neural Information Processing Systems, 2014.
[29] I. Gulrajani, F. Ahmed, M. Arjovsky, V. Dumoulin, and A. C. Courville, Improved Training of Wasserstein GANs. In Advances in Neural Information Processing Systems, 2017.
[30] Y. Hsich, A. Kavis, P. Rolland, and V. Cevher. Mirrored Langevin Dynamics. In Advances in Neural Information Processing Systems, 2018
[31] R. Jordan, D. Kinderlehrer, and E. Otto. The Variational Formulation of the Fokker-Planck Equation. SIAM journal on mathematical analysis, 1998.
[32] Lang. Fundamentals of Differential Geometry. Graduate Texts in Mathematics. Springer- Verlag, 1999.
[33] C.-L. Li, W.-C. Chang, Y. Cheng, Y. Yang. and B. Póczos. MMD GAN: Towards Deeper Understanding of Moment Matching Network. In Advances in Neural Information Processing Systems, 2017.
[34] Y. Li, K. Swersky, and R. Zemel. Generative Moment Matching Networks. In International Conference on Machine Learning, 2015.
[35] C. Liu and J. Nocedal. On the Limited Memory BFGS Method for Large Scale Optimization. Muthematical progranming, 1989.
[36] Q. Liu. Stein Variational Gradient Descent as Gradient Flow. Advances in Neural Information Processing Systems, 2017.
[37] S. Liu, O. Bousquet, and K. Chaudhuri. Approximation and Convergence Properties of Genera- tive Adversarial Learning. In Advances in Neural Information Processing Systems, 2017.
[38] Y.A. Ma, S. Chatterji. X. Cheng, N. Flammarion, P. L. Bartlett, and M. I. Jordan. Is there Analog of Nesterov Acceleration for Gradient basddd MCMC? Bernoulli, 2021.
[39] H. P. McKean. A Class of Markov Processes Associated with Nonlinear Parabolic Equations. Proceedings of the National Academy of Sciences, 1966.
[40] Mei, A. Montanari, and P.M. Nguyen. Mean Field View of the Landscape of Two-layer Neural Networks. Proceedings of the National Academy of Sciences, 2018.
[41] L. Mescheder. S. Nowozin. and A. Geiger. Adversarial Variational Bayes: Unifying Variational Autoencoders and Generative Adversarial Networks. In International Conference Machine Learning, 2017.
12

---

[42] P. Milgrom and Segal. Envelope Theorems for Arbitrary Choice Sets. Econometrica, 2002.
[43] Y. Mroueh, C. Li, T. Sercu, A. Raj, and Y. Cheng. Sobolev GAN. In International Conference Learning Representations, 2018.
[44] Y. Mrouch and M. Rigotti. Unbalanced Sobolev Descent. Advances in Neural Information Processing Systems, 2020.
[45] Mrouch, T. Sercu, and A. Raj. Sobolev Descent. International Conference on Artificial Intelligence and Statisticx, 2019.
[46] A. Müller. Integral Probability Metrics and their Generating Classes of Functions. Advances Applied Probability, 1997.
[47] H. Narayanan and S. Mitter. Sample Complexity of Testing the Manifold Hypothesis. In Advances in Neural Information Processing Systems, 2010.
[48] X. Nguyen, M. J. Wainwright, and M. I. Jordan. Estimating Divergence Functionals and the Likelihood Ratio by Convex Risk Minimization. IEEE Transactions on Information Theory, 2010.
[49] S. Nowozin, B. Cseke, and R. Tomioka. f-GAN:
Training Generative Neural Samplers using Variational Divergence Minimization. In Advances in Neural Information Processing Systems, 2016.
[50] F. Otto and C. Villani. Generalization of an Inequality by Talagrand and Links with the Logarithmic Sobolev Inequality. Journal of Functional Analysis, 2000.
[51] R. T. Rockafellar. Convex Analysis. Princeton University Press, 1970.
[52] G. M. Rotskoff and E. Handen-Eijnden. Neural Networks Interacting Particle Systems: Asymptotic Convexity of the Loss Landscape and Universal Scaling of the Approximation Error.
CoRR, httb:/8805.0091 2018.
[53] A. Salim and P. Richtárik. Primal Dual Interpretation of the Proximal Stochastic Gradient Langevin Algorithm. In Advances in Neural Information Processing Systems, 2020.
[54] F. Santambrogio. Optimal Transport for Applicd Mathematicians Calculus of Variations, PDEs and Modeling, 2015.
[55] F. Santambrogio. (Euclidean, metric, and Wasserstein] Gradient Flows: An Overview. Bulletin of Mathematical Sciences, 2017.
[56] E. Shannon.
Mathematical Theory of Communication. The Bell system technical journal,
1948.
[57] C.-J. Simon-Gabbian and B. Schölkopf. Kemel Distribution Embeddings: Universal Kernels, Characteristic Kernels and Kernel Metries on Distributions. The Journal of Machine Learning Research, 2018
[58] J. Sirignano and K. Spiliopoulos. Mean Field Analysis of Neural Networks: A Central Limit Theorem. Stochastic Processes and their Applications, 2020.
[59] B. Sriperumbudur. On the optimal estimation of probability measures in weak and strong topologies. Bernoulli, 2016.
[60] 1. Steinwart and A. Christmann. Support Vector Machines. Springer Science & Business Media, 2008
[61] I. Tolstikhin, B. K. Sriperumbudur, Mu, and Minimax Estimation of Kernel Mean Embeddings. Journal of Machine Learning Research, 2017.
[62] T. Van Erven and P. Harremos. Rényi divergence and Kullback-Leible divergence. IEEE Transactions on Information Theory, 2014.
13

---

[63] Villani. Trend to Equilibrium Dissipative Equations. Functional Inequalities and Mass Transportation. Contemporary Mathematics, 2004.
[64] C. Villani. Optimal Transport: Old and New. Springer Science & Business Media, 2008.
[65] J. Weed and F. Bach. Sharp Asymptotic and Finite-sample Rates Convergence Empirical Measures in Wasserstein distance. Bernoulli, 2019.
[66] P. Zhang, Q Liu, D. Zhou, T. Xu, and X. He. On the Discrimination Generalization Tradcoff in GANs, International Conference on Learning Representations, 2018.
[67] D.-X. Zhou, Derivative Reproducing Properties for Kernel Methods in Learning Theory. Journal of Computational and Applied Mathematics, 2008.
14


Reviews:
[{"review_id": "_mI2jOwDFJq", "reviewer": "Reviewer_Go8M", "summary": "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).", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The theoretical aspect of the paper is quite solid. While the established properties of KALE as a probability divergence are not surprising - after all it interpolates MMD and KL and hence ought to enjoy the properties of both - the exposition is clear and rigorous. It is perhaps a bit lacking in motivation: what motivates the regularization in (2), and what does the term do intuitively? The choice seems rather arbitrary to me. It would also be nice to have some comparison with Sinkhorn divergence, which also interpolates between MMD and KL: what theoretical advantage does KALE have compared to it? why do we need another interpolating divergence?\n\nUsing KALE as a gradient flow is interesting as it circumvents a lot of the problems that KL flow has and it is curious that it converges slower than KL. The particle descent algorithm is natural and not very novel, but the consistency of the descent algorithm is well justified. I'm wondering if it is possible to avoid space discretization (particle-based methods)? For instance, you can parameterize $\\mathbb{P}_n$'s using normalizing flow, and then using techniques similar to \"Stochastic Optimization for Large-scale Optimal Transport\" by Genevay et al. for solving for finding $h_n^\\star$ in an RKHS. To me, there is this disconnection where a large part of the theory is based on the RKHS theory but the algorithm has no use of it.\n\nFor the experiments, it is nice that the properties of KALE flow are checked empirically, but the setups are simple 2D examples. I will be interested in seeing how KALE would fare in higher dimensional tasks, such as when the density is defined on a submanifold. In general, I found the experiments section a bit weak and lacks applications: for what practical problems would people use KALE gradient flow for, compared to alternatives? Some comparison with other gradient flow algorithms (such as the SVGD algorithm by Chewi et al.) on a larger scale setting would be nice.\n\nSome minor comments:\n- Line 147: what is $h^*$?\n- Line 608: \"litterature\"\n- Figure 2: where is MMD result?\n- Algorithm 1 could use some captions to point back what \"backprop\", \"log_ratio\" refer to in the main text.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors have not addressed the limitations or potential negative societal impact.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "U2acOGyxOLl", "reviewer": "Reviewer_ftSD", "summary": "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. ", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Originality: The paper is in two parts. \n\nThe first part is about defining the KALE, a relaxation of the KL which is not equal to infinity if the two distributions are singular. The KALE depends on a relaxation parameter and on the choice of a RKHS, and is shown to interpolate between the MMD and the KL (by setting the parameter accordingly). The main mathematical remark here is that the derivative at zero of the objective defining the KALE is the witness function whose norm is the MMD. This is hidden is the appendix, I would suggest to state it in the main paper as a Lemma. \nExcept the weak continuity of KALE, this part seems to be mainly an adaptation of [5], for the purpose of defining the KALE flow, which is the topic of the second (and main) part of the paper. \n\nThe second part is about the KALE flow. The authors first establish the existence and uniqueness of the KALE flow, a task which is non trivial due to the fact the KALE doesn't have the form of a free energy.  Then, under a boundedness assumption of the flow, they show a O(1/t) convergence rate of KALE to zero along the flow. This results seems to be mainly an extension of a result of [3]. \nI don't think that the convergence of the KL would be faster (as suggested after Prop 3) because KL does not satisfy log Sobolev inequality.\n\nAdditionally, the authors propose their main algorithm, obtained as a time discretization of the KALE flow. This algorithm is a provable approximation of the Euler discretization of the KALE flow (which is intractable). \n\nFinally, they propose a modified algorithm by introducing a Gaussian noise at each iteration of the algorithm. This seems similar to a technique used in [3]. This is the least convincing part of the paper because Prop 5 relies on an assumption that cannot be checked (and I cannot find a case where this assumption holds, it can be the case that this assumption is super strong), and empirically this noise injection does not lead to large improvements. Therefore, the introduction of noise is not really justified. \n\nThe simulations of the KALE flow are nice and show what the KALE flow brings compare to other flows such as MMD flow of KL flow. The KALE flow improves the convergence in the case where the initial and the target distribution have disjoint supports. \n\nRefs: Various papers have investigated the task of sampling in the case where source and target have disjoint supports, **in the case where the log density is known** in particular Arxiv 1507.02564, 1802.10174, 2006.09270, 2010.16212\n\n\n\n\n\n\n\n\n\nQuality and clarity: The paper is rather clear but would deserve more polishing. There are many typos in the appendix, and also some typos in the main paper (l.13, l. 116, l.214, l.284, Eq (13) etc.). \nIt would be better to move the assumptions in the main paper (and correct the typos in the statement of the assumptions). \nThe bibliography has not been cleaned before submission. Refs [1] and [2] are almost the same and the authors refers to [2] as [1] sometimes in the appendix. Venues, names etc. have to be harmonized. \nFinally, the implementation of Algo 1 should be discussed a bit more (2 more lines would be sufficient).\n\nSignificance: Overall, the content is novel and relevant to the community, but the main results seem close to those in [3] and [5].\n\n\nMINOR:\n\nl. 57: \"generalization properties\"? I don't understand\nl.92: It seems better to write that we work on an open set X endowed with its Borelian sigma-algebra\nl.124: Strong convexity is used to prove that KALE is a probability divergence\nl.151: Why not using directly the primal formulation of KALE?\nl.295: Why are the heart and the spiral supported on disjoint subsets? This is not seen on the figure.\nl.306 \"the target admits a positive density\"\n\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"ethical_concerns": "N/A (theoretical work)", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "7"}, "extra_text": {}}, {"review_id": "I6bsvgE-FhJ", "reviewer": "Reviewer_WiM9", "summary": "The authors propose KALE Flow to transport mass from one distribution to another. They also study the properties and convergence of the proposed algorithm.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "\nThis paper provides a good theoretical analysis of KALE, and introduces a particle descent algorithm.\nHowever,  as the authors stated, Arbel M, Zhou L, Gretton A.  2021 studied KALE previously.  The main contribution in this paper  compared with Arbel M, Zhou L, Gretton A.  2021, is   the optimization guarantees. However, the main idear of this  part is similar to the MMD flow paper. For the numerics,  I am not really convinced by the toy experiments.\n\n It is better to present several concrete examples, in which the kernelized KL divergence has a closed-form solution. Or the authors may provide several toy examples to compare the advantage of the proposed generalized KL divergences.\nSome  related works are missing, for example,  Li et.al, Natural gradient via optimal transport, Information geometry, 2018; Li, Hessian metric via transport information geometry, 2020.\n\nAnother issue is  the exponential dependence on the iteration number $n$ seems unavoidable in the current context. If my understanding is right, KALE flow converges at an extremely slow rate $O(poly(1/\\log n))$ by choosing $n = O(\\log N)$.\nTherefore,  the KALE flow becomes less appealing. \n\n\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "No.", "review_text": "", "extra_scores": {"ethical_concerns": "No.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "G9QwCeubHHl", "reviewer": "Reviewer_2uTN", "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. \n\nFirst, 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)). \n\nThen, 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).\n", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This work is relevant for the NeurIPS community, since prior work have demonstrated that gradient flows can be employed in modern machine learning approaches, including generative modeling. Moreover, the newly introduced probability divergence KALE is clearly motivated: according to the empirical analysis in Section 5, KALE can return useful information when comparing two distributions, even when their support do not overlap (unlike KL), and might offer a higher discriminative power than MMD. These advantages, combined with the topological properties in Theorem 1 and the fact that it can be estimated when the two distributions are accessible through samples only, suggest that KALE has great potential for other applications than gradient flows.  \n\nThe paper is very dense as it provides numerous theoretical contribution on novel tools, and due to lack of time, I was not able to thoroughly check all the proofs in the appendix. The empirical analysis in Section 5 illustrates well the flexibility of KALE gradient flow over the MMD and KL ones, likely due to its interpolation property (Proposition 2). The rest of the claims were not empirically verified, and in my opinion, some of them should be discussed more: \n\n(A1) \"As shown in [5], the KALE's sample complexity scales with $\\frac1{\\sqrt{N}}$\" (l.157-158), where [5] refers to \"Generalized energy based models\" (Arbel et al., 2021). However, the authors also emphasized that the KALE they consider \"differs from the one in [5]\" (l.111). In that sense, it is not clear why the sample complexity in [5] still applies here. On the other hand, that sample complexity is said to be \"similar to the sample complexity of Sinkhorn divergences\". The sample complexity of Sinkhorn divergences also strongly depends on the value of its regularization coefficient; is it the case for KALE?\n\n(A2) The convergence rate of KALE flow derived in Proposition 3 is said to be slower than the one yielded by KL flow (l.192), but it is not clear how it compares with MMD flow (in particular, with Proposition 7 in \"Maximum Mean Discrepancy Gradient Flow\", Arbel et al., 2019). Therefore, I have trouble understanding why \"trajectories of the KALE flow exhibit better convergence properies when compared to the MMD flow\" (l.326-327): it this claim solely supported by the empirical results in Section 5 for now?\n\n(A3) Proposition 4 establishes an approximation error bound, which guarantees that \"given sufficiently many samples of $\\mathbb{P}_0$ and $\\mathbb{Q}$, one can approximate an exact discrete KALE flow ... with arbitrary precision\" (l.240-242). However, the derived bound depends on other constants than the number of samples, whose impact is not discussed at all. In particular, the smaller $\\lambda$, the larger $A$ (l.239), so the larger $N$ needs to be to compensate. This might be a limitation of KALE in practice and it should be addressed in the paper to give a more sensible interpretation. \n\nThe aforementioned points explain why I can't give a higher overall score for now. \n\nThis paper is very well written, but some parts in the main document are not self-contained, which makes them hard to understand. I encourage the authors to address the following points to improve readability.\n\n(B1) Assumptions 1 to 3, which play an important role in most theoretical results in this work, are only presented in the appendix (Section B). The main document should at least provide a small description of them. Besides, it is not clear to which extent these assumptions are met in the experiments.\n\n(B2) The notation $\\dot{H}^{-1}(\\mathbb{P}_t)$ in the negative Sobolev norm is not clear. The definition of this norm (given in Section F) could be recalled in Proposition 3.\n\n(B3) The set $\\mathcal{H}_n$ in equation (13) is not defined.\n\n(B4) The constants in Proposition 4 ($L$, $K$, $\\mu$) are not explained. This point is actually connected to (B1), since these constants come from the assumptions.\n\nSome other minor inaccuracies and typos are given below.\n\nl.97: \"an superscript\" -> \"a superscript\"\n\nEquation (7): Given equation (6), why is it -1 instead of +1?\n\nEquation (11): $\\partial \\mathbb{P}_t$ -> $\\partial_t \\mathbb{P}_t$\n\nEquation (13): What is the norm used here? (Precise the set)\n\nl.214 \"Approxation\" -> \"Approximation\"\n\nEquation (15): the set of $h$ should be precised. \n\nAlgorithm 1: \n- Extra \"max_iter\",\n- The variable $\\texttt{v}$ should depend on $\\texttt{h}$_$\\texttt{star}$,\n- Precise what operation $\\texttt{log}$_$\\texttt{ratio}$ does.\n\nl.257: Given the rest of the statement, $\\mathcal{D}_{\\beta_n}$ should be defined as a function of $\\mathbb{P}_n$.\n\nl.272: \"appearing in the the\"\n\nl.284: \"and thus the the\" \n\nFigure 3: in the legend, \"and and\"\n\nl.624: \"ahd\" -> \"and\"\n\nl.646: \"metrices\" -> \"metrizes\"\n\nEquation below l.651: $\\mathcal{K}(0, \\mathbb{P})$ -> $\\mathcal{K}(0_{\\mathcal{H}}, \\mathbb{P})$\n\nEquation (27): $\\mathbb{P}$ should be $\\mathbb{P}_t$\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors addressed two limitations: (i) the condition on the noise schedule in Proposition 5 might be too strict, given Figures 3 to 5; (ii) KALE flow (Figure 2) does not recover outliers well (Figure 2, l.303-305). \n\nThe societal impact is not discussed.\n", "review_text": "", "extra_scores": {"ethical_concerns": "This paper does not seem to raise any ethical issues.\n", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper develops KALE particle flow, although the paper is incremental and the noise injection scheme is an adaption from MMD flow, reviewers and the AC agreed that the  theoretical contribution of the paper is reasonable  but the experimental side of the paper is rather weak. Weak accept  
Decision: Accept (Poster)