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