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: The paper builds on the recent theoretic results of [Bhattacharya et al. (2020)], which shows that a mean-field variational approximation with carefully chosen approximation family leads to an ELBO $\Psi$ which is sharp up to a constant C(d) which only depends on the dimensionality d of the parameter space. For the proof, [Bhattacharya et al. (2020)] assumes that the approximation family is a Gamma distribution truncated to [0,1].  Therefore the authors of this paper conjecture that the Gamma distribution truncated to [0,1] is a good choice for the source distribution of a normalizing flow. Their experimental results suggest that this choice leads a higher ELBO than using as source distribution a normal distribution.