Loss function based second-order Jensen inequality and its application to particle variational inference

Abstract:
Bayesian model averaging, obtained as the expectation of a likelihood function by a posterior distribution, has been widely used for prediction, evaluation of uncertainty, and model selection. Various approaches have been developed to efficiently capture the information in the posterior distribution; one such approach is the optimization of a set of models simultaneously with interaction to ensure the diversity of the individual models in the same way as ensemble learning. A representative approach is particle variational inference (PVI), which uses an ensemble of models as an empirical approximation for the posterior distribution. PVI iteratively updates each model with a repulsion force to ensure the diversity of the optimized models. However, despite its promising performance, a theoretical understanding of this repulsion and its association with the generalization ability remains unclear. In this paper, we tackle this problem in light of PAC-Bayesian analysis. First, we provide a new second-order Jensen inequality, which has the repulsion term based on the loss function. Thanks to the repulsion term, it is tighter than the standard Jensen inequality. Then, we derive a novel generalization error bound and show that it can be reduced by enhancing the diversity of models. Finally, we derive a new PVI that optimizes the generalization error bound directly. Numerical experiments demonstrate that the performance of the proposed PVI compares favorably with existing methods in the experiment.

heading: 1 Introduction
text: Bayesian model averaging (BMA) has been widely employed for prediction, evaluation of uncertainty, and model selection in Bayesian inference. BMA is obtained as the expectation of a likelihood function by a posterior distribution and thus it contains information of each model drawn from the posterior distribution [21]. Since estimating the posterior distribution is computationally difficult in practice, various approximations have been developed to efficiently capture the diversity in the posterior distribution [21, 1, 2].
One of these recently proposed approaches involves optimizing a set of models simultaneously with interaction to ensure the diversity of the individual models, similar to ensemble learning. One notable example is particle variational inference (PVI) [17, 28], which uses an ensemble as an empirical approximation for the posterior distribution. Such PVI methods have been widely employed in variational inference owing to their high computational efficiency and flexibility. They iteratively update the individual models and the update equations contain the gradient of the likelihood function and the repulsion force that disperses the individual models. Thanks to this repulsion term, the obtained ensemble can appropriately approximate the posterior distribution. When only one model is
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
used in PVI, the update equation is equivalent to that of the maximum a posteriori (MAP) estimation. Other methods have been developed apart from PVI, especially for latent variable models, which have introduced regularization to the MAP objective function to enforce the diversity in the ensemble. A notable example of such methods is the determinantal point process (DPP) [29].
Despite successful performances of these methods in practice [22, 18, 17, 28, 5, 27], a theoretical understanding of the repulsion forces remains unclear. Some previous studies considered PVI as a gradient flow in Wasserstein space with an infinite ensemble size [16, 12] and derived the convergence theory. However, an infinite ensemble size is not a practical assumption and no research has been conducted to analyze the repulsion force related to the generalization.
BMA can be regarded as a special type of ensemble learning [26], and recent work has analyzed the diversity of models in ensemble learning in light of the PAC-Bayesian theory [19]. They reported that the generalization error is reduced by increasing the variance of the predictive distribution. However, the existing posterior approximation methods, such as PVI and DPPs, enhance the diversity with the repulsion of the parameters or models rather than the repulsion of the predictive distribution. We also found that the analysis in a previous work [19] cannot be directly extended to the repulsion of parameters or loss functions (see Appendix F). In addition, when the variance of the predictive distribution is included in the objective function in the variational inference, the obtained model shows large epistemic uncertainty, which hampers the fitting of each model to the data (see Section 5).
Based on these findings, this study aims to develop a theory that explains the repulsion forces in PVI and DPPs and elucidates the association of the repulsion forces with the generalization error. To address this, we derive the novel second-order Jensen inequality and connect it to the PAC-Bayesian generalization error analysis. Our second-order Jensen inequality includes the information of the variance of loss functions. Thanks to the variance term, our bound is tighter than the standard Jensen inequality. Then, we derive a generalization error bound that includes the repulsion term, which means that enhancing the diversity is necessary to reduce the generalization error. We also show that PVI and DPPs can be derived from our second-order Jensen inequality, and indicate that these methods work well from the perspective of the generalization error. However, since these existing methods do not minimize the generalization error upper bound, there is still room for improvement. In this paper, we propose a new PVI that directly minimize the generalization error upper bound and empirically demonstrate its effectiveness.
Our contributions are summarized as follows:
1. We derive a novel second-order Jensen inequality that inherently includes the variance of loss functions. Thanks to this variance term, our second-order Jensen inequality is tighter than the standard Jensen inequality. We then show that enhancing the diversity is important for reducing the generalization error bound in light of PAC-Bayesian analysis.
2. From our second-order Jensen inequality, we derive the existing PVI and DPPs. We demonstrate that these methods work well even at a finite ensemble size, since their objective functions includes valid diversity enhancing terms to reduce the generalization error.
3. We propose a new PVI that minimizes the generalization error bound directly. We numerically demonstrate that the performance of our PVI compares favorably with existing methods.

heading: 2 Background
text: In this section, we briefly review PVI, DPPs, and PAC-Bayesian analysis.

heading: 2.1 Particle variational inference
text: Assume that training dataset D = (x1, . . . , xD) is drawn independently from unknown data generating distribution ν(x). Our goal is to model ν(x) by using a parametrized model ln p(x|θ), where θ ∈ Θ ⊂ Rd. We express p(D|θ) = ∑Dd=1 ln p(xd|θ) and assume a probability distribution over parameters. In Bayesian inference, we incorporate our prior knowledge or assumptions into a prior distribution π(θ). This is updated to a posterior distribution p(θ|D)∝p(D|θ)π(θ), which incorporates the observation D. Let us consider the approximation of p(θ|D) with q(θ). We predict a new data point by a predictive distribution p(x) =Eq(θ)p(x|θ), where E denotes the expectation. This expectation over the posterior is often called BMA [21].
Assume that we draw N models from the posterior distribution and calculate BMA. We denote those drawn models as an empirical distribution ρE(θ) = 1N ∑N i=1 δθi(θ), where δθi(θ) is the Dirac distribution that has a mass at θi. We also refer to these N models as particles. The simplest approach to obtain these N particles is MAP estimate that updates parameters independently with gradient descent (GD) as follows [28]:
θnewi ←− θoldi + η∂θ ln p(D|θoldi )π(θoldi ), (1) where η ∈ R+ is a step size. In BMA, we are often interested in the multi-modal information of the posterior distribution. In such a case, MAP estimate is not sufficient because N optimized particles do not necessarily capture the appropriate diversity of the posterior distribution. Instead, particle variational inference (PVI) [17, 28] approximates the posterior through iteratively updating the empirical distribution by interacting them with each other:
θnewi ←− θoldi + ηv({θoldi′ }Ni′=1), (2) where v({θ}) is the update direction and explicit expressions are summarized in Table 1. Basically, v is composed of the gradient term and the repulsion term. In Table 1, the repulsion terms contain the derivative of the kernel function K, and the Gaussian kernel [23] is commonly used. When the bandwidth of K is h, the repulsion term is expressed as ∂θiK(θi, θj)=−h−2(θi−θj)e−(2h
2)−1‖θi−θj‖2 , where ‖ · ‖ denotes the Euclidean norm. We refer to this as a parameter repulsion. Note that the repulsion term depends on the distance between particles, and the closer they are, the stronger force is applied. This moves θi away from θj , and thus particles tend not to collapse to a single mode.
For over-parametrized models such as neural networks, since the repulsion in the parameter space is not enough for enhancing the diversity, function space repulsion force for supervised tasks was developed [28]. We call it function space PVI (f-PVI). Pairs of input-output data are expressed as D = {(xd, yd)}Dd=1. We consider the model p(y|x, θ) = p(y|f(x; θ)) where f(x; θ) is a cdimensional output function parametrized by θ and x is an input. Furthermore, we consider the distribution over f and approximate it by a size-N ensemble of f , which means that we prepare N parameters (particles) {θi}Ni=1. We define fi(x) := f(x; θi). When we input the minibatch with size b into the model, we express it as fi(x1:b) = (fi(x1), . . . , fi(xb)) ∈ Rcb. Then the update equation is given as
θnewi ←− θoldi + η ∂fi(x1:b)
∂θi ∣∣∣ θi=θoldi v({fi(x1:b)}Ni=1), (3)
where v({fi}) is obtained by replacing ln p(D|θ)π(θ) with (D/b) ∑b d=1 ln p(yd|fi(xd))π(fi), where π(f) is a prior distribution over f and the Gram matrix K(θi, θj) is replaced with K(f bi , f b j ) in Table 1. See appendix C.1 for details. Then, f-PVI modifies the loss signal so that models are diverse. We refer to the repulsion term of f-PVI as a model repulsion. We express f bi := fi(x1:b) for simplicity. When we use the Gaussian kernel, the model repulsion is expressed as
∂θiK(f b i , f b j ) = −h−2(f bi − f bj )e−‖f b i−f b j ‖ 2/(2h2)∂θif b i . (4)
Thus, the model repulsion pushes model fi away from fj .

heading: 2.2 Regularization based methods and determinantal point processes
text: Another common approach for enhancing the diversity for latent variable models is based on regularization. A famous example is the determinantal point process (DPP) [29], in which we maximize
EρE ln p(D|θ)π(θ) + ln detK, (5) where K is the kernel Gram matrix defined by Ki,j = K(θi, θj). This log-determinant term is essentially a repulsion term that enhances the diversity in the parameter space.

heading: 2.3 PAC-Bayesian theory
text: Here, we introduce PAC-Bayesian theory [7]. We define the generalization error as the cross-entropy:
CE := Eν(x)[− lnEq(θ)p(x|θ)], (6)
which corresponds to the Kullback-Leibler (KL) divergence. Our goal is to find q(θ) that minimizes the above CE. In many Bayesian settings, we often minimize not CE but a surrogate loss [19] that is obtained by the Jensen inequality:
Eν(x)[− lnEq(θ)p(x|θ)] ≤ Eν(x),q(θ)[− ln p(x|θ)]. (7) Since the data generating distribution is unknown, we approximate it with a training dataset as Eν(x),q(θ)[− ln p(x|θ)] ≈ Eq(θ) 1D ∑D d=1[− ln p(xd|θ)]. The PAC-Bayesian generalization error analysis provides the probabilistic relation for this approximation as follows: Theorem 1 (Germain [7]). For any prior distribution π over Θ independent of D and for any ξ ∈ (0, 1) and c > 0, with probability at least 1− ξ over the choice of training data D ∼ ν⊗D(x), for all probability distributions q over Θ, we have
Eν(x),q(θ)[− ln p(x|θ)] ≤ Eq(θ) 1
D D∑ d=1 [− ln p(xd|θ)] + KL(q, π) + ln ξ−1 + Ψπ,ν(c,D) cD , (8)
where Ψπ,ν(c,D) := lnEπED∼ν⊗D(x)exp[cD(−Eν(x) ln p(x|θ) +D−1 ∑D d=1 ln p(xd|θ))].
The Bayesian posterior is the minimizer of the right-hand side of the PAC-Bayesian bound when c = 1. Recently, the PAC-Bayesian bound has been extended so that it includes the diversity term [19]. Under the same assumptions as Theorem 1 and for all x, θ, p(x|θ) <∞, we have
CE≤−Eν,q[ln p(x|θ) + V (x)] ≤ −Eq 1
D D∑ d=1 [ln p(xd|θ) + V (xd)] + KL(q, π) + ln ξ−1+Ψ′π,ν(c,D) 2 cD , (9)
where V (x) := (2 maxθ p(x|θ)2)−1Eq(θ) [ (p(x|θ)− Eq(θ)p(x|θ))2 ] is the variance of the predictive distribution and Ψ′π,ν(c,D) is the modified constant of Ψπ,ν(c,D) (see Appendix C.2 for details). A similar bound for ensemble learning, that is, q(θ) as an empirical distribution, was also previously proposed [19] (see Appendix C.3). This bound was derived directly from the second-order Jensen inequality derived in another work [14]. Furthermore, the diversity comes from the variance of the predictive distribution, which is different from PVI and DPPs because their repulsion is in the parameter or model space. Note that we cannot directly change the variance of the predictive distribution to that of PVI or DPPs because it requires an inequality that is contrary to the Jensen inequality. We also found that directly optimizing the upper bound of Eq.(9), referred to as PAC2E, results in a too large variance of the predictive distribution which is too pessimistic for supervised learning tasks (see Section 5).

heading: 3 Method
text: Here, we derive our novel second-order Jensen inequality based on the variance of loss functions and then derive a generalization error bound. Then, we connect our theory with existing PVI and DPPs.

heading: 3.1 A novel second-order Jensen inequality
text: First, we show the second-order equality, from which we derive our second-order Jensen inequality.
Theorem 2. Let ψ be a twice differentiable monotonically increasing concave function on R+, Z be a random variable on R+ that satisfies E[Z2] <∞, and its probability density be pZ(z). Define a constant µ := ψ−1(E[ψ(Z)]). Then, we have
E[Z] = µ− ( 2 dψ(µ)
dz )−1 ∫ R+ [ d2ψ(c(z)) dz2 (z − µ)2 ] pZ(z)dz, (10)
where c(z) is a constant between z and µ that is defined from the Taylor expansion (see Appendix D.1 for details).
Proof sketch: There exists a constantc(z)betweenz andµs.t. ψ(z)=ψ(µ)+ dψ(µ)dz (z−µ)+ 12 d2ψ(c(z)) dz2 (z−µ)2 from the Taylor expansion. Then we take the expectation. Full proof is given in Appendix D.1.
This theorem states the deviation of EZ from µ when ψ is applied to Z. By setting ψ(·)= ln(·) and Z=p(x|θ) and applying ln to both hand sides of Eq.(10), we have the following equality: Corollary 1. If for all x and θ, p(x|θ) <∞, we have
Eq(θ) ln p(x|θ)=lnEq(θ)p(x|θ)− ln ( 1 + Eq(θ)(2g(θ, x)2)−1(eln p(x|θ)−eEq(θ) ln p(x|θ))2 ) , (11)
where g(θ, x) is a constant between p(x|θ) and eEq(θ) ln p(x|θ)that is defined from the Taylor expansion (see Appendix D.2 for details). Remark 1. Recall that the standard Jensen inequality is Eq ln p(x|θ)≤ lnEqp(x|θ), and its gap is called the Jensen gap. In Eq.(11), the second term of the right-hand side is always positive. Thus, this term corresponds to the Jensen gap. Remarkably when we use the standard Jensen inequality, this information is lost. We clarify the meaning of this term below. Also note that our second-order equality is different from those of the previous works [19, 14] (see Appendix F for details).
Next, we show our first main result, loss function based second-order Jensen inequality: Theorem 3. Under the same assumption as Corollary 1,
Eq(θ) ln p(x|θ) ≤ lnEq(θ)p(x|θ)− Eq(θ) (
ln p(x|θ)− Eq(θ) ln p(x|θ) 2h(x, θ) )2 ︸ ︷︷ ︸
:= R(x, h)
, (12)
where h(x, θ)−2 = exp ( ln p(x|θ) + Eq(θ) ln p(x|θ)− 2 max
θ ln p(x|θ)
) . (13)
Proof sketch: Apply √ αβ≤ α−βlnα−ln β for any α, β> 0 to Eq.(11). Full proof is given in Appendix D.3.
Remark 2. R is the weighted variance of loss functions, and it is always positive. Thus, this inequality is always tighter than the Jensen inequality, and the equality holds if the weighted variance is zero. Compared to the results of the previous works [19, 14] that used the predictive variance in the inequality, our bound focuses on the variance of loss functions. We refer to our repulsion term R as a loss repulsion.
Then, by rearranging Eq.(12) and taking the expectation, we have the following inequality:
CE ≤ −Eq(θ),ν(x)[ln p(x|θ)]− Eν(x)R(x, h) ≤ −Eq(θ),ν(x)[ln p(x|θ)]. (14) Using this inequality, we obtain the second-order PAC-Bayesian generalization error bound: Theorem 4. (See Appendix D.4 for the complete statement) Under the same notation and assumptions as Theorems 1 and 3, with probability at least 1− ξ, we have
CE≤−Eν,q[ln p(x|θ)+R(x,h)]≤−Eq 1
D D∑ d=1 [ln p(xd|θ)+R(xd,hm)]+ KL(q, π)+ ln ξ−1+Ψ′′π,ν(c,D) 3 cD , (15)
where Ψ′′π,ν(c,D) is the modified constant of Ψπ,ν(c,D) andR(x,hm) isR(x,h) in Eq.(12) replacing h(x, θ)−2 of Eq.(13) with hm(x, θ)−2 =exp(ln p(x|θ)+minθln p(x|θ)−2 maxθln p(x|θ)).
Proof sketch. We express Eq(θ)[ln p(x|θ) + R(x,hm)] as Eq(θ)q(θ′)q(θ′′)L(x,θ,θ′,θ′′) where L(x,θ,θ′,θ′′) :=ln p(x|θ)+(2hm(x,θ))−2(ln p(x|θ)2−2 ln p(x|θ) ln p(x|θ′)+ln p(x|θ′) ln p(x|θ′′)). Then, we apply the same proof technique as Theorem 1 [7] to the loss function L(x,θ,θ′,θ′′) with λ = 3cD. Full proof is given in Appendix D.4. Remark 3. To reduce the upper bound of the generalization error, Eq.(15), we need to control the trade-off between the data fitting term of the negative log-likelihood and enhancing the diversity of the models based on the loss repulsion term R. Remark 4. In the definition of hm, minθln p(x|θ) is too pessimistic in some cases. If we additionally assume that there exists a positive constant M s.t. Eq(θ)[ln p(x|θ)]2 < M < ∞, we can replace minθln p(x|θ) with Medianθ(ln p(x|θ))−M1/2 in hm in Theorem 4 (see Appendix D.4.1).
Compared to Eq.(9), our bound focuses on the variance of loss functions, which has a direct connection to the repulsion of PVI and DPPs (see Section 3.2). Furthermore, we show that optimizing the upper bound in Eq.(15) shows competitive performance with the existing state-of-the-art PVI (see Section 5). Note that this inequality is not restricted to the case where q(θ) is an empirical distribution. We can also use this for parametric variational inference [21].

heading: 3.2 Diversity in ensemble learning and connection to existing methods
text: In the following, we focus on the ensemble setting and use a finite set of parameters as ρE(θ) := 1 N ∑N i=1 δθi(θ) and discuss the relationship of our theory and existing methods. We show the summary of the relationships in Appendix J.

heading: 3.2.1 Covariance form of the loss repulsion
text: To emphasize the repulsion between models, we upper-bound Eq.(12) using the covariance:
Theorem 5. Under the same assumption as Corollary 1, we have
EρE(θ) ln p(x|θ) ≤ lnEρE(θ)p(x|θ)− 1
2(2hw(x, θ))2N2 N∑ i,j=1 (ln p(x|θi)− ln p(x|θj))2 , (16)
where hw(x, θ)−2 = exp ( mini ln p(x|θi) + 1N ∑ i=1 ln p(x|θi)− 2 maxj ln p(x|θj) ) .
See Appendix D.5 for the proof. We can also show a generalization error bound like Theorem 4 for the ensemble learning setting (see Appendix E.1). In existing PVI and DPPs, the repulsion is based not on the loss function but the parameters or models, as seen in Section 2. We derive the direct connection between our loss repulsion and the model and parameter repulsion below.

heading: 3.2.2 Relation to w-SGLD and model repulsion
text: First, from Eq.(16), we derive the direct connection to w-SGLD, which is an instance of PVI introduced in Section 2. Let us define an N × N kernel Gram matrix G whose (i, j) element is defined as
Gij := exp ( −(8h2w)−1 (ln p(x|θi)− ln p(x|θj))2 ) . (17)
Applying the Jensen inequality to Eq.(16), we obtain
lnEρE(θ)p(x|θ) ≥ EρE(θ) ln p(x|θ)− 1
N N∑ i=1 ln N∑ j=1 Gij N ≥ EρE(θ) ln p(x|θ). (18)
This is tighter than the standard Jensen inequality. To derive the relation to w-SGLD, we optimize the middle part of Eq.(18) by gradient descent. We express L(θi) := ln p(x|θi) and do not consider the dependency of hw on θ for simplicity. By taking the partial derivative with respect to θi, we have
∂θi ln p(x|θi) +  N∑ j=1 ∂L(θj)Gij∑N k=1Gjk + ∑N j=1 ∂L(θj)Gji∑N k=1Gik  ∂θiL(θi). (19) The second term corresponds to the repulsion term, which is equivalent to that of w-SGLD shown in Table 1. The difference is that our Gram matrix G in Eq.(19) depends on the loss function rather than the parameter or model. Using the mean value theorem, it is easy to verify that there exists a constant C such that ‖ ln p(x|θi)− ln p(x|θj)‖2 = ‖C(θi − θj)‖2 (see Appendix G for details), thus we can easily transform the loss repulsion to the parameter or model repulsion.
However, since we cannot obtain the explicit expression of the constant C, it is difficult to understand the intuitive relation between our loss repulsion and the parameter or model repulsion. Instead, here we directly calculate ∂θiGi,j and discuss the relation. Due to the space limitation, we only show the relation to the model repulsion in the regression task of f-PVI. See Appendix G for the complete statement including the classification setting of f-PVI and the parameter repulsion of PVI. Following the setting in Section 2.1, for a regression problem, we assume that p(y|f(x; θ)) is the
Gaussian distribution with unit variance for simplicity. We define L(fi) := ln p(y|f(x; θi)) and dLij := ∂fiL(fi) + ∂fjL(fj). The derivative of the Gram matrix G is expressed as
∂θiGij = −((fi − fj)‖dLij‖2︸ ︷︷ ︸ i) + ∂fiL(fi)dLij‖fi − fj‖2︸ ︷︷ ︸ ii) )(4hw) −2Gij∂θifi. (20)
The first term i) corresponds to the model repulsion of f-PVI shown in Eq.(30) and the second term ii) is the correction term based on the loss function. Thus, our loss repulsion can be translated to the model repulsion of f-PVI plus the correction term using the loss function.
In conclusion, we have confirmed that w-SGLD is directly related to our theory. For Eq.(18), we can also derive a generalization error bound like Theorem 4 (see Appendix E.2 for details). This means that w-SGLD controls the trade-off between the model fitting and enhancing the diversity by optimizing the generalization error bound. This explains the reason why w-SGLD still works well even with a finite ensemble size.

heading: 3.2.3 Relation to other PVIs
text: Here, we derive other PVI and DPPs from our theory. First, we derive GFSF shown in Table 1. We express the identity matrix with size N as I . We obtain the following upper bound from Eq.(16): Theorem 6. Under the same assumption as Corollary 1, for any real value ∈ (1, N − 1) and for any positive constant h̃, we have
lnEρE(θ)p(x|θ) ≥ EρE(θ) ln p(x|θ)− 2
h̃N ln det( I +K) +
2 lnN
h̃N ≥ EρE(θ) ln p(x|θ), (21)
where K is an N ×N kernel Gram matrix of which (i, j) element is defined as Kij := exp ( −h̃ lnN(4hw)−2 (ln p(x|θi)− ln p(x|θj))2 ) . (22)
The proof is shown in Appendix D.6. This is tighter than the standard Jensen inequality.
Remark 5. Eq.(21) holds for any positive bandwidth h̃. h̃ defines the tightness Eq.(21). Remark 6. In Eq.(22), we used the scaling of lnN to define K. This is motivated by the median trick of the existing PVI, which tunes the bandwidth as lnN/median2. This scaling implies that, for each i, ∑ jKij≈1+ 1N holds. We found that using this scaling is necessary to obtain the bound Eq.(21). We conjecture that this is the reason why scaling the bandwidth is important for PVI in practice.
In the same way as Eq.(19) for w-SGLD, we also optimize the middle term in Eq.(21) by GD. By taking the partial derivative, we have the following update equation:
∂θiEρE(θ) log p(x|θi) + 2
h̃N ∑ j (K + I)−1ij ∇L(θj)Kij∂θiL(θi). (23)
See appendix D.7 for the proof. The second term is the repulsion force, which is the same as that in the update equation of GFSF in Table 1.
Next we consider the relation to DPPs. Using the trace inequality [8] to Eq.(16) and using a Gram matrix G̃ whose (i, j) element is G
1 2 ij in Eq.(17), we obtain
lnEρEp(x|θ) ≥ EρE ln p(x|θ) + 2
N ln detG̃−lnN (24)
The proof is shown in Appendix D.8. The lower bound term of Eq.(24) is equivalent to the objective function of the DPP introduced in Eq.(5).
For Eqs.(21) and (24), we can derive a PAC-Bayesian generalization error bound like Theorem 4 (see Appendix E.2 for details). Moreover, we can connect our loss repulsion in Eqs.(21) and (24) to model and parameter repulsions in the same way as w-SGLD. Accordingly, GFSF and DPP are closely related to the second-order generalization error bound.
In conclusion, we derived PVI and DPPs from our theory based on the second-order Jensen inequality. On the other hand, we found it difficult to show its relation to SVGD, since it has the kernel smoothing coefficient in the derivative of the log-likelihood function. We leave it for future work.

heading: 4 Discussion and Related work
text: In this section, we discuss the relationship between our theory and existing work.

heading: 4.1 Theoretical analysis of PVI
text: Originally, PVI was derived as an approximation of the gradient flow in Wasserstein space and its theoretical analysis has been done with an infinite ensemble size [16, 12]. In practice, however, an infinite ensemble size is not realistic and various numerical experiments showed that PVI still works well even with a finite ensemble size [17, 28]. Accordingly, our analysis aimed to clarify why PVI works well with such a finite-size ensembles. On the other hand, as discussed in Section 3.2.2, there is a difference between our loss repulsion and the parameter and model repulsion used in the existing works, and thus it would be an interesting direction to extend our theory to fill such a gap.

heading: 4.2 Relation to existing second-order Jensen inequality
text: Recently, some works derived tighter Jensen inequalities [14, 6]. Liao and Berg [14] derived a second-order Jensen inequality, and Gao et al. [6] worked on a higher-order inequality. Masegosa [19] combined the earlier result [14] with the PAC-Bayesian theory. Although our bound is also a second-order Jensen inequality, its derivation and behavior are completely different from them. In Liao [14], their second-order Jensen inequality includes the term of a variance of a random variable, and Masagosa [19] considered p(x|θ) to be a corresponding random variable that depends on θ. Thus, the second-order inequalities depend on the variance of the predictive distribution Ep(θ)[p(x|θ)]. On the other hand, our bound is based on our novel second-order equality shown in Theorem 2, which leads to the variance of a loss function as shown in Theorem 3. By using the variance of loss functions, we can directly connect our theory to existing PVI and DPPs as shown in Section 3.2. Moreover, as shown in Section 5, including the predictive variance in the objective function results in a large epistemic uncertainty, which means that individual models do not fit well. On the other hand, ours does not show this phenomenon. Consequently, our result can be regarded as an extension of the earlier work [14, 6, 19] that directly focuses on a loss function in machine learning.
Masagosa [19] showed that the second-order PAC-Bayesian generalization error is especially useful under misspecified models, i.e., for any θ, p(x|θ) 6= ν(x). Our theories can also be extended to such a setting (see Appendix H for further discussion).
Other closely related work is an analysis of the weighted majority vote in multiclass classification [20], which uses a second-order PAC-Bayesian bound. While their analysis is specific to the majority vote of multiclass classification, our analysis has been carried out in a more general setting based on the second-order Jensen inequality derived from Theorem 2 and Theorem 3.

heading: 5 Numerical experiments
text: According to our Theorem 4, it is important to control the trade-off between the model fitting and diversity enhancement in order to reduce the generalization error. Therefore, we minimize our generalization error bound Eq.(15) directly and confirm that the trade-off is controlled appropriately. Our objective function is
F({θi}Ni=1) := − 1
N N∑ i=1 D∑ d=1 [ln p(xd|θi)+R(xd,h)] + KL(ρE, π), (25)
which is obtained setting q(θ) = ρE(θ) and c= 1 in Eq.(15). We call our approach as Variance regularization, and express it as VAR. We compared the performances of our methods with MAP, PAC2E, and f-PVI(f-SVGD) on toy data and real data. Our experimental settings are almost same as that of the previous work [28] and the detailed settings and additional results are shown in Appendix I.
As for the KL divergence, since ρE is an empirical measure, we need to carefully define a prior distribution so that the KL divergence between ρE and π can be defined properly. We used a certain prior distribution as introduced previously [19] (see Appendix C.3 for details). Moreover, to eliminate such a limitation, we considered smoothing the gradient using SVGD. We call this approach
VAR-SVGD, of which update equation is given as
θnewj ←− θoldj + η 1
N N∑ i=1 Kij∂θoldi ( log p(D|θoldi )π(θoldi ) +R(D, h) ) + ∂θoldi Kij , (26)
where R(D, h) = ∑Dd=1R(xd, h) and Kij is the Gram matrix using the Gaussian kernel defined similarly to the model repulsion in Section 2.1. We tuned its bandwidth by the median trick.

heading: 5.1 Toy data experiments
text: First, using toy data, we visualized the model fitting and diversity in various methods. We considered a regression task and we randomly generated the 1-dimensional data points and fit them by using a feed-forward neural network model with 2 hidden layers and ReLU activation, which has 50 units. We used 50 ensembles for each method. The results are shown in Figure 1, which visualizes 95% credible intervals for the prediction and mean estimate corresponding to aleatoric and epistemic uncertainties.
Hamilton Monte Carlo (HMC)[21] is the baseline method used to express the Bayesian predictive uncertainty properly. MAP and SVGD methods give an evaluation of the uncertainty that is too small. Our method and f-SVGD showed very similar evaluations for uncertainty. PAC2E [19] showed large epistemic uncertainty, which is expressed as dark shaded area in Figure 1. We conjecture it is because PAC2E includes the predictive variance in the objective function and the enhanced diversity is too large. Described below, PAC2E shows slightly worse results in real data experiments than those of other approaches. This might be because including the predictive variance in the objective function does not result in a better trade-off of between model fitting and enhancing diversity in practice.

heading: 5.2 Regression task on UCI
text: We did regression tasks on the UCI dataset [4]. The model is a single-layer network with ReLU activation and 50 hidden units except for Protein data, which has 100 units. We used 20 ensembles. Results of 20 repetition are shown in Table 2. We found that our method compares favorably with f-SVGD. We also found that PAC2E shows worse performance than those of other methods. We conjectured that this is because the predictive variance in the objective function enhances too large diversity as shown in Figure 1, which hampers the fitting of each model to the data.

heading: 5.3 Classification task on MNIST and CIFAR 10
text: We conducted numerical experiments on MNIST and CIFAR 10 datasets. For MNIST, we used a feedforward network having two hidden layers with 400 units and a ReLU activation function and used 10 ensembles. For CIFAR 10, we used ResNet-32 [9], and we used 6 ensembles. The results are shown
in Table 3. For both datasets, our methods show competitive performance compared to f-SVGD. For CIFAR 10, as reported previously [28], f-SVGD is worse than the simple ensemble approach. We also evaluated the diversity enhancing property by using the out of distribution performance test. It is hypothesized that Bayesian models are more robust against adversarial examples due to their ability to capture uncertainty. Thus, we generated attack samples and measured the vulnerability to those samples. We found that, as shown in Figure 4 in Appendix I, our method and f-SVGD showed more robustness compared to MAP estimation in each experiment.

heading: 5.4 Contextual bandit by neural networks on the real data set
text: Finally, we evaluated the uncertainty of the obtained models using contextual bandit problems [24]. This problem requires the algorithm to balance the trade-off between the exploitation and exploration, and poorly evaluated uncertainty results in larger cumulative regret. We consider the Thompson sampling algorithm with Bayesian neural networks having 2 hidden layers and 100 ReLU units, and we used 20 particles for each experiment. Results of 10 repetition are shown in Table 4. We can see that our approach outperform other methods.

heading: 6 Conclusion
text: In this work, we derived a novel second-order Jensen inequality that includes the variance of loss functions. We also derived a PAC-Bayesian generalization error bound. Our error bound shows that both model fitting and enhancing diversity are important for reducing the generalization error. Then, we derived the existing PVI and DPPs from our new Jensen inequality. Finally, we proposed a new PVI that directly minimizes our PAC-Bayesian bound. It shows competitive performance with the current state-of-the-art PVI. In future work, it would be interesting to apply our second-order Jensen inequality to general variational inference or optimal control problems.
Other interesting direction is to derive an upper-bound of the Jensen gap. Gao et al. [6] derived lnEp(x|θ) − E ln p(x|θ) ≤ Variance(p(x|θ)), which uses the predictive distribution. This means that the larger the predictive variance is, the larger upper bound of the Jensen gap we have. We leave it for future work to upper-bound the Jensen gap with the variance of loss function using our theory.
Acknowledgments and Disclosure of Funding
FF was supported by JST ACT-X Grant Number JPMJAX190R.

Review summaries:
This paper derives a tighter Jensen Inequality and applied it into the framework of particle variational inference. Different from previous work, it tackles this problem in light of PAC-Bayesian analysis, and provides a new second-order Jensen inequality, which has the repulsion term based on the loss function. In addition, the authors derive the existing PVI and DPPs and demonstrate that these methods work well even at a finite ensemble size, since their objective functions includes valid diversity enhancing terms to reduce the generalization error. Finally, the numeral experiments show the better performance than the most related work PAC2.

This paper studies particle variational inference (PVI), which is a mix between ensemble and VI (the variational posterior is assumed to be a mixture of Dirac distributions, and an additional "repulsion term" is added to give better diversity), and proposes a novel generalization to Jensen's inequality. From here, the authors propose an upper bound to the first term of the ELBO (the expected log-likelihood term), which is tighter compared to the standard Jensen inequality. Empirical findings show that the proposed loss performs well in various cases.

The paper presents a tighter PAC-Bayes bound based on the second-order Jensen inequality. Remarkably, minimizing the new bound under various relaxations/approximations recovers various PVIs and imposing a DPP prior. This novel view of PVIs explains why their repulsive force helps generalization (also, this is a valuable finite-sample regime analysis of PVIs), and also inspires new algorithms for PVIs. Experiment verifies the utility and the benefits of the new methods.