$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: The paper presents a new bound as the objective for variational inference. The bound combines the recent progress of thermodynamic variational object which gives a tighter bound than the conventional ELBO, and the f-divergence which induces more possible distribution metrics. Experiments show its better performance on some Bayesian inference tasks.