Reliable Estimation of KL Divergence using a Discriminator in Reproducing Kernel Hilbert Space
Year: 2022

Abstract:
Estimating Kullback–Leibler (KL) divergence from samples of two distributions is essential in many machine learning problems. Variational methods using neural network discriminator have been proposed to achieve this task in a scalable manner. However, we noted that most of these methods using neural network discriminators suffer from high fluctuations (variance) in estimates and instability in training. In this paper, we look at this issue from statistical learning theory and function space complexity perspective to understand why this happens and how to solve it. We argue that the cause of these pathologies is lack of control over the complexity of the neural network discriminator function space and could be mitigated by controlling it. To achieve this objective, we 1) present a novel construction of the discriminator in the Reproducing Kernel Hilbert Space (RKHS), 2) theoretically relate the error probability bound of the KL estimates to the complexity of the discriminator in the RKHS space, 3) present a scalable way to control the complexity (RKHS norm) of the discriminator for a reliable estimation of KL divergence, and 4) prove the consistency of the proposed estimator. In three different applications of KL divergence – estimation of KL, estimation of mutual information and Variational Bayes – we show that by controlling the complexity as developed in the theory, we are able to reduce the variance of KL estimates and stabilize the training.

heading: 1 Introduction
text: Estimating Kullback–Leibler (KL) divergence from data samples is an essential component in many machine learning problems including Bayesian inference, calculation of mutual information or methods using information theoretic objectives. Variational formulation of Bayesian Inference requires KL divergence computation, which could be challenging when we only have finite samples from two distributions. Similarly, computation of information theoretic objectives like mutual information requires computation of KL divergence between the joint and the product of marginals.
KL divergence estimation from samples was studied thoroughly by Nguyen et al. [1] using a variational technique, convex optimization and RKHS norm regularization, while also providing theoretical guarantees and insights. However, their technique requires handling the whole dataset at once and is not scalable. Many modern models need to use KL divergence with large scale data, and often with neural networks, for example total correlation variational autoencoder (TC-VAE) [2], adversarial variational Bayes (AVB) [3], information maximizing GAN (InfoGAN) [4], and amortized MAP [5] all need to compute KL divergence in a deep learning setup. These large scale models have imposed new requirements on KL divergence estimation like scalability (able to handle large amount of data samples) and minibatch compatibility (compatible with minibatch-based optimization).
Methods like Nguyen et al. [1] are not suitable in the large scale setup. These modern needs were later met by modern neural network based methods such as variational divergence minimization
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
(VDM) [6], mutual information neural estimation (MINE) [7], and discriminator based KL estimation with GAN-type objective [8, 5]. A key attribute of these methods is that they are based on updating a neural-net based discriminator to estimate KL divergence from a subset of samples making them scalable and minibatch compatible. We, however, noticed that even in simple examples, these methods exhibited pathologies like unreliability (high fluctuation of estimates) or instability during training (KL estimates blowing up). Similar observations of instability of VDM and MINE have also been reported in the literature [8, 9].
Why are these techniques unreliable? In this paper, we attempt to understand the core problem in the KL estimation using discriminator network. We look at it from the perspective of statistical learning theory and discriminator function space complexity and draw insights. Based on these insights, we propose that these fluctuations are a consequence of not controlling the smoothness and the complexity of the discriminator function space. Measuring and controlling the complexity of function space itself becomes a difficult problem when the discriminator is a deep neural network. Note that naive approaches to bound complexity by the number of parameters would neither be guaranteed to yield meaningful bound [10], nor be easy to implement.
Therefore, we present the following contributions to resolve these challenges. First, we propose a novel construction of the discriminator function using deep network such that it lies in a smooth function space, the Reproducing Kernel Hilbert Space (RKHS). By utilizing the learning theory and the complexity analysis of the RKHS space, we bound the probability of the error of KLdivergence estimates in terms of the radius of RKHS ball and kernel complexity. Using this bound, we propose a scalable way to control the complexity by penalizing the RKHS norm. This additional regularization of the complexity is still linear, (O(m)) in time complexity with the number of data samples. Then, we prove consistency of the proposed KL estimator using ideas from empirical process theory. Experimentally, we demonstrate that the proposed way of controlling complexity significantly improves KL divergence estimation and significantly reduce the variance. In mutual information estimation, our method is competitive with the state-of-the-art method and in Variational Bayesian application, our method stabilizes training of MNIST dataset leading to sharp reconstruction.

heading: 2 Related Work
text: Nguyen et al. [1] used variational method to estimate KL divergence from samples of two distribution using convex risk minimization (CRM). They used the RKHS norm as a way to both measure and penalize the complexity of the variational function. However, their work required handling all data at once and solving a convex optimization problem which has time complexity in the order of O(m3) and space complexity in the order of O(m2) . Ahuja [11] used similar convex formulation in RKHS space and found it difficult to scale. VDM reformulated the f-Divergence objective using Fenchel duality and used a neural network to represent the variational function [6]. Although close in concept to [1], it is scalable since it uses a separate discriminator network and adversarial optimization. It, however, did not control the complexity of the neural-net function, and faced issues with stability.
One area of modern application of KL-divergence estimation is in computing mutual information, which is useful in applications such as stabilizing GANs [7]. MINE [7] also optimized a lower bound to KL divergence (Donsker-Varadhan representation). Similar to VDM, MINE used a neural network as the dual variational function: it is thus scalable, but without complexity control and is unstable. Another use of KL divergence is scalable variational inference (VI) as shown in AVB [8]. VI requires KL divergence estimation between the posterior and the prior, which becomes nontrivial when a sample based scalable estimation is required. AVB solved it using GAN-type adversarial formulation and a neural network discriminator. Similarly, [5] used GAN-type adversarial formulation to obtain KL divergence in amortized inference.
Chen et al. [2] proposed TC-VAE to improve disentanglement by penalizing the KL divergence between the marginal latent distribution and the product of marginals in each dimension. The KL divergence was computed by a minibatch-based sampling strategy that gives a biased estimate. Our work is close to Song et al. [9] who investigated the high variance in existing mutual information estimators and found that clipping the discriminator output is helpful in reducing variance. In our work, we take a principled way to connect variance to the complexity of discriminator function space and constrain it by penalizing its RKHS norm instead. None of the existing works considered looking at the discriminator function space, connecting its complexity to the unreliable KL-divergence estimation, or mitigating the problem by controlling the complexity.

heading: 3 Reproducing Kernel Hilbert Space
text: LetH be a Hilbert space of functions f : X → IR defined on non-empty space X . It is a Reproducing Kernel Hilbert Space (RKHS) if ∀x ∈ X , the evaluation functional, δx : H → IR, δx : f 7→ f(x), is linear and continuous at every f . Every RKHS, HK , is associated with a unique positive definite kernel, K : X × X → IR, called the reproducing kernel [12], such that it satisfies: 1. ∀x ∈ X ,K(., x) ∈ HK , 2. ∀x ∈ X ,∀f ∈ HK , 〈f,K(., x)〉HK = f(x) RKHS is studied using a specific integral operator. Let L2(dρ) be a space of functions f : X → IR that are square integrable with respect to a Borel probability measure dρ on X , we denote an integral operator LK : L2(dρ)→ L2(dρ) [13, 14]: (LKf)(x) = ∫ X f(y)K(x, y)dρ(y). This operator will be important in constructing a function in RKHS and in computing sample complexity.

heading: 4 Problem Formulation and Contribution
text: GAN-type Objective for KL Estimation: Let p(x) and q(x) be two probability density functions in space X and we want to estimate their KL divergence using finite samples from each distribution in a scalable and minibatch compatible manner. As shown in [8, 5], this can be achieved by using a discriminator function. First, a discriminator f : X → IR is trained with the objective:
f∗ = argmax f [Ep(x) log σ(f(x)) + Eq(x) log(1− σ(f(x)))] (1)
where σ is the Sigmoid function given by σ(x) = e x
1+ex . Then it can be shown [8, 5] that the KL divergence KL(p(x)||q(x)) is given by: KL(p(x)||q(x)) = Ep(x)[f∗(x)] Sources of Error: Eq. (1) is ambiguous in the sense that it is silent about the discriminator function space over which the optimization is carried out. Typically, a neural network is used as the discriminator. This implies that we are considering the space of functions represented by the neural network of given architecture as the hypothesis space, over which the maximization occurs in eq. (1). Hence, we must rewrite eq. (1) as
f∗h = argmax f∈h [Ep(x) log σ(f(x)) + Eq(x) log(1− σ(f(x)))] (2)
where h is the discriminator function space. Furthermore, we also approximate integrals in eq. (2) with the Monte Carlo estimate using finite number of samples, say m, from the distribution p and q.
fmh = argmax f∈h [ 1 m ∑ xi∼p(xi) log σ(f(xi)) + 1 m ∑ xj∼q(xj) log(1− σ(f(xj))) ]
(3)
Similarly, we write KL estimate obtained from, respectively, infinite and finite samples as:
KL(f) = Ep(x)[f(x)], KLm(f) = 1
m ∑ xi∼p(xi) [f(x)] (4)
Each of these steps introduce some error in our estimate. We can now start our analysis by first decomposing the total estimation error as:
KLm(f m h )−KL(f∗) = KLm(fmh )−KL(fmh )︸ ︷︷ ︸
Deviation-from-mean error +KL(fmh )−KL(f∗h)︸ ︷︷ ︸ Discriminator induced error +KL(f∗h)−KL(f∗)︸ ︷︷ ︸ Bias
(5)
This equation decomposes total estimation error into three terms: 1) deviation from the mean error, 2) error in KL estimate by the discriminator due to using finite samples in optimization eq. (3), and 3) bias when the considered function space does not contain the optimal function. Here, we concentrate on quantifying the probability of deviation-from-mean error which is directly related to observed variance of the KL estimate.
Summary of Technical Contributions: Since the deviation is the difference between a random variable and its mean, we can bound the probability of this error using concentration inequality and the complexity of the function space of fmh . To use smooth function space, we propose to construct a function out of neural networks such that it lies on RKHS (Section 5). Then, we bound the probability of deviation-from-mean error through the covering number of the RKHS space (Section 6.1), then control complexity (Section 6.2) and prove consistency of the proposed estimator (Section 7).

heading: 5 Constructing f in RKHS
text: The following theorem due to Bach [15] paves a way for us to construct a neural function in RKHS. Theorem 1. [[15] Appendix A] A function f ∈ L2(dρ) is in Reproducing Kernel Hilbert Space,HK , if and only if it can be expressed as
∀x ∈ X , f(x) = ∫ W g(w)ψ(x,w)dτ(w), (6)
for a certain function g :W → R such that ||g||2L2(dτ) <∞,W is a compact space and functions w 7→ ψ(x,w) are measurable for all x. The RKHS norm of f satisfies ||f ||2HK ≤ ||g|| 2 L2(dτ) and the kernel K is given by
K(x, t) = ∫ W ψ(x,w)ψ(t, w)dτ(w) (7)
Theorem 1 not only gives us a condition when a square integrable function is guaranteed to lie in RKHS, it also provides us with a recipe to construct a function in RKHS. We construct f using this theorem where ψ and g are realized with the neural networks and dτ is a probability measure with Gaussian distribution. We sample w ∼ N (0, γI) and pass it through two neural networks, ψ and g, where ψ takes x and w as two arguments and g takes only w as an argument. More precisely, we consider ψ(x,w) = φθ(x)Tw, where φθ is a neural network with parameters, θ. The kernel K, as defined in eq. (7), can be obtained as:
Kθ(x ∗, t∗) = ∫ W φθ(x ∗)TwwTφθ(t ∗)dτ(w) = γφθ(x ∗)Tφθ(t ∗) (8)
where Ew∼N (0,γI)[wwT ] = γI. We sometimes denote the kernel K by Kθ to emphasize that it is a function of neural network parameters, θ. Furthermore, representation of f as in Theorem 1 provides us an important upper bound on the RKHS norm of f as ||f ||2HK ≤ ||g|| 2 L2(dτ) which we will use later to bound the complexity of the discriminator function space.
Traditionally, kernel K remains fixed and the norm of the function f determines the complexity of the function space. In our formulation, both the RKHS kernel and the norm of f with respect to the kernel changes during training since the kernel depends on neural network parameters, θ. Therefore, the challenge is to tease out how neural parameters, θ, affect the deviation-from-mean error in eq. (5).

heading: 6 Error Analysis and Control
text: Assumptions: Before starting our analysis, we list assumptions upon which our theory is based. A1. The input domains X andW are compact. A2. The functions φθ and g are Lipschitz continuous with Lipschitz constants Lφ and Lg respectively. A3. Higher order derivatives DαxK(x, t) up to some high order ν/2 of kernel K exist.
Assumptions A1 is satisfied in our experiments since we consider a bounded set in Rn and RD as our domains. Similarly, A2 is satisfied since we enforce Lipschitz continuity of φ and g by using spectral normalization [16]. Assumption A3 is a bit subtle. By the definition of K in eq.(8), higher order derivative of K exists if and only if higher order derivative of φθ exists. This is readily satisfied by deep networks with smooth activation functions, and is true everywhere except at origin for ReLU activation. Using the boundedness of the input domain and Lipschitz continuity, we show the following proposition which will be useful later in the error bounds. Proposition 1. Under the assumptions A1, A2, we have sup
x,t Kθ(x, t) <∞ and ||g||2L2(dτ) <∞.

heading: 6.1 Bounding the Error Probability of KL Estimates
text: Bounding the probability of deviation-from-mean error (eq. (5)) is tricky since, in our case, the kernel is not fixed and we are also optimizing over them. We bound it in two steps: 1) we derive a bound for a fixed kernel, 2) we take supremum of this bound over all the kernels parameterized by θ.
For a fixed kernel, we first bound the probability of deviation-from-mean error in terms of the covering number in Lemma 1. We then use an estimate of the covering number of RKHS due to [14] to relate the bound to kernel Kθ in Theorem 2, identifying the role of neural networks in this error bound.
Lemma 1. Let fmHK be the optimal discriminator function in an RKHSHK which is bounded by M with respect to ||.||∞ . Let KLm(fmHK ) = 1 m ∑ i f m HK (xi) and KL(f m HK ) = Ep(x)[f m HK (x)] be the estimate of KL divergence from m samples and that by using the true distribution p(x) respectively. Then the probability of error at some accuracy level, , is lower-bounded as:
Prob.(|KLm(fmHK )−KL(f m HK )| ≤ ) ≥ 1− 2N (HK , 4 √ SK
) exp(− m 2
4M2 )
where N (HK , η) denotes the covering number of an RKHS space HK with disks of radius η, and SK = sup
x,t K(x, t) which we refer to as kernel complexity.
Proof Sketch. We cover RKHS with discs of radius η = 4 √ SK . Within this radius, the deviation does not change too much. So, we can bound deviation probability at the center of disc and apply union bound over all the discs. To bound deviation probability at the center, we apply Hoeffding’s inequality and applying union bound simply leads to counting number of discs which is exactly the covering number. See supplementary materials for the full proof.
Lemma 1 bounds the probability of error in terms of the covering number of the RKHS space. Note that the radius of the disc is inversely related to SK which indicates how complex the RKHS space defined by the kernel Kθ is. Here Kθ depends on the neural network parameters θ. Therefore, we denote SK as a function of θ as SK(θ) and term it kernel complexity. Next, we use Lemma 2 due to [14] to obtain an error bound in estimating KL divergence with finite samples in Theorem 2. Lemma 2 ([14]). Let K : X ×X → R be a C∞ Mercer kernel and the inclusion IK : HK ↪−→ C(X ) be the compact embedding defined by K to the Banach space C(X ). Let BR be the ball of radius R in RKHSHK . Then ∀η > 0, R > 0, ν > n, we have
lnN (IK(BR), η) ≤ ( RCν η ) 2n ν
(9)
whereN gives the covering number of the space IK(BR) with discs of radius η, and n represents the dimension of the input space X . Cν is given by Cν = Cs √ ||Ls|| where Ls is a linear embedding from square integrable space L2(dρ) to the Sobolev space Hν/2, ||Ls|| denotes operator norm and Cs is a constant.
To prove Lemma 2 [14], the RKHS space is embedded in the Sobolev Space Hν/2 using Ls and then the covering number of the Sobolev space is used. Thus the norm of Ls and the degree of Sobolev space, ν/2, appears in the covering number of a ball inHK . In Theorem 2, we use Lemma 1 and 2 to bound the estimation error of KL divergence. Theorem 2. Let KL(fmH ) and KLm(fmH ) be the estimates of KL divergence obtained by using true distribution p(x) and m samples respectively as described in Lemma 1, then the probability of error in the estimation at the error level is given by:
Prob.(|KLm(fmH )−KL(fmH )| ≤ ) ≥ 1− 2 exp
[( 4RCp √ Sp||Lp|| ) 2n ν
− m 2
4M2
]
where Cp √ Sp||Lp|| = sup
Kθ
Cs √ SK(θ)||Ls||, i.e. Cp, Sp,Lp correspond to a kernel for which the
bound is maximum.
Proof. We prove this in two steps: First we obtain an error bound for a fixed kernel space and apply supremum over all θ. For any RKHSHKθ , with fixed kernel Kθ, we have
Prob.(|KLm(fmHKθ )−KL(f m HKθ )| ≥ ) ≤ 2 exp
[( 4RCs √ SK(θ)||Ls|| ) 2n ν
− m 2
4M2
] (10)
We prove this error bound as follows. Lemma 2 gives the covering number of an RKHS ball of radius R, which we apply to Lemma 1. We fix the radius of discs to η =
4 √ SK in Lemma 1 and substitute Cν = Cs √ ||Ls(θ)|| to obtain eq.(10).
Since we are continuously changing θ during training, the kernel also changes. Hence, to find the upper bound over all possible kernels, we take the supremum over all kernels.
Prob.(|KLm(fmH )−KL(fmH )| ≥ ) ≤ sup Kθ Prob.(|KLm(fmHKθ )−KL(f m HKθ )| ≥ ) (11)
≤ 2 exp
[( 4RCp √ Sp||Lp|| ) 2n ν
− m 2
4M2
] (12)
where Sp = SK(θp) and Lp = LK(θp), i.e., Sp and Lp correspond to kernel complexity and Sobolev operator norm corresponding to optimal kernel Kθp that extremizes eq. (11). Theorem statement readily follows from eq. (12) Theorem 2 shows that the error increases exponentially with the radius of the RKHS space, R, complexity of the kernel SK(θp), and the norm of the Sobolev space embedding operator ||Lp||. The Sobolev embedding operator, Lp, is a mapping from L2(dρ) to the Sobolev space Hν/2. It can be shown [14] that the operator norm can be bounded as ||Lp|| ≤ ρ(X ) ∑ |α|≤ν/2 sup
x,t∈X (DαxKθp(x, t)) 2, where ρ is the measure of the input space X . Therefore,
the norm ||Lp|| directly measures smoothness of Kθp while SK(θp) only depends on the supremum value of Kθp . Cp is a constant not depending on Sp, Lp or R, and it is always finite (see [14, 17] for more details). Sp and Lp are related to φ, γ through the kernel Kθp . Sp is defined as Sp = sup
x,t
K(x, t) = sup x,t γφTθp(x)φθp(t), so it is directly related to the network φθp . We show that this is finite due to Lipschitz constraint on φθp as described in Assumption A2 (see Supplementary material for proof). Lp is upper bounded as follows: ||Lp|| ≤ ρ(X ) ∑ |α|≤ν/2 sup
x,t∈X (DαxKθp(x, t)) 2 = ρ(X ) ∑ |α|≤ν/2 sup
x,t∈X (Dαxγφ T θp (x)φθp(t)) 2 where D is a differentiation operator. Therefore, Lp is
finite if the higher order derivatives of φθp exist and are finite.

heading: 6.2 Complexity Control
text: From Theorem 2, we see that the error probability could be decreased by decreasing R, ||Lp|| and SK(θp). Using argument similar to the proof of Proposition 1, we can show that the Lipschitz constraint on φθ also affects SK and may affect ||Lp||. In our experiments, however, we fix the Lipschitz constraints during optimization and do not change SK and ||Lp|| dynamically. Here, we focus on the norm, R from Theorem 2. From Theorem1, we know that the RKHS norm is upper bounded by the norm of function g as ||f ||2HK ≤ ||g|| 2 L2(dτ). We have access to finite approximation of ||g||2L2(dτ) by construction and hence, we can use it to penalize the complexity during optimization. To obtain the optimal discriminator fmh , we optimize the following objective with an extra penalization of the upper bound, i.e. ||g||L2(dτ) on the RKHS norm of f :
fmh = argmax f∈h
1
m ∑ xi∼p(xi) log σ(f(xi)) + 1 m ∑ xj∼q(xj) log(1− σ(f(xj)))− λ0 m ||g||2L2(dτ) (13)
The regularization term prevents the radius of RKHS ball from growing, maintaining a low error probability. Optimization of eq. (13) w.r.t. neural network parameters θ allows dynamic control of the complexity of the discriminator function on the fly in a scalable and efficient way. Note that, computation of ||g||L2(dτ) requires randomly sampling w ∼ N (0, γI) and passing through neural network g independent of the data xi, xj . Therefore, if the computational complexity of optimization is O(m), it will remain the same after incorporating this additional term, i.e. regularization does not increase asymptotic time complexity which is linear with the number of samples, m.

heading: 7 Variance and Consistency of the Estimate

heading: 7.1 Variance Analysis
text: Theorem 2 gives an upper bound on the probability of error. Intuitively, the variance and probability of error behave similarly for many distributions, i.e. higher variance might indicate higher probability of error. Below we quantify this intuition for a Gaussian distributed estimate:
Theorem 3. Let X = KLm(fmH ) be the estimated KL divergence using m samples as described in Theorem 2. Assuming that X follows a Gaussian distribution X ∼ N (µ, ς2), we can obtain an upper bound on the standard deviation of the estimate as follows:
ς ≤ erf−1 [ − 4 exp [( 4RCp √ Sp||Lp|| ) 2n ν − m 24M2 ] + 1 ]
where erf is the Gauss error function and is a monotonic function.
Theorem 3 suggests that by decreasing R, the radius of the RKHS ball, the variance of the estimate could be decreased. Experimentally, we observe that the variance decreases as we penalize the RKHS norm more, consistent with the spirit of Theorem 3.
Note that Theorem 3 makes a strong assumption that the estimate is distributed as a Gaussian distribution. While it gives us good intuition about the decay of variance as the complexity increases, it is natural to inquire about the validity of this type of relation in a more general sense without assumption on the probability distribution of the estimate. To make a general statement, the key idea is to understand how the cumulative distribution function (CDF) is related to the variance. To clarify this point further, let’s look at the eq.(33) in the proof of Theorem 3 in supplementary material:
1− Φµ,ς(µ+ ) ≤ 2 exp
[( 4RCs √ SK ||Ls|| ) 2n ν
− m 2
4M2
] . This equation connects the CDF to the
variables like SK , R, Ls of the discriminator function space without assuming anything about the shape of the distribution. For a Gaussian distribution, we plug in CDF of a Gaussian distribution, Φµ,ς(x̂) = 1 2 [ 1 + erf ( x̂−µ ς √ 2 )] and obtain the result of Theorem 3. For any other distribution, we can carry out similar analysis. A key factor that determines the behavior between the variance and the discriminator complexity is how variance appears in the CDF expression. For example, in both Gaussian distribution and in exponential distribution, we know that the relation between CDF function and variance is inversely related. We believe that as long as this inverse type of relation between CDF and variance holds, we can obtain statements like Theorem 3 for other distributions as well. This leads to a key insight: similar to the Gaussian case, the decaying behavior of the variance with decreasing complexity holds as long as the estimate is distributed such that its CDF function has inverse relation with the variance.

heading: 7.2 Consistency of Estimates
text: In eq.(13), we use a regularized objective to obtain an optimal discriminator instead of an unregularized objective as in eq.(3). It is important to show that the estimator of KL divergence obtained by using this regularized objective is consistent and approaches the true estimate in the limiting case as the number of data samples grows to infinity. In the following theorem, we show this.
Theorem 4. Let f∗ and fm be optimal discriminators as described in eq. (1) and eq. (13) respectively, and the KL estimate is given by KL(f) = Ep(x)[f(x)], KLm(f) = 1m ∑ xi∼p(xi)[f(x)]. Then, in the limiting case as m→∞, |KLm(fmh )−KL(f∗)| → 0 in probability.
Proof Sketch. The difference between the true KL divergence and the estimated KL divergence can be divided into three terms as shown in eq. (5). We assume that our function space is rich enough to contain the true solution, driving bias to zero. From Theorem 2, we see that in the limiting case of m→ 0, the deviation-from-mean error goes to 0. Therefore, the key step that remains to be shown is that the discriminator induced error (second term in eq.(5)) also goes to 0 as m→∞. It can be shown if we can prove that the optimal discriminator in eq. (13) approaches the optimal discriminator in eq. (2). To prove this, we show that the argument being maximized by fmh approaches the argument being maximized by f∗h in the limiting case. To show this, we need to show that the function space, log σf , is Glivenko Cantelli [18], which we prove in following steps: 1. We show that f is Lipschitz continuous by definition and due to Lipschitz continuity of φθ. Then we show that log σf is Lipschitz continuous if f is Lipschitz continuous. 2. Then we show that for a class of functions with Lipschitz constant L, the metric entropy, logN , can be obtained in terms of L and entropy number of the bounded input space, X . 3. Since the metric entropy does not grow with the number of samples m, we show that 1m logN → 0 which lets us show that log σf belongs to Glivenko Cantelli class of functions by using Theorem 2.4.3 from [18]. See supplementary material for the complete proof.

heading: 8 Experimental Results
text: We present results on three applications of KL divergence estimation: 1. KL estimation between simple Gaussian distributions, 2. Mutual information estimation, 3. Variational Bayes. In our experiments, the RKHS discriminator is constructed with ψ and g networks as described in Section 5, where the network ψ is very close to a regular neural network. In two experiments, we compare our results with the models using regular neural net discriminator to ensure that the difference in performance between RKHS and regular neural network is not due to architectural difference. Our code is publicly available at https://github.com/sandeshgh/Reliable-KL-estimation
KL Estimation between Two Gaussians We assume that we have finite sets of samples from two distributions. We further assume that we are required to apply minibatch based optimization. We consider estimating KL divergence between two Gaussian distributions in 2D, where we know the analytical KL divergence between the two distributions as the ground truth. We consider three different pairs of distributions corresponding to true KL divergence values of 1.3, 13.8 and 38.29, respectively and use m = 5000 samples from each distribution to estimate KL in the finite case. We repeat the estimation experiments with random initialization 30 times and report the mean, standard deviation, scatter and box plots.
Fig. 1 top row compares the estimation of KL divergence with regular neural net and RKHS discriminator with complexity control based on eq. (13). With our proposed RKHS discriminator, the KL estimates are significantly more reliable and accurate: error reduced from 0.5 to 0.04, 5.8 to 1.07 and 60.6 to 9.7 and variance reduced from 0.2 to 0.002, 223 to 4.4 and 3521 to 33 for true KL 1.3, 13.8 and 38.29 respectively. In Fig. 1 bottom row, we investigate our complexity control method on the effect of varying the regularization parameter λ = λ0/m. As expected, increasing regularization parameter penalizes more on the RKHS norm and therefore reduces variance. This is consistent with our theory. Regarding bias, however, as we increase the λ, the bias decreases and then starts to increase. Hence, one needs to strike a balance between bias and variance while choosing λ.
Mutual Information Estimation Computation of mutual information is a direct use case of KL divergence computation. We replicate the experimental setup of [20, 9] to estimate mutual information between (x, y) drawn from 20-d Gaussian distributions, where the mutual information is increased by step size of 2 from 2 to 10. We compare the performance of our method with traditional KL divergence computation methods like contrastive predictive coding (CPC) [19], convex risk minimization (NWJ) [1] and SMILE [9]. In Fig.2, our method with RKHS discriminator (with λ = 1e−5) performs better than CPC [19] and NWJ [1], and is competitive with the state-of-the-art, SMILE [9]. In the bottom row, we also show the effect of regularization parameter λ in our method. Similar to the previous experiment, increasing the regularization parameter decreases the variance and increases the bias. It is consistent with our theoretical insights about the effect of reducing RKHS norm on variance.
Adversarial Variational Bayes Variational Bayes requires KL divergence estimation. When we do not have access to analytical form of the posterior/prior distributions, but only have access to the samples, we need to estimate KL divergence from samples. Adversarial Variational Bayes (AVB) [8] presents a way to achieve this using a discriminator network. We adopt this setup and demonstrate that the training becomes unstable if we do not constrain the complexity of the discriminator. First, we
train AVB on MNIST dataset with a simple neural network discriminator architecture. As the training progresses, the KL divergence blows up after about 500 epochs (Fig. 3(b)) and the reconstruction starts to get worse (Fig. 3(a)). We modify the same architecture according to our construction such that the discriminator lies in the RKHS and then penalize the RKHS norm as in eq. (13). It stabilizes the training for a large number of epochs as shown in Fig. 3(c) and the reconstruction does not deteriorate as the training progresses, resulting into sharp reconstruction (Fig. 3(a)). To make this comparison more precise and quantitative, we compute FID score between reconstruction and the ground truth after each epoch as tabulated in Table 1. In the earlier epochs, the FID score ( and the reconstruction in Fig.3(a)) is okay even for the neural net discriminator, However, this score is worse than using the RKHS discriminator, where the scores are in the range of 30-40 upto epoch 400. For the neural network discriminator, the score increases (worsens) in the mid epochs and becomes completely unstable with FID score shooting up to 262 at epoch 900. On the other hand, for the RKHS discriminator, the score steadily and smoothly decreases as the epoch increases reaching the best 28.17 at epoch 900. This experiment demonstrates that the proposed RKHS discriminator with norm regularization is both reliable and effective in terms of standard metric like FID.
We want to clarify that this instability in training neural net discriminator is present if we use a basic discriminator architecture. It does not mean that there exists no other method to design a stable neural net discriminator. In fact, AVB [8] presents a discriminator that adds additional inner product structure to stabilize the discriminator training. Our point here is that we can stabilize the training by ensuring that the discriminator lies in a well behaved function space (the RKHS) and controlling its complexity, consistent with our theory.

heading: 9 Limitations, Discussion and Conclusion
text: Limitations: The proposed construction of neural function in RKHS exhibits good properties of both the deep learning and kernel methods. However, it requires constructing two separate deep networks, ψ and g. It makes our model a bit bulky and also requires more parameter due to additional g. Moreover, currently our RKHS discriminator’s output is scalar; generalizing this function to a multivariate output could make our model bulkier and increase parameters even more. Second limitation is the requirement of higher order derivative of kernel K in assumption A3. While this requirement is satisfied if smooth activation function is used in φθ, for activations like ReLU or LeakyReLU, the derivatives exist everywhere except at the origin. In these cases, we need to carefully investigate if we can use subgradients to define operator norm ||Lp||.
Discussion and Conclusion: We have shown that using a regular neural network as a discriminator in estimating KL divergence results in unreliable estimation if the complexity of the function space is not controlled. We then showed a solution by constructing a discriminator function in RKHS space using neural networks and penalizing its complexity in a scalable way. Although the idea to use RKHS norm to penalize complexity is not new (see for example [1]), it is not clear how to use this idea directly on the function f . In traditional kernel methods, algorithms often do not work with RKHS function f directly, but rather work with kernel matrix, K by using, for example, the Representer Theorem [21]. In the case of big data, working with the big kernel matrix is computationally expensive although some methods have been proposed to speed up the computation, like Random Fourier Feature [22]. We propose a different view by directly constructing a function in RKHS space, which led us to scalable algorithm while incorporating the advantages of neural networks. Moreover, our representation could also be seen as an improvement over RFF by using neural basis, ψ, instead of Fourier basis. The idea of constructing a neural-net function in RKHS and complexity control could also be useful in stabilizing GANs in general. Currently, the most successful way to stabilize GANs is to enforce smoothness by gradient penalization [23, 24, 25]. On the light of the present analysis, gradient penalty could also be thought as a way to control the complexity of the discriminator.

heading: 10 Acknowledgements and Disclosure of Funding
text: We would like to express our deep gratitude to Prof. Linwei Wang and Dr. Prashnna Kumar Gyawali for the helpful discussion at the initial stage of this work. Their discussion was especially important in the inception of this project, in identifying the need for KL estimation from samples and limitation of existing KL estimation approaches. We would like to thank Prof. Dana H. Brooks for the helpful discussions and giving exposure to this work in the early stage. We are grateful to Prof. Octavia Camps for her support, encouragements and resources. We are very thankful to Zulqarnain Khan for his feedback on the early draft of the paper.
Funding disclosure: We are thankful to the support from NIH/NCI R01CA240771 and NIH/NHLBI U01HL089856 .

Review summaries:
This papers studies the problem of estimating KL divergence from samples via variational methods with neural network discriminators. They propose a novel way to construct the discriminator in some RKHS and prove the consistency of the resulting estimator. By penalizing the RKHS norm of the discriminator, they reduce the variance of their estimator and stabilize the training.

The authors consider the problem of KL divergence estimation. Modern machine learning methods for KL divergence estimation make use of neural networks as function approximators. The authors assert that many of the challenges with using such methods is the uncontrolled complexity of the function space. They propose to resolve this by constructing their neural network functions to be RKHS functions, and control the RKHS norm. They prove learning-theoretic results about their method, showing that it is consistent and exhibits controllable fluctuations around its mean (deviation-from-mean error).

An improved method for estimating the KL divergence between two distributions based on samples. To estimate the log density ratio, the authors proposed to train a discriminator while constraining the family of the discriminator function. It constructs the function using the defined in F. Bach (2017), which is a convolution between two network functions. The authors did a careful (though no very thorough) performance analysis of this method, and demonstrated the performance in simple datasets for three different tasks involving the KL divergence.