$f$-Divergence Thermodynamic Variational Objective: a Deformed Geometry Perspective

In this paper , we propose a f -divergence Thermodynamic Variational Objective ( f -TVO ) . f -TVO generalizes the Thermodynamic Variational Objective ( TVO ) by replacing Kullback–Leibler ( KL ) divergence with arbitary differeitiable f - divergence . In particular , f -TVO approximates dual function of model evidence f∗ ( p ( x ) ) rather than the log model evidence log p ( x ) in TVO . f -TVO is derived from a deformed χ-geometry perspective . By defining χ-exponential family exponential , we are able to integral f -TVO along the χ-path , which is the deformed geodesic between variational posterior distribution and true posterior distribution . Optimizing scheme of f -TVO includes reparameterization trick and Monte Carlo approximation . Experiments on VAE and Bayesian neural network show that the proposed f -TVO performs better than cooresponding baseline f -divergence variational inference . 1 INTRODUCTION . Variational inference ( VI ) is a core technique in probabilistic machine learning ( Mnih & Gregor , 2014 ; Mnih & Rezende , 2016 ; Tucker et al. , 2017 ; Maddison et al. , 2017 ; Paisley et al. , 2012 ; Salimans et al. , 2013 ; Ranganath et al. , 2014 ; Kucukelbir et al. , 2015 ) . Normal Variational inference tries to perform a tractable posterior distribution approximation for log model evidence estimation log p ( x ) . Commonly , this is done by introducing a divergence D [ q ( z|x ) ||p ( z|x ) ] between variational distribution q ( z|x ) and true posterior p ( z|x ) as a regularization term . Then the evidence lower bound ( ELBO ) : log p ( x ) −D [ q ( z|x ) ||p ( z|x ) ] is used as a surrogate optimization objective for maximum likelihood parameter estimation ( Hoffman & Johnson , 2016 ; Duan et al. , 2017 ; Yang , 2017 ) . Recently , there are serveral directions that generalize normal variational inference . Interested readers can refer to VI ’ s background and progression in ( Blei et al. , 2017 ; Zhang et al. , 2018 ) . One direction is to explore tighter variational bounds ( Burda et al. , 2016 ; Chen et al. , 2018 ; Masrani et al. , 2019 ; Brekelmans et al. , 2020 ) . Among these efforts , the recent Thermodynamic Variational Objective ( TVO ) approach connects variational inference with thermodynamic integration . In paricular , TVO interprets the log model evidence estimation as 1D integral along the geodesic path between joint distribution of data and hidden variable p ( x , z ) and variational hidden variable distribution q ( z|x ) . By properly choosing intergral intervals of Riemannian sum approximations , the TVO approximates the log likelihood with a series of upper and lower bounds , which are tighter surrogates than the normal variational inference objectives for model parameter likelihood maximization ( Masrani et al. , 2019 ) . Another direction is trying to extend the normal VI framework by replacing KL divergence with other statistical divergences . Among these efforts , Rényi ’ s α-divergence and χ-divergence can be regarded as the root divergences , which makes many existing divergence as special cases including KL-divergence . Rényi ’ s α-VI ( Li & Turner , 2016 ) and χ-VI ( Dieng et al. , 2017 ) are stochastic VI that generalized from stochastic VI ( Hoffman et al. , 2013 ) and black-box VI ( Ranganath et al. , 2014 ) , and outperform the classical KL-VI in Bayesian regressions and image reconstruction on some benchmarks . More generally , recent work f -divergence variational inference ( f -VI ) ( Wan et al. , 2020 ) gives a general framework that includes all divergences from f -divergence family . The variational objective of f -divergence VI can be regarded as the surrogate bound of dual function of model evidence f∗ ( p ( x ) ) . In this paper , we combine the above two directions together by proposing f -TVO that generalizes the original TVO to arbitary f -divergence in a unified framework . f -TVO approximates dual function of model evidence f∗ ( p ( x ) ) rather than the log model evidence log p ( x ) in TVO . The estimation of f -TVO is established from a deformed χ-geometry prospective . By defining χ-exponential family distribution , the computation of f -TVO is integrated over the χ-path , which is the geodesic between p ( x , z ) and q ( z|x ) under χ-geometry . The proposed f -TVO is a tighter bound of dual function of model evidence f∗ ( p ( x ) ) rather than the normal variational lower bound in TVO . Our contributions include the following : – We enrich the TVO by proposing a unified f-TVO framework that compatible with an amount of existing divergence , such as KL , α divergence , χ divergence ; – We derive a unified f-TVO bound equiped with the upper/lower bound criteria that improves many existing bounds , such as ELBO , CUBO , RVB ; – Experiments show that our proposed f-TVO is a better optimization surrogate than corre- sponding f-VI . The rest of the paper is organized as follows : Sec . 2 reviews preliminary knowledges for f -divergence and Thermodynamic Variational Objective , Sec . 3 presents our f -TVO , Sec . 4 extends our approach to more variants of f -TVO , and Sec . 5 presents our experiment results . 2 PRELIMINARY . In this section , we will review the preliminary knowledge for f -divergence and Thermodynamic Variational Objective . 2.1 f -DIVERGENCE f -divergence is a measures of difference between two probability distributions . It can be defined as follows ( Sason & Verdú , 2016 ) . Definition 1 . Given a convex function f with f ( 1 ) = 0 , then f -divergence from probability distribution q ( z ) to p ( z ) is defined as : Df ( p||q ) = ∫ f ( q ( z ) p ( z ) ) p ( z ) dz = Ep [ f ( q ( z ) p ( z ) ) ] ( 1 ) Many existing well-know divergence can be regarded as special cases of f -divergence by properly choose the generation function f . Table 1 shows the relationship between the sample divergences and their corresponding generation function f . Wan et al . ( Wan et al. , 2020 ) proposed a f -divergence variational inference ( f -VI ) that generalizes normal variational inference to all f-divergences in a unified framework . f -VI framework is primarily based on reverse f -divergence ( Wan et al. , 2020 ) . One can connect forward f -divergence Df ( p||q ) with the reverse f -divergence Df ( q||p ) via dual function f∗ as follows . Definition 2 . Given a function f , its dual funtion1 f∗ is defined as : f∗ ( t ) = t · f ( 1/t ) ( 2 ) The dual function f∗ has the following properties : ( 1 ) ( f∗ ) ∗ = f , ( 2 ) f∗ is convex iff f is convex , ( 3 ) f∗ ( 1 ) = 0 iff f ( 1 ) = 0. f -divergence in equation 1 can be reformulated with the dual function f∗ : Df ( p||q ) = ∫ f ( q ( z ) p ( z ) ) p ( z ) dz = ∫ f∗ ( p ( z ) q ( z ) ) q ( z ) dz = Df∗ ( q||p ) ( 3 ) As shown in the theorem below , f -variational bound is an approximation of the dual function f∗ of the model evidence p ( x ) : Theorem 1 . The dual function of model evidence f∗ ( p ( x ) ) is bounded by the f -variational bound Lf ( q , x ) Lf ( q , x ) = Eq ( z|x ) [ f∗ ( p ( z , x ) q ( z|x ) ) ] ≥ f∗ ( p ( x ) ) ( 4 ) where equality is obtained when q ( z|x ) = p ( z|x ) . The proof of the theorem is straightforward as f∗ is convex ( Wan et al. , 2020 ) . Many well-known variation bounds can be represented in the f -VI framework . For example , f -VI recovers to ELBO when f = t log t , and f -VI recovers to CUBOn when f = t1−n − t. More details are shown in the Table 1 . 2.2 THERMODYNAMIC VARIATIONAL OBJECTIVE . The evidence lower bound ( ELBO ) is a lower bound of the log evidence of a generative model p ( x , z ) . It can be formulated as the difference between log evidence of model evidence log p ( x ) and KL divergence of variational posterior distribution q ( z|x ) and true posterior distribution p ( z|x ) : ELBO = log p ( x ) −DKL ( q ( z|x ) ||p ( z|x ) ) ( 5 ) The Thermodynamic Variational Objective ( TVO ) gives a tighter bound of log p ( x ) than ELBO . As shown in Fig . 1 , TVO is a Riemannian sum of integral along the geometric path between the model and inference network : 1 K ( ELBO + K−1∑ k=1 Eπβk [ log p ( x , z ) q ( z|x ) ] ) ︸ ︷︷ ︸ TVO ≤ ∫ 1 0 Eπβ [ log p ( x , z ) q ( z|x ) ] dβ = log p ( x ) ( 6 ) 1Note : dual function here is different from Fenchel conjugate function : fc ( y ) = supx∈dom f ( y > x− f ( x ) ) . where β = [ β0 , β1 , · · · , βK ] is a partition between [ 0,1 ] , where β0 = 0 , and βK = 1 , πβ is the normalized geometric combination of p ( x , z ) and q ( z|x ) , i.e . π̃β = q ( z|x ) 1−βp ( x , z ) β and πβ = π̃β/Zβ , where Zβ is a normalization constant for each β . One can verify that πβ is a member of exponential family distribution , no matter what q ( z|x ) and p ( x , z ) are : πβ ( z|x ) = π0 ( z|x ) exp { β · T ( x , z ) − φ ( x ; β ) } ( 7 ) where T ( x , z ) = log p ( x , z ) q ( z|x ) , φ ( x ; β ) is the log partition function that normalize over z . 3 χ-THERMODYNAMIC VARIATIONAL OBJECTIVE In this section , we establish our theory for f -Thermodynamic Variational Objective ( f -TVO ) . We first connect f -divergence with χ-geometry . Then , we present f -TVO as an integration along the χ-path . Lastly , we specify the computation process of f -TVO and its gradients . 3.1 χ-LOGARITHM AND χ-EXPONENTIAL We first define a generalized χ-logarithm function for a positive non-increasing function χ : logχ ( u ) = ∫ u 1 χ ( v ) dv ( 8 ) which is different from the definition in ( Naudts , 2011 ; Amari , 2016 ) . Notice that equation 8 gives normal log function log u when χ ( v ) = 1v . The inverse of the χ-logarithm is called χ-exponential : expχ ( u ) = 1 + ∫ u 0 ψ ( v ) dv ( 9 ) where ψ is an auxiliary function . One can determine ψ for each logχ ( u ) , e.g . ψ ( v ) = exp ( v ) when logχ ( u ) = log u. χ-Logarithm and χ-Exponential have the following properties : ( 1 ) d du expχ ( u ) = ψ ( u ) , ( 2 ) u = expχ ( logχ ( u ) ) , ( 3 ) 1 χ ( u ) = ψ ( logχ ( u ) ) . For any divergence with convex diffrentiable generation function f with f ( 0 ) = 1 , there exists a continuous positive non-increasing function χ that connects the dual function f∗ : f∗ ( u ) = − logχ ( u ) = − ∫ u 1 χ ( v ) dv ( 10 ) then f -divergence can be represented as Df ( q||p ) = Df∗ ( p||q ) = − ∫ p ( z ) · logχ ( q ( z ) p ( z ) ) dz ( 11 )

Summary: This paper proposed new variational inference that combines f-divergence variational inference and the thermodynamic variational objective. The authors introduced several new concepts of exponential families to extend TVO. Finally, the authors provided the estimator of the gradient of the objective function based on the reparametrization trick.