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

1 Introduction . 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 ] can not 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 . 2 Background . In this section , we briefly review PVI , DPPs , and PAC-Bayesian analysis . 2.1 Particle variational inference . 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 . 2.2 Regularization based methods and determinantal point processes . 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 . 2.3 PAC-Bayesian theory . 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 can not 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 ) .

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