Year: 2021

Abstract:
Predicated on the premise that neural networks are best viewed as singular statistical models, we set out to propose a new variational approximation for Bayesian neural networks. The approximation relies on a central result from singular learning theory according to which the posterior distribution over the parameters of a singular model, following an algebraic-geometrical transformation known as a desingularization map, is asymptotically a mixture of standard forms. From here we proceed to demonstrate that a generalized gamma mean-field variational family, following desingularization, can recover the leading order term of the model evidence. Affine coupling layers are employed to learn the unknown desingularization map, effectively rendering the proposed methodology a normalizing flow with the generalized gamma as the source distribution.

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

heading: 2 VARIATIONAL INFERENCE FOR SINGULAR MODELS
text: 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.

heading: 3 SINGULAR LEARNING THEORY
text: That the posterior distribution in singular models can be written, under quite general conditions, as a mixture of standard forms is predicated on the monomialization of K(w). The following theorem from Watanabe (2009), adapted for notational consistency, gives precise conditions for the existence of the resolution map, an algebraic geometrical transformation such that K(w) can be written as a monomial. The result is itself based on Hironaka’s resolution of singularities, a celebrated result in modern algebraic geometry. To prepare, let W = {w ∈ W : K(w) ≤ } for some small positive constant and W (R) be some real open set such that W ⊂W (R) . The theorem below will make use of the multi-index notation: for a given w = (w1, . . . , wd) ∈ Rd, define wk := wk11 · · ·w kd d where the multi-index k = (k1, . . . , kd) with each kj a nonnegative integer. Theorem 3.1 (Theorem 6.5 of Watanabe (2009)). Suppose the model-truth-prior triplet (p, p0, ϕ) satisfies Fundamental Conditions I and II with s = 2 in Watanabe (2009). We can find a real analytic manifold M (R) and a proper and real analytic map g : M (R) →W (R) such that
1. M = g−1(W ) is covered by a finite set M = ∪αMα where Mα = [0, b]d.
2. In each Mα, K(g(ξ)) = ξ2k = ξ2k11 · · · ξ 2kd d , (9)
where kj ∈ N such that not all kj are zero.
3. There exists C∞ function b(ξ) such that
ϕ(g(ξ))|g′(ξ)| = ξhb(ξ) = ξh11 · · · ξ hd d b(ξ), (10)
where hj ∈ N, |g′(ξ)| is the absolute value of the determinant of the Jacobian and b(ξ) > c > 0 for ξ ∈ [0, b]d.
If we “plug in" equation 9 and equation 10 into the transformed posterior p(ξ|Dn), we obtain the first display of the paper, equation 1. Theorem 6.5 of Watanabe (2009) holds for regular statistical models as well, e.g., by the transform w = g(ξ) = w0 + I(w0)1/2ξ, we can put a regular model-truth-prior triplet into the standard form. It is worth noting that neither the resolution map g nor the multi-indices k and h are unique for a given triplet (p, p0, ϕ).
A crucial quantity that appears in singular learning theory is the real log canonical threshold (RLCT). Let {Mα : α} be as in Theorem 3.1 and λj = hj+12kj , j = 1, . . . , d where hj and kj are the entries of the multi-indices h and k in a local coordinate Mα. (For brevity, the dependence on α has been suppressed.) When kj = 0, λj is taken to be infinity. Uniquely associated to a triplet (p, p0, ϕ) are its real log canonical threshold (RLCT) and its multiplicity defined, respectively, as
λ = min α min j∈1,...,d λj , m = max α
#{j : λj = λ}. (11)
Let {α∗} be the set of those local coordinates in which both the min and max in equation 11 are attained. Watanabe (2000) calls such a set the essential coordinates.
Assuming the prior is proper, the RLCT of a model-truth-prior triplet is at most d/2 (Watanabe, 2009, Theorem 7.2). When the model is regular, the RLCT is exactly equal to d/2 and the multiplicity m = 1 (Watanabe, 2009, Remark 1.15). Murfet et al. (2020) argues that (twice) the RLCT might be a most natural count of parameters in singular models. In Appendix A, we recall a simple toy example where the resolution map, the RLCT, and the multiplicity can be calculated explicitly. Remark. In Section 6, we will learn the resolution map using affine coupling layers. However as indicated by Theorem 3.1, the resolution map g(ξ) as well as the multi-indices k and h in fact depend on α. Despite this, it is unclear if learning multiple resolution maps gα would have any practical advantages since the RLCT is determined entirely by the essential local coordinates.

heading: 4 RELATED WORK
text: Bayesian learning in neural networks well precede the advent of modern deep learning MacKay (1992); Neal (1996). The resurgence of interest in Bayesian learning for deep neural networks, sometimes called Bayesian deep learning Wilson & Izmailov (2020), has been prompted by concerns of overconfidence and miscalibration. Since exact inference for Bayesian neural networks is intractable, all methods proceed by approximate inference. A major class of approximate inference techniques is based on scaling classic MCMC to modern settings of large datasets and deep neural networks. Some scalable variants of MCMC suitable for deep neural networks include Welling & Teh (2011); Chen et al. (2014); Zhang et al. (2019). Another major approximate inference technique for Bayesian neural network is represented by variational inference which learns the target posterior via optimization. The various flavors of variational inference can be commonly characterized by two ingredients: 1) an approximating family, e.g., a class of distributions over the neural network weights and 2) a criterion for measuring closeness of two distributions. The most commonly employed approximating family is undoubtedly the mean-field family of fully factorized Gaussian distributions. Many variational inference techniques share this approximating family even if they use different criterion to measure closeness to the target (Graves, 2011; Blundell et al., 2015; Hernandez-Lobato et al., 2016; Li & Turner, 2016; Khan et al., 2018; Sun et al., 2019). The limitations of the mean-field Gaussian approximating family are well known however (MacKay, 1992). The desire to move beyond
mean-field Gaussian has motivated many recent methods to make use of more realistic covariance structures (Zhang et al., 2018) or more expressive approximating families, e.g., via normalizing flows (Louizos & Welling, 2017). Monte Carlo dropout (Gal & Ghahramani, 2016) is another popular approximate inference technique for Bayesian neural networks which, despite first appearances, can in fact be viewed as variational inference. Another strain of work is based on the idea of using stochastic gradient descent as a sampler for the underlying posterior distribution of interest. Works in this spirit include Mandt et al. (2018), Izmailov et al. (2018) and Maddox et al. (2019). Finally there are various approximate inference techniques for Bayesian neural networks that are not easily classifiable according to the distinctions above, e.g., temperature scaling (Guo et al., 2017) and deep ensembles (Lakshminarayanan et al., 2017).

heading: 5 THE GENERALIZED GAMMA MEAN-FIELD APPROXIMATION
text: Consider the mean-field variational family proposed in Bhattacharya et al. (2020) Q(0,1] = {qλ,k,β = d∏ j=1 qλj ,kj ,βj (ξj) : λ = (λ1, . . . , λd) ∈ Rd>0,k = (k1, . . . , kd) ∈ Rd>0,β = (β1, . . . , βd) ∈ (0,∞)d}
qj(ξj) := qλj ,kj ,βj (ξj) ∝ ξ 2kjλj−1 j exp(−βjξ 2kj j )1(0,1](ξj), (12)
where each univariate density qj supported on (0, 1] is the density of a truncated generalized gamma random variable. Let j∗ ∈ {1, . . . , d} be the dimension that attains λ̃j∗ = λ̃ where λ̃ is the RLCT of the underlying (unknown) model-truth-prior triplet. It was established in Bhattacharya et al. (2020) that if g is a resolution map, then Q(0,1] satisfies equation 7 in particular by setting λ and k to their respective true values, and βj∗ = n and all other βj = 1.
In Theorem 5.1 below we will instead establish equation 8. In the proof we show this can be accomplished by setting λ and k to their true values as in Bhattacharya et al. (2020) though the optimal values of β will be different. Finally, although we assume below that m = 1, this is not necessary; as long as m d, we can set βj = n1/m for all j such that λ̃j = λ̃. Theorem 5.1. Suppose the model-truth-prior triplet is such that Theorem 3.1 holds with [0, b] = [0, 1], K(g(ξ)) = ξ2k̃ and ϕ(g(ξ))|g′(ξ)| = ξh̃. Let λ̃ denote the RLCT of the triplet and assume the multiplicity m = 1. Then supq∈Q(0,1] ΨK(q, g) = −λ̃ log n + logC(d) where C(d) is the term in the leading coefficient C that depends on d.
Proof. If ξ ∈ [0, 1]d and ϕ(g(ξ))|g′(ξ)| ∝ ξh̃, then necessarily b(ξ) = ∏d j=1 1/(h̃j + 1). Let
λ̃j = (h̃j + 1)/(2k̃j). Applying Corollary 5.9 in Lin (2011), we obtain
Z̄K(n) = ∫ [0,1]d exp(−nξ2k̃)ξh̃ dξ ≈ Cn−λ̃(log n)m−1 = Cn−λ̃ (13)
where
C = Γ(λ̃) (m− 1)! ∏d j=1(2k̃j) ∏d j=m+1(2k̃j)(λ̃j − λ̃) . (14)
If we denote by C(d) the terms in the leading coefficient that depend on d, then we have
logC(d) = − d∑ j=1 log(2k̃j)− d∑ j=m+1 log(2k̃j)− d∑ j=m+1 log λ̃j − d∑ j=m+1 log(1− λ̃/λ̃j).
Returning to our variational distribution, if the kj’s are well specified and we additionally set βj∗ = n where j∗ is the dimension that attains λ̃j∗ = λ̃, we then have EqnK(g(ξ)) = λj∗ ∏ j 6=j∗ G(λj , βj),
where G(λ, β) is as in B.1. Next, we have
KL(q(ξ)||ϕ(g(ξ))|g′(ξ)|) = KL ( q ∥∥∥ ξh̃b(ξ)) = d∑
j=1
KL ( qj ∥∥∥ ξh̃jj /(h̃j + 1))
= d∑ j=1 Eqj log qj − h̃jEqj log ξj + log(h̃j + 1).
If in addition the λj’s are all well-specified, then
KL(q||ϕ(g(ξ))|g′(ξ)|) = d∑ j=1 [ −βjG(λ̃j , βj)− logB(k̃j , h̃j , βj) + log(2k̃j) + log λ̃j ] ,
where B(k, h, β) is as in B.1. Now let us make use of the fact that logB(k, h, β) −λ log β and G(λ, β) λ/β for large β. Since βj∗ = n is large, we get ΨK(q, g) = −λ̃ ∏ j 6=j∗ G(λ̃j , βj) + d∑ j=1 [ βjG(λ̃j , βj) + logB(k̃j , h̃j , βj)− log(2k̃j)− log λ̃j ]
−λ̃ log n+ λ̃(1− d∏
j=m+1
G(λ̃j , βj)) + d∑ j=m+1 [βjG(λ̃j , βj) + logB(k̃j , h̃j , βj)]− d∑ j=1 log(2k̃j)− d∑ j=1 log(λ̃j).
If βj are sufficiently large for j 6= j∗, we get ΨK(q, g) −λ̃ log n− d∑ j=1 log(2k̃j)− d∑ j=1 log(λ̃j) + d∑ j=m+1 λ̃j(1− log βj) + λ̃(1− d∏ j=m+1 λ̃j βj ).
(15)
Then there exist βj , j 6= j∗ so that equation 15 matches −λ̃ log n+ logC(d).
Theorem 5.1 critically assumes that b(ξ) ∝ 1. We do not expect this to hold in reality. For example, even for the simple one-hidden layer tanh network considered in Section 7, it does not appear that b(ξ) ∝ 1 (Watanabe, 2000). Currently, we are prevented from stating a more general version of Theorem 5.1 because there is no off-the-shelf derivation of the leading coefficient when the singular integral is of the general form ∫ [0,b]d
exp(−nξ2k̃)ξh̃b(ξ) dξ. We expect the generalization to be technically feasible but as its development requires advanced knowledge of algebraic geometry, it is beyond the scope of this paper and best left as separate investigation.
In the next section, we proceed to learn the resolution map rather than presume it is known. This is an improvement over Bhattacharya et al. (2020) which, due to the difficulty of deriving resolution maps, was limited in its single experiment to the toy neural network fw(x) = b tanh(ax) with weight w = (a, b) ∈ R2. Again, though Theorem 5.1 is stated in terms of the truncated generalized gamma mean-field family, we do not believe the assumption b(ξ) ∝ 1 is critical and proceed henceforth to work with the untruncated generalized gamma mean-field family, which we denote Q.

heading: 6 LEARNING TO DESINGULARIZE
text: When the resolution map g is known, the preceding results suggest to 1) apply the change-of-variables g(ξ) = w and 2) maximize ELBO(q, g) over the untruncated generalized gamma mean-field family Q via e.g., stochastic gradient descent. However, theoretically, the resolution map g is notoriously difficult to derive. Computationally, the recursive blow-up procedure in algebraic geometry is entirely not scalable to high dimensions.
We propose to learn the resolution map via a normalizing flow, which is commonly used to model complex distributions as the push-forward of a simple source distribution through an invertible neural network G. An interesting direction of future work would be to exploit properties of the resolution
map that may aid in the design of the normalizing flow architecture. For now, we simply make use of a common type of invertible architecture consisting of affine coupling layers. With r denoting a binary mask, a so-called affine coupling layer acts as follow:
u, v ∈ Rd, u 7→ v = (1− r) u+ r (u exp(s(r u)) + t(r u)).
The binary mask r must alternate from one affine coupling layer to the next for otherwise there would be little expressive power in the resulting network. Let Gθ be a network consisting of alternating affine coupling layers, where θ denotes the collective parameters. Since the resolution map is to be learned, there is no loss in generality to assuming j∗ = 1. We shall need the gradient with respect to the variational parameters and the normalizing flow weights as part of employing stochastic gradient descent, i.e. ∇λ,k,β,θELBO(qλ,k,β, Gθ). Note that here we abuse the notation slightly as we do not need to update β1 which should be set to the sample size n according to the proof of Theorem 5.1.
Although the source distribution in a normalizing flow can have its own trainable parameters, it is common practice to adopt a parameter-less source distribution. In particular if the source distribution is itself reparametrizable, then the learning of the associated parameters can be absorbed into the invertible transformation. The generalized gamma distribution is not easily reparametrizable for general values of λ,k,β. However for certain settings of λ, we can in fact avail ourselves to the reparametrization trick, at least approximately. Let Vj be a gamma random variable with shape λj and rate βj , then V 1/(2kj) j ∼ qj(ξj) := ξ 2kjλj−1 j exp(−βjξ 2kj j )1(0,∞)(ξj). We will mostly be interested in settings where λj is large, in which case Vj is approximately Gaussian with mean λj/βj and variance λj/β2j . Letting T ( ) := (F −1 1 ( 1) 1/(2k1), . . . , F−1d ( d) 1/(2kd)), where F−1j ( ) ≈
(λj + √ λj )/βj , the reparametrization trick then leads to the objective function
E ∼N(0,I) [ n∑ i=1 log p(yi|xi, Gθ(T ( ))) + logϕ(Gθ(T ( ))) + log |G′θ(T ( ))| ] − Eq log q. (16)
Note that equation 16 is not the same as ELBO(qλ,k,β, Gθ) because we have made use of the Gaussian approximation for Vj in the case of large λj .
The entropy component of equation 16, −Eqλ,k,β log qλ,k,β, can be derived analytically, see Appendix B.2. Next, the specific architecture of Gθ has rendered the log Jacobian term, log |G′θ(·)|, computationally tractable. The final piece is to replace E ∼N(0,I) with an empirical average over M samples. In the experiments that follow we will consider either learning the variational parameters λ,k,β or fixing them at some initial value since other values can be learned through a transformation that gets absorbed into G.

heading: 7 EXPERIMENTS
text: In this section, we compare the effect of two source distributions for a normalizing flow given by the affine coupling network Gθ consisting of 2 pairs of alternating couplings with scaling and translation networks each consisting of 2 hidden layers with 16 hidden units, see Appendix D for exact details of the architecture. We denote by nf_gamma_λ0_k0_β0_flag the variational family that results from pushing forward the untruncated generalized gamma distribution with variational parameters initialized at λ0 = (1, λ0, . . . , λ0),k0 = (k0, . . . , k0), and β0 = (n, β0, . . . , β0), and a boolean flag indicating whether the variational parameters are subject to updating. For the other approximation resulting from pushing forward a Gaussian source distribution N(µ0, v0), we write analogously nf_gaussian_µ0_v0 where µ0 and v0 are the fixed mean and variance. For each combination of λ0 and β0 considered, we set µ0 = λ0/β0 and v0 = λ0/β20 .
We employ the widely adopted variant of (minibatch) stochastic gradient descent known as Adam. The number of epochs and batch size were set to 2000 and n/10, respectively. Different constant learning rates were employed according to the parameter type: 1e−3 for affine coupling layer weights and k, and 1e−1 for λ and β. Expectations in the objective function that are not analytically tractable are replaced with an average over M = 5 samples. At the end of training, we evaluate Ψ(q, g) in nf_gamma without recourse to the (approximate) reparametrization and by using the analytic expression for the entropy component of Ψ and 100 samples from q to approximate other components in Ψ under expectation.
To date, there is a very small collection of model-truth-prior triplets where the RLCT λ̃ and multiplicity mare known. We shall limit our experiments to two such triplets. This will allow us to compare the achieved Ψ(q̂, Ĝ) following training to −λ log n+ (m− 1) log log n. Note that even in such triplets where the RLCT and multiplicity are known, the exact value of the leading coefficient C is still usually unknown. Now, to compare nf_gamma and nf_gaussian to each other, we can simply see which achieves higher Ψ(q̂, Ĝ) after training.
Looking ahead to downstream tasks, it is natural to ask whether the variational posterior predictive distribution, pvb(y|x,Dn) = 〈p(y|x,Gθ̂(ξ))〉q̂(ξ), inherits the desirable properties of the Bayes posterior predictive distribution, p(y|x,Dn) = 〈p(y|x,w)〉p(w|Dn). The answer turns out to depend on the relationship between the variational real log canonical threshold λvb and the RLCT of the model-truth-prior triplet. It may very well be that a variational family which is closer to the true posterior (in the KL sense) than another variational family may induce a worse approximation of the true Bayes posterior predictive distribution. We provide an in-depth account of this phenomenon through the lens of singular learning theory in Appendix E.
Our first example concerns the tanh model-truth-prior triplet. Consider input x ∈ R following the uniform distribution on [−1, 1], and response variable y ∈ R modeled as p(y|x,w) =
1√ 2π exp(− 12 (y − f(x,w)) 2), where fw(x) = ∑H h=1 bh tanh(ahx) is a tanh network with H hidden units and w is the collection of neural network weights {(ah, bh)}Hh=1. If the true distribution is given by p0(y|x) = p(y|x, 0) = 1√2π exp(− 1 2y 2) and the prior ϕ is a C∞ function of w with compact support, satisfying ϕ(0) > 0, Aoyagi & Watanabe (2006) showed that λ = H+i 2+i
4i+2 and m = 2 if i2 = H , and m = 1 if i2 < H where i is the maximum integer satisfying i2 ≤ H . In contrast, were this a regular statistical model, we would have λ = H . We consider two settings for p0(y|x) = p(y|x,w0) in the experiment: 1) w0 = 0 and 2) w0 = 5. Note for the latter, the corresponding RLCT and multiplicity are unknown.
Our next example is the reduced rank model-truth-prior triplet. Consider input x ∈ RM and response variable y ∈ RN modeled as p(y|x,w) = (2π)−N/2 exp{− 12 ||y−BAx||
2}where {w = (A,B)|A ∈ RH×M , B ∈ RN×H}. This model is readily seen to be a special case of a neural network with hidden units H and identity activation function. We shall set M = H + 3 and N = H . In the realizable case, i.e., p0(y|x) = p(y|x,A0, B0), if the prior ϕ is a C∞ function with compact support satisfying ϕ(A0, B0) > 0, the RLCT was derived in Aoyagi & Watanabe (2005) for various values of N,M,H and r = rank(B0A0). Below we set B0 = IN×N and A0 = [IH×H ; JH×3]. The rank r for B0A0 equals H . This then falls under Case (3) in Aoyagi & Watanabe (2005) since N + H < M + r, leading to λ = (NH −Hr+Mr)/2,m = 1. Note that were this a regular model, we would instead have λ = (MH +NH)/2.
Table 1 provides a summary of the values ofH considered in each of the triplets and the corresponding RLCT and dimension. In the experiments, we tested various combinations of values of H , settings for w0 and priors. Throughout, the priors were designed to be mis-specified. The full results of these experiments can be found in Tables 2–4. Note that the comparison between different variational approximations have been made subject to the value of H , the prior, and the relationship λ0/β0 = µ0. We should also point out that not all runs reached convergence as can be seen from the count column in Tables 2–4. This seems largely to be a result of the constant learning rate we used for all training.
The first setting (Table 2) is the realizable tanh network with w0 = 0. In this case, the true RLCT is known as discussed above. As Table 2 demonstrates, there is no clearly discernible difference between nf_gamma and nf_gaussian where the source distribution for both starts near zero. However when the source distribution has λ0/β0 = µ0 = 5, i.e. very far from w0, nf_gamma is more robust than nf_gaussian. The second setting (Table 3) is analogous to the previous except the true generating mechanism is given by w0 = 5. In this setting, the untruncated generalized gamma source distribution is seen to provide a better variational approximation than the Gaussian source distribution. This appears to hold quite independent of the initial values of λ,k,β as well as the gradient flag. Interestingly, in contrast to the previous setting, there is a clear benefit to learning λ,k,β. The third setting (Table 4) is the realizable reduced rank regression. In this experiment, we do not see a discernible difference between the source distributions until we look at λ0/β0 = µ0 = 5, which is far from w0 = (A0, B0). As in the first setting, we see nf_gamma is more robust than nf_gaussian.
Finally, though we have cautioned against judging the quality of a variational approximation according to the approximate posterior predictive distribution induced, it is nonetheless informative to visualize downstream uncertainty quantification. In Figure 1, we choose one line from the tanh w0 = 0 experiment (Table 2) to further examine. For the two methods illustrated in in Figure 1, their performance in terms of Ψ is quite close as indicated by Table 2. It is then comforting that the respective confidence bands are quite similar. Next, in Figure 2, we choose one line from the tanh w0 = 5 experiment (Table 3) to further examine. This time, nf_gamma achieves significantly better Ψ than nf_gaussian. We see a corresponding relationship between their respective posterior predictive distributions. Namely, the confidence band resulting from nf_gaussian are far too conservative while that from nf_gamma is less so.

heading: 8 CONCLUSION
text: In this work we propose a variational approximation for Bayesian neural networks by leveraging insights from singular learning theory. Namely, for large n, the posterior distribution over neural network weights is not Gaussian but rather can be put into a mixture of standard forms. From this, we demonstrate that the generalized gamma mean-field family, following desingularization, can in theory achieve the leading order term of the log normalized evidence. Because we choose to learn the desingularization map using affine coupling layers, the proposed work can be cast as a normalizing flow with an unconventional source distribution. Interestingly, for large values of the variational parameter λ, the source distribution in each dimension is approximately N(λ, λ/β)1/(2k), which is reparametrizable in terms of the conventional source distribution N(0, 1). Though learning the variational parameters λ,k,β in the source distribution goes against conventional wisdom in normalizing flows, our experiments suggest some performance may be gained by learning the optimal untruncated generalized gamma source distribution at the same time as learning the resolution map.
ETHICS STATEMENT
Bayesian learning for neural networks is often touted as a panacea to the challenges of uncertainty quantification in deep learning. In this work, we propose a variational approximation that performs well in terms of the ELBO achieved. However, we make no claim that our variational approximation is superior in downstream uncertainty quantification. We explain extensively in the appendix that a variational approximation which achieves higher ELBO may induce a worse approximate posterior predictive distribution than a variational approximation which achieves lower ELBO. According to singular learning theory, the performance of a variational approximation in terms of the posterior predictive distribution it induces critically depends on the relationship between the variational RLCT and the underlying RLCT of the model-truth-prior triplet.
REPRODUCIBILITY STATEMENT
A zip file of the source code has been submitted as supplementary materials. For theoretical results, explanations of assumptions and proof of the claims are included in the main text. No datasets are used; the experiments are solely based on simulated data.

heading: A TOY EXAMPLE OF RLCT CALCULATION
text: We recall Example 27 from Watanabe (2018) to illustrate the concepts of resolution map, RLCT and multiplicity for a simple model-truth-prior triplet. For univariate input x ∈ [0, 1] and univariate output y ∈ R, consider the model with parameter w = (a, b) ∈ [0, 1]2 given by
p(x, y|w) = 1√ 2π exp(−1 2 (y − a tanh(bx))2) (17)
Suppose the prior is uniform, i.e., ϕ(w) = 1 and the truth is given by p0(x, y) = p(x, y|0, 0). Then we can easily see that
K(w) = b2a2 1
2 K0(w),
where
K0(w) = ∫ 1 0 ( tanh(bx) b )2 dx
The following desingularization map puts the triplet in standard form:
ξ1 =
√ K0(w)
2 a
ξ2 = b
Furthermore we have ϕ(g(ξ)) = ξh where h = (0, 0) and b(ξ) = |g′(ξ)|. Since (k1, k2) = (1, 1) and (h1, h2) = (0, 0) we have (λ1, λ2) = (1/2, 1/2). Therefore for this particular model-truth-prior triplet, the RLCT is 1/2 with multiplicity 2.

heading: B BASIC CALCULATIONS FOR THE GENERALIZED GAMMA DISTRIBUTION
text: B.1 TRUNCATED
We summarize some basic calculations for the [0, 1] truncated generalized gamma distribution from Bhattacharya et al. (2020). Let qj be the univariate truncated generalized gamma density given in equation 12.
• The normalizing constant of qj is given byB(λj , kj , βj) whereB(λ, k, β) = β−λΓ(λ)γ(λ,β)
2k and γ(a, x) = 1Γ(a) ∫ x 0 ta−1e−t dt is the (regularized) lower incomplete gamma function.
• The quantity Eqjξ 2kj = G(λj , βj) where G(λ, β) = λβ γ(λ+1,β) γ(λ,β) .
B.2 UNTRUNCATED
Consider the univariate density
q(ξ) ∝ ξh exp(−βξ2k)
supported on (0,∞) where h, k, β > 0. Let λ = (h+ 1)/(2k). Elementary calculations give
• Eqξ2k̃ = β− k̃ k
Γ( k̃k+λ)
Γ(λ)
• Eq log ξ = 12k (ψ(λ)− log β) where ψ is the digamma function.
• Eq log q = h2k (ψ(λ)− log β)− λ− logZ where, Z = β−λΓ(λ)
2k is the normalizing constant of q.

heading: C GENERALIZATION OF (BHATTACHARYA ET AL., 2020, THEOREM 3.1)
text: Lemma C.1 below states that any variational distribution qλ,k,β in the untruncated generalized gamma mean field family Q with well-specified k where we additionally set βj∗ = n is capable of achieving ΨK(q, g) ≥ −λ̃ log n + A where A is some constant that does not depend on n. Lemma C.1 is a straightforward extension of Theorem 3.1 of Bhattacharya et al. (2020)from the [0, 1]d truncated generalized gamma mean field variational family to the untruncated case. We also relax the condition in Theorem 3.1 of Bhattacharya et al. (2020) that the variational parameters λ and k (and hence h) are well-specified; here we only require that k be well-specified.
Lemma C.1. Suppose the model-truth-prior triplet is such that Theorem 3.1 holds with K(g(ξ)) = ξ2k̃ and ϕ(g(ξ))|g′(ξ)| = ξh̃. Consider the untruncated generalized gamma mean-field family Q, we have supq∈QΨK(q, g) ≥ −λ̃ log n+A, where A is a constant that does not depend on n.
Proof. If kj = k̃j for all j = 1, . . . , d, we have
EqnK(g(ξ)) = n d∏ j=1 β − k̃j kj j Γ( k̃j kj + λj) Γ(λj) = n 1− k̃j∗ kj∗ Γ( k̃j∗ kj∗ + λj∗) Γ(λj∗) ∏ j 6=j∗ β − k̃j kj j Γ( k̃j kj + λj) Γ(λj) = λj∗ ∏ j 6=j∗ λj βj .
Next, we have
KL(q(ξ)||ϕ(g(ξ))|g′(ξ)|)
= d∑ j=1 Eqj [log qj − h̃j log ξj ]− Eq log b(ξ)
= d∑ j=1
[ hj − h̃j
2kj (ψ(λj)− log βj)− λj + λj log βj + log(2kj)− log Γ(λj)
] − Eq log b(ξ)
= d∑ j=1
[ hj − h̃j
2kj ψ(λj)− λj +
h̃j + 1
2kj log βj + log(2kj)− log Γ(λj)
] − Eq log b(ξ)
= d∑ j=1
[ hj − h̃j
2k̃j ψ(λj)− λj + λ̃j log βj + log(2k̃j)− log Γ(λj)
] − Eq log b(ξ)
= λ̃ log n+ ∑ j 6=j∗ λ̃j log βj + d∑ j=1
[ hj − h̃j
2k̃j ψ(λj)− λj + log(2k̃j)− log Γ(λj)
] − Eq log b(ξ).
The last line follows from setting βj∗ = n where j∗ ∈ {1, . . . , d} is such that λ̃j∗ = λ̃, breaking ties arbitrarily. Finally, the desired inequality follows from the fact that Eq log b(ξ) is bounded below by some constant.

heading: D EXPERIMENTS
text: In the experiments, we train an affine coupling network with two pairs of alternating couplings. The translation t is a feedforward (leaky) ReLU neural network with tanh output activation function. The scaling t is another feedforward (leaky) ReLU neural network with identity output activation function. The models implemented are given below.
Translation network
Sequential( (0): Linear(in_features=dim, out_features=16, bias=True) (1): LeakyReLU(negative_slope=0.01) (2): Linear(in_features=16, out_features=16, bias=True) (3): LeakyReLU(negative_slope=0.01) (4): Linear(in_features=16, out_features=16, bias=True) (5): LeakyReLU(negative_slope=0.01) (6): Linear(in_features=16, out_features=dim, bias=True) )
Scaling network
Sequential( (0): Linear(in_features=dim, out_features=16, bias=True) (1): LeakyReLU(negative_slope=0.01) (2): Linear(in_features=16, out_features=16, bias=True) (3): LeakyReLU(negative_slope=0.01) (4): Linear(in_features=16, out_features=16, bias=True) (5): LeakyReLU(negative_slope=0.01) (6): Linear(in_features=16, out_features=dim, bias=True) (7): Tanh() )

heading: E APPROXIMATE POSTERIOR PREDICTIVE DISTRIBUTION
text: Bayesian prediction proceeds by marginalization, i.e., averaging over all possible values of the model parameter. This results in the posterior predictive distribution1,
p(y|x,Dn) := ∫ w p(y|x,w)p(w|Dn) dw = 〈p(y|x,w)〉p(w|Dn). (18)
According to singular learning theory, Bayesian prediction via equation 18 is superior to MAP and MLE, in the sense that the expected generalization error of the posterior predictive distribution is smaller than that of the MAP or MLE. Specifically, let
Gn(p̂n(y|x)) := KL(p0(y|x)p(x)||p̂n(y|x)p(x))
be the generalization error of a predictive distribution p̂n(y|x). According to Theorems 1.2 and 7.2 in Watanabe (2009), we have
EGn(p(y|x,Dn)) = λ/n+ o(1/n) (19)
where the expectation is taken with respect to Dn and λ is the RLCT of the model-truth-prior triplet. On the other hand, Theorem 6.4 of Watanabe (2009) shows that the expected generalization error of MLE (and similarly of MAP) is
EGn(p(y|x, ŵmle)) = C/n+ o(1/n) (20)
where C, the maximum of a Gaussian process, can be much larger than λ. Such a distinction cannot be made in regular models in which the difference between the three estimators becomes negligible in the large n regime.
Now, consider the variational approximation to the posterior predictive distribution given by
pvb(y|x,Dn) = 〈p(y|x,w)〉q∗(w),
where q∗ is the optimal variational distribution in some variational family Q. It may be tempting to compare different variational approximations according to how well their respective variational posterior predictive distributions approximate the true Bayes posterior predictive distribution, e.g., Blundell et al. (2015); Louizos & Welling (2017). This turns out to be a thorny issue, as documented in various works on Bayesian neural networks Heek (2018); Krishnan & Tickoo (2020); Foong et al. (2020). In particular, pvb(y|x,Dn) does not necessarily inherit the desirable properties of the Bayes posterior predictive distribution p(y|x,Dn). Here, we offer insights from singular learning theory to account for this.
Consider the (normalized) variational free energy
F̄vb(n) = EqnKn(w) + KL(q(w)||ϕ(w))
where q is some variational distribution. If the minimum (normalized) variational free energy F̄ ∗vb(n) = minq F̄vb(n) admits an asymptotic expansion 2, then it would be of the form
F̄ ∗vb(n) = λvb log n+ (mvb − 1) log log n+Rn.
Note that λvb ≥ λ necessarily. Now, if the generalization error of pvb(y|x,Dn) would admit an asymptotic expansion3, it would be of the form
EGn(pvb(y|x,Dn)) = λ̃vb/n+ o(1/n). 1Though this may seem fundamentally distinct from the maximum likelihood estimator (MLE) or maximum a posterior (MAP) solution commonly employed in training deep networks, both MLE and MAP may in fact be regarded as impoverished estimates of equation 18 whereby p(w|Dn) is approximated with δ(w = ŵ), a point mass at ŵ.
2We should disclose that general conditions for such an asymptotic expansion of the minimum variational free energy is still an open problem. The issue has so far been addressed on a case-by-case basis, e.g., reduced rank regression Nakajima & Watanabe (2007), nonnegative matrix factorization Kohjima & Watanabe (2017); Hayashi (2020), normal mixture model Watanabe & Watanabe (2006), hidden Markov model Hosino et al. (2005).
3Similar to the minimum variational free energy, an expansion of EGn(pvb(y|x,Dn)) has not been established in full generality at this point.
Importantly, λ̃vb 6= λvb in general. Examples can be found where sometimes one is bigger, sometimes the other Nakajima & Watanabe (2007).
This development is to be contrasted with the Bayes free energy and the Bayes generalization error. Recall the normalized Bayes free energy, F̄n := − log Z̄(n), admits the following asymptotic expansion F̄n = λ log n+ (m− 1) log log n+Rn. (21) The very same λ coefficient appears in the asymptotic expansion of the Bayes generalization error in equation 19. This means that minimizing the Bayes free energy is equivalent to minimizing the Bayes generalization error. In contrast, minimizing the variational free energy (equivalent to maximizing the ELBO) does not necessarily lead to a lower generalization error since a variational family with higher variational free energy (higher λvb) may have lower variational generalization error (lower λ̃vb).

Review summaries:
Inspired by asymptotic results from singular learning theory, the authors of the paper propose using a generalised gamma mean-field distribution with a normalising flow (that targets the "desingularization map") to perform variational inference. The authors additional build on prior work to derive a tighter bound on the log normalised evidence for variational inference using this approximate posterior combined with the "correct" desingularization map.   Note: while I am familiar with the Bayesian Neural Network literature, I am not familiar with singular learning theory. This is reflected in my confidence score. As such, I was unable to verify the theory in the paper, and assess the significance of the theoretic results included in the paper.

The paper analyzes the approximation of the log-ratio (called normalized evidence in the paper) (the log-probability under the model minus the log-probability under the true model). Since the model is singular, the Laplacian approximation of the normalized evidence needs a new method. The mean-field distribution and a resolution map allow us to estimate two important numbers (lambda and m) in the approximation.   The paper has made several contributions to the study of this problem in several aspects. First, the paper improves previous results by relaxing the assumption. Second, the paper actually uses a normalizing flow to actually learn the transformation.    

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.

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.