A Variance Reduction Method for Neural-based Divergence Estimation

1 INTRODUCTION . Divergences such as Kullback-Leibler ( KL ) divergence , f -divergences , Hellinger divergence , αdivergences and Rényi divergences , which were initially developed in the fields of information theory and statistical physics , are indispensable tools in a growing number of machine learning applications . They have been used in adversarial training of generative models ( Goodfellow et al. , 2014 ; Nowozin et al. , 2016 ) , in the estimation of generalization errors ( Esposito et al. , 2021 ) and hypothesis testing ( Broniatowski & Keziou , 2009 ) , to name a few . Mutual information ( MI ) , in particular , which is defined as the KL divergence between the joint distribution of a pair of variables and their marginals ( and can be generalized to divergences other than KL ) , plays a crucial role in Bayesian networks and ( conditional ) independence ( Cheng et al. , 2002 ) , self-supervised learning via contrastive losses ( van den Oord et al. , 2018 ; Le-Khac et al. , 2020 ) as well as in representation learning ( Hjelm et al. , 2019 ; Chen et al. , 2016 ) . Classical divergence estimators perform reasonably well for low dimensional cases , however they scale poorly to large , high dimensional datasets which are typically encountered in modern machine learning . The most compelling estimation approach of a divergence is via the optimization of a lower variational bound parametrized by neural networks . These lower bounds , which are likelihood-free approximations , are maximized in order to compute the divergence value at the optimizer . Wellknown variational representations are the Legendre transformation of an f -divergence ( Broniatowski & Keziou , 2006 ; Nguyen et al. , 2010 ) as well as the Donsker-Varadhan ( DV ) variational formula ( Donsker & Varadhan , 1983 ) for KL divergence and its extension to Rényi divergence ( Birrell et al. , 2020b ) . Their tractability stems from their objective functionals , which are computed from expected values and approximated using statistical averages from the available or generated samples . Despite the scalability and tractability , the estimation of a divergence based on variational formulas is a notoriously difficult problem . One challenge stems from the potentially high bias , since any approximation for the worst case scenario requires an exponential number of samples in order to attain the true divergence value ( McAllester & Stratos , 2020 ) . Additionally , the statistical variance , which scales exponentially with respect to the divergence ’ s value for certain variational estimators ( Song & Ermon , 2019 ) , is often prohibitively high . Focusing on the elevated MI , there are several further lower bounds ( Barber & Agakov , 2003 ; Belghazi et al. , 2018 ; van den Oord et al. , 2018 ; Poole et al. , 2019 ; Guo et al. , 2021 ) and a few upper bounds ( Cheng et al. , 2020 ; Poole et al. , 2019 ) which aim to provide more reliable estimates of MI in the low sample size regime . However , the majority of these MI estimators are not transferable to the general estimation of divergences and frequently produce instabilities during training which are further magnified by the small batch and/or sample size . In this paper , we propose to reduce a divergence estimator ’ s variance via an explicit variance penalty ( VP ) which is added to the objective functional . Our contributions are summarized as follows : • We present a novel variance reduction penalty for f -divergence and expand it via the delta method to the nonlinear setting , including the DV formula for KL divergence as well as the variational formula for the Rényi divergences . The proposed VP is able to flexibly trade off bias and variance . • We present numerical evidence on synthetic datasets that the proposed approach improves both mean squared error ( MSE ) and median absolute error ( MedAE ) in a range of sample sizes and types of divergences . Furthermore , we implemented the proposed VP in several other lower and upper bounds of MI , showing that our variance reduction approach is not restricted to particular variational formulas but it is generic and applicable to the majority of existing variational representations . • When applied to real datasets , we demonstrate the ability of the proposed approach to reduce the variance of the estimated Rényi divergence , thus enabling the detection of rare biological sub-populations which are otherwise difficult to identify . Interestingly , the baseline estimator is unstable when the order value is above one , but it becomes stable when the VP is added . • We also applied the VP to the disentangled representation learning of speech into its text , speaker , and style components . Results on objective evaluation metrics showed that the addition of the VP generally improves the training performance , as much as 18 % relative to the baseline systems . 1.1 RELATED WORK . There are several general-purpose variance reduction techniques in Monte Carlo stochastic sampling , with the most popular approaches being antithetic sampling or more broadly coupling methods , control of variates and importance sampling ( Robert & Casella , 2005 ; Glasserman , 2004 ; Srinivasan , 2013 ) . These methods have not been explicitly applied for the variational divergence estimation problem . We speculate that either they are not applicable due to the unavailability of analytical probability density formulas or they are inefficient ( e.g. , the control of variates approach requires a second estimator and potentially a second parametric model in order to be applied ) . Another way to reduce the variance is to restrict the function space to more smooth and/or controlled test ( or critic ) functions , balancing again between bias and variance . For instance , the restriction to Lipschitz continuous functions has the potential to reduce the variance since there exist favorable concentration inequality results for the Lipschitz space ( Wainwright , 2019 ) . In the GAN literature , Wasserstein GAN ( Gulrajani et al. , 2017 ) and spectral normalization ( Miyato et al. , 2018 ) impose Lipschitz continuity which resulted in signigicant gains in terms of training stability . Similarly , the restriction of test functions to an appropriately designed reproducing kernel Hilbert space could reduce the variance ( Sreekar et al. , 2020 ) . Such approaches can be combined with our proposed variance penalties , as our formulation allows for general test-function spaces . However , we do not focus on this point here . Given the importance of MI , several estimators aim towards improved statistical properties . Lower bounds such as MINE ( Belghazi et al. , 2018 ) , which uses the DV variational formula with an expo- nential moving average , NWJ estimator ( Nguyen et al. , 2010 ) and BA estimator ( Barber & Agakov , 2003 ) as well as upper bounds such as CLUB ( Cheng et al. , 2020 ) still have high variance . InfoNCE ( van den Oord et al. , 2018 ) is one of the few MI estimators that has low variance , but at the cost of either high bias or high computational cost due to the need for many negative samples and thus large batch size . Poole et al . ( 2019 ) and Guo et al . ( 2021 ) aim to clarify the relationships and trade-offs between those variational bounds . A different approach to reducing variance is by appropriately working on the gradients of the objective function ( Wen et al. , 2020 ; 2021 ) . Finally , we discuss the approach of truncating the test function inside a bounded region as proposed in ( Song & Ermon , 2019 ) . The determination of the truncation threshold is quite difficult since it requires an a priori understanding of the log-likelihood ratio . Moreover , a high truncation threshold will not affect the estimation since a high threshold implies no real benefit in terms of variance reduction . On the other hand , a low threshold will result in large bias . Overall , using a high truncation threshold in order to avoid extreme values is a good practice even though it will have a limited impact on variance reduction . 2 BACKGROUND ON VARIATIONAL FORMULAS FOR RÉNYI AND f -DIVERGENCES . While our variance reduction method can be applied to any divergence that possesses a variational formula , here our focus will be on the Rényi and f -divergences , including the KL divergence . For Rényi divergences an appropriate objective functional can be constructed from a difference of cumulant generating functions ( Birrell et al. , 2020b ) Rα ( Q‖P ) = sup g∈Mb ( Ω ) { 1 α− 1 logEQ [ e ( α−1 ) g ] − 1 α logEP [ eαg ] } , α 6= 0 , 1 . ( 1 ) Here Q and P are probability distributions on the set Ω , EQ and EP denote the expectations with respect to Q and P respectively , andMb ( Ω ) is the space of bounded measurable real-valued functions on Ω . For f divergences , f being a lower semicontinuous convex function with f ( 1 ) = 0 , one has the well-known Legendre transform variational formula ( Broniatowski & Keziou , 2006 ; Nguyen et al. , 2010 ) Df ( Q‖P ) = sup g∈Mb ( Ω ) { EQ [ g ] − EP [ f∗ ( g ) ] } , ( 2 ) where f∗ ( y ) = supx∈R { yx − f ( x ) } is the Legendre transform of f . Here and in the following , the function of g that is being optimized will be called the objective functional . Equation ( 2 ) can be generalized to the ( f , Γ ) -divergences ( Birrell et al. , 2020a ) , where Γ ⊂ Mb ( Ω ) is a restricted test-function space DΓf ( Q‖P ) = sup g∈Γ { EQ [ g ] − ΛPf [ g ] } , ( 3 ) ΛPf [ g ] = inf ν∈R { ν + EP [ f∗ ( g − ν ) ] } . ( 4 ) In particular , if fKL ( x ) = x log ( x ) corresponds to the KL divergence then ΛPfKL [ g ] = log ( EP [ exp ( g ) ] ) ≡ Λ P [ g ] ( 5 ) is the classical cumulant generating function and equation ( 3 ) ( with Γ = Mb ( Ω ) ) becomes the Donsker-Varadhan variational formula ( Dupuis & Ellis. , 1997 , Appendix C.2 ) DKL ( Q‖P ) = sup g∈Mb ( Ω ) { EQ [ g ] − logEP [ eg ] } . ( 6 ) For general f , we will often write equation ( 3 ) as DΓf ( Q‖P ) = sup g∈Γ , ν∈R { EQ [ g − ν ] − EP [ f∗ ( g − ν ) ] } ( 7 ) and if Γ is closed under the shifts g 7→ g − ν , ν ∈ R then we can write it simply as DΓf ( Q‖P ) = sup g∈Γ { EQ [ g ] − EP [ f∗ ( g ) ] } . ( 8 ) In particular , if Γ = Mb ( Ω ) then DΓf = Df . The generalizations of Rényi and KL divergence obtained by using a restricted space Γ in place of Mb ( Ω ) in equation ( 1 ) or equation ( 6 ) will be denoted by RΓα and D Γ KL , respectively .

Summary: This paper takes the variational view of the statistical estimation of statistical divergences. Standard variational formulas suffer from high variance when estimated with finite samples, this paper proposes to address the problem by adding a variance penalty regularisation term to the objective functional, thus trading variance for bias. The bias-variance trade-off is controlled via the regularisation coefficient. From this new formulation, the paper derives a new low-variance Neural Divergence Estimation Algorithm with the variance regularisation technique.