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: The paper proposes to adopt a variance penalty term to manipulate the divergence estimators. Experiments first investigated the correlations of the variance penalty and the bias and variance trade-offs. Then, the authors tested the proposed approaches on two real-world application (biological and speech synthesis) datasets. The results show that the proposed approach can effectively reduce the variance as compared to systems without including the variance penalty.