Variational Inference via Resolution of Singularities

1 INTRODUCTION . Singular statistical models are ubiquitous in modern machine learning . In contrast to their regular counterpart , singular models need not be identifiable nor possess a positive-definite Fisher information matrix . These departures from classic regularity conditions have considerable implications for both theory and practice . Since neural networks are ( strictly ) singular models , accounting for such differences may be critical to some of the most pressing challenges in deep learning theory including model selection and the generalization puzzle ( Murfet et al. , 2020 ) . In this work , singular learning theory ( Watanabe , 2009 ) is brought to bear on the challenge of performing inference in Bayesian neural networks ( MacKay , 1992 ; Neal , 1996 ) . In particular , singular learning theory helps us understand the large-sample properties of the posterior distribution over neural network weights w ∈ Rd . By ( Watanabe , 2018 , Chapter 6 ) , there exists a so-called desingularization map ( also known as a resolution map ) , g : Rd → Rd , g ( ξ ) = w , such that when n is large , the posterior distribution in a local weight set is proportional to exp ( −nξ2k11 ξ 2k2 2 · · · ξ 2kd d ) |ξ h1 1 · · · ξ hd d |b ( ξ ) . ( 1 ) The k ’ s and h ’ s are natural numbers , and b ( · ) is a real-valued C∞ function . We say the posterior distribution has been put into standard form if coordinates ξ have been found that allow the posterior to be locally written as in equation 1 . Remarkably , it is possible under very general conditions to write the posterior distribution over the parameters of a singular model as just such a mixture of standard forms . Based on equation 1 , we proceed to demonstrate that a certain mean-field variational family , following desingularization , can recover the leading order term of the log ( normalized ) evidence , up to constant terms that do not depend on sample size n or dimension d. Let ( x , y ) denote the input-target pair modeled jointly as p ( x , y|w ) = p ( y|x , w ) p ( x ) . Let us assume the parameter space W is a compact set in Rd and p0 ( x , y ) = p0 ( y|x ) p ( x ) is the true data-generating mechanism . Throughout we suppose there exists w0 ∈ W such that p0 ( y|x ) = p ( y|x , w0 ) . In the parlance of singular learning theory , this condition is known as realizability . Let ϕ ( w ) be a compactly-supported prior . We shall refer to ( p ( · , · ) , p0 ( · , · ) , ϕ ( · ) ) as a model-truth-prior triplet . Define K ( w ) to be the Kullback-Leibler divergence between the truth and the model , as a function of the model parameters w : K ( w ) = KL ( p0 ( x , y ) ||p ( x , y|w ) ) = Ep0 log p0 ( y|x ) p ( y|x , w ) . Following Watanabe ( 2009 ) , we say a model is regular if it is 1 ) identifiable i.e. , { w : K ( w ) = 0 } is a singleton , and 2 ) its Fisher information matrix I ( w ) is positive definite for arbitrary w ∈W . We call a model strictly singular if it is not regular . The term singular will refer to either regular or strictly singular models . Henceforth let p ( y|x , w ) be a neural network model . We wish to approximate the intractable posterior distribution over neural network weights , p ( w|Dn ) = ∏n i=1 p ( yi|xi , w ) ϕ ( w ) Z ( n ) , whereDn = { ( xi , yi ) } n i=1 is a dataset of n input-output pairs . The normalizing constant , Z ( n ) = ∫ ∏n i=1 p ( yi|xi , w ) ϕ ( w ) dw , is variously known as the ( model ) evidence and the marginal likelihood . The asymptotic expansion of Z̄K ( n ) , a variant of Z ( n ) defined in equation 5 , will prove crucial to justifying the proposed variational approximation . For strictly singular models , quantities such as Z ( n ) and Z̄K ( n ) manifest as a singular integral , i.e. , an integral of the form ∫ W e−nf ( w ) ϕ ( w ) dw where W ⊂ Rd is a compact semi-analytic subset , and f and ϕ are real analytic functions . The behavior of a singular integral critically depends on the zeros of f . 2 VARIATIONAL INFERENCE FOR SINGULAR MODELS . Variational inference is an approximate inference technique in which a family of densities Q , often called the variational family , is first posited and a member of the variational family , some q∗ ∈ Q , is then found via optimization according to some criterion that measures closeness to the desired target density . To begin , let us write the posterior distribution p ( w|Dn ) in two alternate forms : Ln ( w ) = − 1 n n∑ i=1 log p ( yi|xi , w ) p ( w|Dn ) ∝ e−nLn ( w ) ϕ ( w ) Z ( n ) = ∫ W e−nLn ( w ) ϕ ( w ) dw Kn ( w ) = 1 n n∑ i=1 log p0 ( yi|xi ) p ( yi|xi , w ) p ( w|Dn ) ∝ e−nKn ( w ) ϕ ( w ) Z̄ ( n ) = ∫ W e−nKn ( w ) ϕ ( w ) . On the left , the posterior is written in terms of the average negative log likelihood Ln ( w ) and on the right in terms of the average log likelihood ratio Kn ( w ) , which is the empirical counterpart to K ( w ) . Since Z̄ ( n ) = Z ( n ) / ∏n i=1 p0 ( yi|xi ) , we refer to Z̄ ( n ) as the normalized evidence . Now , consider some general change of variables g : Rd → Rd , g ( ξ ) = w. In the new coordinate ξ , the KL divergence between a variational distribution q ( ξ ) and the desired posterior target , p ( ξ|Dn ) ∝ e−nKn ( g ( ξ ) ) ϕ ( g ( ξ ) ) |g′ ( ξ ) | , is given by KL ( q ( ξ ) ||p ( ξ|Dn ) ) = EqnKn ( g ( ξ ) ) + KL ( q ( ξ ) ||ϕ ( g ( ξ ) ) |g′ ( ξ ) | ) + log Z̄ ( n ) . As long as the support of q is contained in the support of the posterior , we have KL ( q ( ξ ) ||p ( ξ|Dn ) ) ≥ 0 and hence the following bound : Ψ ( q , g ) : = −EqnKn ( g ( ξ ) ) −KL ( q ( ξ ) ||ϕ ( g ( ξ ) ) |g′ ( ξ ) | ) ≤ log Z̄ ( n ) . ( 2 ) Equality in equation 2 is achieved if and only if q ( ξ ) = p ( ξ|Dn ) . It is easy to recognize that maximizing Ψ ( q , g ) is equivalent to maximizing the so-called evidence lower bound , ELBO ( q , g ) : =Eq [ n∑ i=1 log p ( yi|xi , g ( ξ ) ) − log q ( ξ ) + log ( ϕ ( g ( ξ ) ) |g′ ( ξ ) | ) ] , ( 3 ) since Ψ ( q , g ) = ELBO ( q , g ) + nSn where Sn = − 1n ∑n i=1 log p0 ( yi|xi ) is the empirical entropy . Indeed , just as Ψ ( q , g ) is a lower bound on the log ( normalized ) evidence log Z̄ ( n ) , so too is ELBO ( q , g ) a lower bound on the log ( unnormalized ) evidence logZ ( n ) . To facilitate theoretical analysis , we will work with the deterministic counterparts to Ψ ( q , g ) and Z̄ ( n ) , respectively given by ΨK ( q , g ) : = −EqnK ( g ( ξ ) ) −KL ( q ( ξ ) ||ϕ ( g ( ξ ) ) |g′ ( ξ ) | ) , ( 4 ) and Z̄K ( n ) : = ∫ W e−nK ( w ) ϕ ( w ) dw . ( 5 ) It is plain to see that just as Ψ ( q , g ) ≤ log Z̄ ( n ) , we have ΨK ( q , g ) ≤ log Z̄K ( n ) . Our theoretical results address the possibility that supq∈QΨK ( q , g ) ≈ log Z̄K ( n ) for some variational family Q and change-of-variables g. By ( Watanabe , 2009 , Theorem 6.7 ) , for the leading order term in the asymptotic expansion of Z̄K ( n ) , we have Z̄K ( n ) ≈ Cn−λ ( log n ) m−1 . ( 6 ) The rational number λ ∈ [ 0 , d/2 ] is an important quantity in singular learning theory known as the real log canonical threshold ( RLCT ) and the integer m ≥ 1 is its associated multiplicity . These two quantities , to be defined in equation 11 , are uniquely associated to a model-truth-prior triplet . Equipped with equation 6 , ( Bhattacharya et al. , 2020 , Theorem 3.1 ) established that when g is a resolution map , the mean-field variational family Q ( 0,1 ] of equation 12 can achieve sup q∈Q ( 0,1 ] ΨK ( q , g ) ≥ −λ log n− constant independent of n. ( 7 ) ( In Appendix C , we prove this result under more general conditions than assumed in Theorem 3.1 of Bhattacharya et al . ( 2020 ) . ) Unfortunately , the bound in equation 7 provides little reassurance that Q ( 0,1 ] is a desirable variational family . Contribution Following Lin ( 2011 ) , let us call C in equation 6 the leading coefficient of Z̄K ( n ) . We go beyond the analysis in Bhattacharya et al . ( 2020 ) by taking into account those terms in the leading coefficient that depend on the dimension d , call it C ( d ) . Our main result , Theorem 5.1 , shows that if g is a resolution map , then the same variational family Q ( 0,1 ] of equation 12 can achieve sup q∈Q ( 0,1 ] ΨK ( q , g ) = −λ log n+ logC ( d ) . ( 8 ) Next , rather than presuming the resolution map theoretically tractable as in Bhattacharya et al . ( 2020 ) , we employ a normalizing flow to learn the unknown resolution map g at the same time as learning the variational parameters in q . We are aided by the fact that a resolution map can attain the optimal value of ΨK ( q , g ) and therefore justifies learning the resolution map via optimization of Ψ ( q , g ) . Finally let us note that although the result in equation 8 is stronger than that in equation 7 , it does come at the cost of additional assumptions as we will discuss at the end of Section 5 . Remark . When the model is regular , we need not bother with singular learning theory and may obtain Z̄K ( n ) ≈ ϕ ( w0 ) √ ( 2π ) d detH ( w0 ) n−d/2 via the Laplace approximation . The Laplace approximation , however , is egregiously inappropriate for singular models , in particular neural network models . Since λ = d/2 and m = 1 in regular models , equation 6 is a true generalization of the Laplace approximation , holding for both regular and strictly singular models .

Summary: Working with Bayesian neural networks, this paper proposed a variational algorithm to approximate posterior distribution of the network weights. To overcome model singularities, the authors used the idea of normalizing flow by transforming the weights through an affine coupling network, and subsequently worked on the desingularized parameter space. In addition, they derived an asymptotic expression for the ELBO, and compared the Gaussian and generalized gamma approximating families in the experiments.