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

1 Introduction . 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 . 2 Related Work . 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 . 3 Reproducing Kernel Hilbert Space . 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 . 4 Problem Formulation and Contribution . 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 ) .

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