Year: 2021

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

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

heading: 1 INTRODUCTION
text: 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.

heading: 2 PRELIMINARY
text: 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.

heading: 2.2 THERMODYNAMIC VARIATIONAL OBJECTIVE
text: 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)

heading: 3.2 THERMODYNAMIC INTEGRATION
text: Using the χ-logarithm function and χ-exponential function defined in equation 8 and equation 9, we can define a exponential family distribution under χ-geometry:
Definition 3. χ-Exponential Family Distribution:
logχ p(x, θ)
c(x) = θ · x− Φ(θ) (12)
or equivalently:
p(x, θ) = c(x) expχ{θ · x− Φ(θ)} (13)
where x is random variable, θ is parameter, Φ(θ) is the partition function, c(x) is the normalization term.
Given two unnormalized distribution π̃0(z) and π̃1(z), we define an unnormalized distribution that is the log χ interpolation between π̃0(z) and π̃1(z):
logχ(π̃β(z)) = (1− β) logχ(π̃0(z)) + β logχ(π̃1(z)) (14) Its normalized distribution is defined as:
πβ(z) = π̃β(z)
Zβ (15) where Zβ = ∫ π̃β(z)dz. It’s easy to show that: π̃β(z) is on the log χ geodesic between π̃0(z) and π̃1(z) with χ-Exponential Family Distribution as: Theorem 2. πβ is also an χ-exponential family distribution, no matter π̃0(z) and π̃1(z).
Proof.
π̃β(z) = expχ((1− β) logχ π̃0(z) + β logχ π̃1(z)) = expχ(logχ π̃0(z) + β(logχ π̃1(z)− logχ π̃0(z))) (16)
By defining the potential energy function: Uβ(z) = logχ π̃β(z), then we can estimate the logχ of the ratio of the normalizing constants via thermodynamic integration:
logχ Z1 − logχ Z0 = ∫ 1 0 ∫ χ(Zβ) χ(π̃β(z)) (logχ π̃0(z)− logχ π̃1(z))dzdβ (17)
Notice that variational inference can be connected with thermodynamic integration if we set: π̃0(z) = q(z|x), Z0 = ∫ q(z|x)dz = 1
π̃1(z) = p(x, z), Z1 =
∫ p(x, z)dz = p(x) (18)
then the logχ of model evidence is the thermodynamic integration (TI):
logχ p(x) = ∫ 1 0 Sβdβ (19)
where Sβ = ∫ χ(Zβ) χ(π̃β(z)) (logχ q(z|x)− logχ p(x, z)))dz
Recall the relation between f -divergence and χ function in equation 10, then logχ p(x) = f ∗(p(x)) is dual function f of the model evidence, which corresponds to the f -VI in (Wan et al., 2020).
To approximate f -TVO defined in equation 17, we pick the a list of ordered partitions β = [β0, β1, · · · , βK ], where β0 ≤ β1 ≤ · · · ≤ βK , β0 = 0, βK = 1, then our f -TVO is the left Riemannian approximation of TI:
f -TVOL(K) = K∑ k=1 ∆kSβk−1 , (20)
where ∆k = βk − βk−1.

heading: 3.3 COMPUTATION
text: In this section, we show how to efficiently compute the f -TVO.
We first rewrite Sβ in equation 17 into
Sβ =
∫ χ(Zβ)
χ(π̃β(z)) (logχ π̃0(z)− logχ π̃1(z))dz
= Eq(z)[ χ(Zβ)
χ(π̃β(z)) (logχ π̃0(z)− logχ π̃1(z))/q(z|x)] (21)
We apply the reparameterization trick to estimate Sβ . We assume there is a mapping gφ( ) that satisfies: z = gφ( ). Then the expectation of arbitrary function F (z) over distribution q(z) can be computed as Eq(z)[F (z)] = E [F (gφ( ))]. In particular, we use the prevalent Gaussian reparameterization: z ∼ N (µ,Σ), i.e. z = µ + Σ 12 , ∼ N (0, I). Then Sβ can be reformulated as:
Sβ = E [ χ(Zβ)
p( )χ(π̃β(gφ( ))) (logχ p( )− logχ p(x, gφ( )))] (22)
where we define q(z|x) = p( ) and p(x, z) = p(x, gφ( )) in equation 19. Empirically, we can compute Sβ as:
Ŝβ = 1
N N∑ n=1 [wn(logχ p( n)− logχ p(x, gφ( n)))] (23)
where wn = χ( ∑N n=1 π̃β(gφ( n)))/p( n)χ(π̃β(gφ( n))).
In practice, we split the computation of Ŝβ into two parts: Ŝβ = L−R:
L = 1
N N∑ n=1 [wn logχ p( n)]
R = 1
N N∑ n=1 [wn logχ p(x, gφ( n))] (24)
This is because we use log p(x, gφ( n)) and log p( n)) during our computation to avoid numerical overflow or underflow. All the intermediate sum operations are performed by logsumexp operation until we get logL and logR, then we use exponential operation to complete the computation.

heading: 4 GENERALIZING THERMODYNAMIC OBJECTIVES
text: Recall f -TVO in equation 20 is left Riemannian approximation of Thermodynamic integration. Similarly, We can also have right Riemannian approximation of f -TI:
f -TVOR(K) = K∑ k=1 ∆kSβk (25)
Both left and right Riemannian approximation corresponds to existing variational bounds. For example, left single Riemannian approximation of the KL-TVO equals the ELBO, while the right single Riemannian approximation equals to the EUBO. Also, it is easy to verify if the dual function f∗ is increasing (decreasing), then left Riemannian approximation is an upper (lower) bound of Thermodynamic integration, while right Riemannian approximation is a lower (upper) bound of Thermodynamic integration.
In addition to the left or right Riemannian approximation of Thermodynamic integration, we can also define a “zig-zag” (z) approximation of Thermodynamic integration, which combines left and right Riemannian approximation. As shown in FIg. 2, the “zig-zag” Thermodynamic objective can be formulated as:
f -TVOZ = K∑ k=1 ∆k[1(k%2 = 1)Sβk−1 + 1(k%2 = 0)Sβk ] (26)
where 1(A) is indicator function, i.e. 1(A) = 1 when condition A is true, 1(A) = 0 when condition A is false.

heading: 5 EXPERIMENTS
text: In this section, we evaluate the effectiveness and the wide applicability of f -TVO on two tasks. f -TVO is first used as a surrogate optimization objective of f -VI in a Bayesian neural network for linear regression, then used in a VAE for image reconstruction and generation. For both tasks, we use Adam as our optimizer with recommended parameters in (Da, 2014). For all experiments, we use even partition of the temperature for Riemannian integration.

heading: 5.1 BAYESIAN VARIATIONAL AUTOENCODER
text: f -TVO is used in Bayesian VAE for image reconstruction task on datasets including Caltech 101 Silhouettes (Cal), Frey Face (Fre), MNIST (MNI), and Omniglot (Omn). We have replaced the conventional ELBO loss function of VAE with the proposed f -TVOs that correspond to flexible f - divergences. We test and compare the f -TVO VAEs associated with three well-known f -divergences (KL-divergence, Renyi’s α-divergence with α = 3, and χ-divergence with n = 2). We also try a customized divergence which is defined by its dual function f∗c1(t) = log 2 t+ log t.
We use the left Riemannian approximation with 5 partitions, i.e. f -TVOL(5) in our setting. As shown in Table 2, the proposed f -TVOs outperforms the corresponding baselines in most cases.
Effects of numbers of partitions. Here, we study the effects of different numbers of partitions on Bayesian variational autoencoder in our experiments. We try partition numbers K = 1, 2, 5, 10 in the f -TVO, respectively. As shown in Fig. 3, we can see that even partition number is K = 2 can improve the baseline performance.
Effects of f -TVO variants. We also study the effects of f -TVO variants defined in Sec. 4. We try all 3 variants f -TVOL,f -TVOR and f -TVOZ on Bayesian variational autoencoder in our experiments. As shown in Fig. 4, three f -TVO variants achieve similar performance.

heading: 5.2 BAYESIAN NEURAL NETWORKS
text: Our experiments of Bayesian neural network are performed on regression tasks. We generally follow the experiments settings as in (Wan et al., 2020). In particular, We use twelve datasets that are collected from the UCI dataset repository, where each dataset is randomly split into 90%/10% for training and testing. The model is a single-layer neural network with 50 hidden units (ReLUs). We use θ ∼ N (θ; 0, I) as a Gaussian prior of the network weights and Gaussian approximation to the true posterior.
In the experiments, we try different f -TVOs, including corresponds to KL TVO, α-divergence TVO, where α = 3, χ-square divergence TVO. Following (Wan et al., 2020), we also define a custom divergence, which is defined by its dual function fc2 = f̃∗(t) − f̃∗(1), where f̃∗(t) =
−1/6 · (log t + t0)3 − 1/2 · (log t + t0) − 1. We use the left Riemannian approximation with 5 partitions, i.e. f -TVOL(5) in our setting. As shown in Table 3, the proposed f -TVOs outperform the corresponding baselines.

heading: 6 CONCLUSION
text: In this paper, we have proposed a general f -TVO framework that is a tighter evidence bounds than f -VI. f -TVO unifies all differential f -divergence into TVO from a χ-geometry perspective. By connecting f -divergence with χ-geometry, f -TVO is the Riemanian sum of thermodynamic
integration along the geodesic between q(z|x) and p(x, z) under χ-geometry. Empirical experiments of some instances of f -TVO on the popular benchmarks show the superior performance than stateof-the art results, which imply the flexibility and effectiveness of f -TVO. Future work on f -TVO may include more efficient f -TVO optimization methods, and more explorations on the properties of χ-geometry.

Review summaries:
The paper proposes a novel f-divergence Thermodynamic Variational Objective (f-TVO) framework for VI, that extends the TVO towards, a more general, family of f-divergences. The authors propose to use a $\chi$-deformed exponential distribution, which casts the f-TVO objective as integral along the $\chi$-path between p(x,z) and q(z|x) (rather than the geometric path in TVOs) under $\chi$-geometry. The authors propose different variants of f-TVO, that vary between the type of the f-divergance used, as well as how this integral is approximated (K-partitioned ) left-Riemann sum (related to ELBO), right-Riemann sum (related to EUBO) or a 'zig-zag' that alternates between the two. Besides theoretical justifications, results from two sets of experiments show that, in general, the proposed f-TVO perform comparable to or slightly better than the f-VI counterparts, but without clear conclusion wrt the choice of the f-divergence.

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.

This paper proposed an $f$-divergence TVO, which includes some existing works e.g., RVB, CUBO, ELBO into a unified framework. This paper's main idea is to transform $f$-divergence into a generalized $\chi$-exponenetial family and integral TVO along the $\chi$-path. The paper provides some theoretical analysis and the optimization methods of the suggested framework and supports the proposed $f$-TVO with numerical results. 

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.