Variational Bayesian Neural Networks via Resolution of Singularities
In this work, we advocate for the importance of singular learning theory
(SLT) as it pertains to the theory and practice of variational inference in
Bayesian neural networks (BNNs). To begin, using SLT, we lay to rest some of
the confusion surrounding discrepancies between downstream predictive
performance measured via e.g., the test log predictive density, and the
variational objective. Next, we use the SLT-corrected asymptotic form for
singular posterior distributions to inform the design of the variational family
itself. Specifically, we build upon the idealized variational family introduced
in \citet{bhattacharya_evidence_2020} which is theoretically appealing but
practically intractable. Our proposal takes shape as a normalizing flow where
the base distribution is a carefully-initialized generalized gamma. We conduct
experiments comparing this to the canonical Gaussian base distribution and show
improvements in terms of variational free energy and variational generalization
error.
Susan Wei, Edmund Lau
# Variational Bayesian Neural Networks via Resolution of Singularities

###### Abstract

In this work, we advocate for the importance of singular learning theory (SLT) as it pertains to the theory and practice of variational inference in Bayesian neural networks (BNNs). To begin, using SLT, we lay to rest some of the confusion surrounding discrepancies between downstream predictive performance measured via e.g., the test log predictive density, and the variational objective. Next, we use the SLT-corrected asymptotic form for singular posterior distributions to inform the design of the variational family itself. Specifically, we build upon the idealized variational family introduced in Bhattacharya et al. (2020) which is theoretically appealing but practically intractable. Our proposal takes shape as a normalizing flow where the base distribution is a carefully-initialized generalized gamma. We conduct experiments comparing this to the canonical Gaussian base distribution and show improvements in terms of variational free energy and variational generalization error.

Normalizing Flow Real Log Canonical Threshold Singular Learning Theory Singular Models Test log-likelihood Variational Free Energy Variational Inference Variational Generalization Error

## 1 Introduction

A Bayesian neural network (BNN) Mackay (1995) is a neural network endowed with a prior distribution \(\varphi\) on its weights \(w\). Despite their theoretical appeal Lampinen and Vehtari (2001); Wang and Yeung (2020), applying BNNs in practice is not without significant challenges. MCMC and its variants, while widely considered the gold standard, can be prohibitively expensive in terms of computation. On the other hand, fast alternatives such as variational inference may result in _uncontrolled_ approximations.

In this work, we mine insights from **singular learning theory** (SLT) Watanabe (2009) to explain and improve upon certain aspects of BNNs. Roughly speaking, a model is (strictly) **singular** if the parameter-to-model mapping is not one-to-one and the likelihood function does not look Gaussian1. That neural networks are singular is well documented Sussmann (1992); Watanabe (2000, 2001); Fukumizu (2003); Watanabe (2007). We refer the readers to Wei et al. (2022) for a detailed proof in the case of a standard feedforward network. The singular nature of BNNs has interesting implications for the posterior distribution, see Figure 1.

Footnote 1: These features should not be viewed as pathological, see “Deep learning is singular and that’s good” by Wei et al. (2022).

Let \((x,y)\) denote the input-target pair modeled jointly as \(p(x,y|w)=p(y|x,w)p(x)\) where \(w\in\mathbb{R}^{d}\) is the model parameter. Let \(p(y|x,w)\) be a neural network model with functional model \(f\), by which we mean \(y=f(x,w)+\epsilon\) where \(\epsilon\) is some random variable. For example, if we have Gaussian additive noise \(\epsilon\), the conditional distribution could be modelled as \(\mathcal{N}(y|f(x,w),\sigma^{2}I)\) where \(f\) is a feedforward ReLU network with weights \(w\).

The central quantity of interest in BNNs is the intractable posterior distribution over the neural network weights,

\[p(w|\mathcal{D}_{n})=\frac{\prod_{i=1}^{n}p(y_{i}|x_{i},w)\varphi(w)}{Z(n)},\]
where \(\mathcal{D}_{n}=\{(x_{i},y_{i})\}_{i=1}^{n}\) is a dataset of \(n\) input-output pairs. The normalizing constant,

\[Z(n)=\int\prod_{i=1}^{n}p(y_{i}|x_{i},w)\varphi(w)\,dw,\]

is variously known as the **model evidence** and the **marginal likelihood**. Define the **empirical entropy** of the training data,

\[S_{n}=-\frac{1}{n}\sum_{i=1}^{n}\log p_{0}(y_{i}|x_{i}).\]

We shall call

\[\bar{Z}(n)=-\log Z(n)-nS_{n}\]

the **normalized evidence**. Let us call \(F(n):=-\log Z(n)\) the **Bayes free energy** and \(\bar{F}(n):=-\log\bar{Z}(n)\) its normalized version.

Unlike prediction in traditional neural networks, prediction in BNNs proceeds by marginalization, i.e., averaging over all possible values of the network weights. Namely, prediction in BNNs makes use of the **Bayes posterior predictive distribution**,

\[p(y|x,\mathcal{D}_{n}):=\int p(y|x,w)p(w|\mathcal{D}_{n})\,dw. \tag{1}\]

With (1), we can calculate prediction uncertainties as well as obtain better calibrated predictions Heek (2018); Osawa et al. (2019); Maddox et al. (2019).

In Section 3, we recapitulate from the perspective of SLT the predictive advantages of BNNs over traditional neural networks. Specifically, SLT shows that the Bayes posterior predictive distribution in (1) has lower generalization error compared to MLE or MAP point estimates.

Despite compelling arguments for employing BNNs, we must reckon with the fact that they can only ever be applied _approximately_. Among approximate techniques, a major class is represented by scaling classic MCMC to modern settings of large datasets and deep neural networks Welling and Teh (2011); Chen et al. (2014); Zhang et al. (2020). In this paper, we instead turn our focus to variational inference, which is particularly suited to scaling BNNs to large datasets.

All variational inference techniques are characterized by two ingredients. First, a family of densities \(\mathcal{Q}\), often called the variational family, is posited. Second, some \(q^{*}\in\mathcal{Q}\) is found via optimization according to some criterion that measures

Figure 1: Posterior density contour plot for a 2D \(\tanh\)-regression model, \(p(y|x,a,b)\propto\exp(y-a\tanh(bx))\). The white diamond marks the true parameter \((a_{0},b_{0})\) used to generate the dataset \(\mathcal{D}_{n}\). Each row shows a different true distribution, while each column shows a different sample size \(n\). When \(a_{0}b_{0}=0\) as in the second row, the set of true parameters \(W_{0}\) is not a singleton and contains a singularity at the origin. It is worth noticing that, for a singular model, even when the truth is not at a singularity (first row), the posterior is still far from being locally Gaussian even at sample size \(n=5000\).
closeness to the desired target density. In this work, we seek to approximate the posterior density and we will employ the conventional Kullback-Leibler divergence. This leads to the optimization problem,

\[\min_{q\in\mathcal{Q}}\mathrm{KL}(q(w)\parallel p(w|\mathcal{D}_{n})). \tag{2}\]

This is equivalent to minimizing the so-called **normalized2 variational free energy (VFE)**,

Footnote 2: Throughout this paper, we work with normalized quantities for ease of exposition. The asymptotics presented hold equally for the unnormalized counterparts.

\[\bar{F}_{vb}(n):=\mathbb{E}_{q}nK_{n}(w)+\mathrm{KL}(q(w)\parallel\varphi(w)).\]

It is easy to see that \(\bar{F}_{vb}(n)\geq\bar{F}(n)\) with equality if and only if the variational distribution is exactly equal to the posterior. Readers are likely more familiar with the variational objective of maximizing the so-called **evidence lower bound (ELBO)** which is simply related to the (normalized) VFE via \(\mathrm{ELBO}=-\bar{F}_{vb}(n)\).

Let \(q^{*}\in\mathcal{Q}\) be a minimizer of (2). et us call the variational approximation to (1) given by

\[p_{vb}(y|x,\mathcal{D}_{n}):=\int p(y|x,w)q^{*}(w)\,dw, \tag{3}\]

the **induced predictive distribution**. We can measure the predictive accuracy of \(p_{vb}\) using once again the KL divergence, i.e.,

\[G_{n}(p_{vb}(y|x,\mathcal{D}_{n})):=\mathrm{KL}(p_{0}(y|x)\parallel p_{vb}(y| x,\mathcal{D}_{n})),\]

which we shall call the **variational generalization error (VGE)**. Per the discussion in Section 3, this is, up to a constant and a sign flip, nothing more than the typical **test log predictive density**Gelman et al. (2014) commonly employed in variational inference evaluation.

We shall see in Section 4 that, surprisingly, the VGE may be arbitrarily high even for a variational family whose minimum VFE is close to optimality. In other words, it is not guaranteed that minimizing (2) results in good downstream predictive performance. The outlook is not entirely bleak. Depending on the relationship between two critical quantities of variational inference - the **MVFE coefficient**\(\lambda_{\textbf{vfe}}\) and the **VGE coefficient**\(\lambda_{\textbf{vge}}\) - the generalization error of the induced predictive distribution may be controllable via minimizing the VFE.

Clarification of the relationship between the two variational coefficients for most common variational learning problems is an open problem, which we leave aside for future work. We will assume the variational coefficients are related _favorably_, in a manner which will be made clear in Section 4, and proceed to design a variational family whose **variational approximation gap** is small. The proposal is predicated on an important SLT result which states that, roughly speaking, the posterior distribution over the parameters of a singular model is not asymptotically Gaussian, but can still be put into an explicit standard form via the **resolution of singularities**.

## 2 Singular learning theory

In this section, we give a succinct overview of key concepts from SLT. We focus in particular on what SLT has to say about the behavior of the posterior distribution in strictly singular models. Let us assume the parameter space \(W\) is a compact set in \(\mathbb{R}^{d}\) and \(p_{0}(x,y)=p_{0}(y|x)p(x)\) is the true data-generating mechanism. Throughout, we suppose there exists \(w_{0}\in W\) such that \(p_{0}(y|x)=p(y|x,w_{0})\). In the parlance of SLT, this condition is known as **realizability**. Let \(\varphi(w)\) be a compactly-supported prior. We shall refer to \((p(\cdot,\cdot),p_{0}(\cdot,\cdot),\varphi(\cdot))\) as a **model-truth-prior triplet**. The roles played by compactness and realizability in singular learning theory are discussed in Appendix A.

Define \(K(w)\) to be the Kullback-Leibler divergence between the truth and the model, i.e.,

\[K(w):=\mathrm{KL}(p_{0}(x,y)\parallel p(x,y|w)).\]

Following Watanabe (2009), we say a model is **regular** if 1) it is identifiable, i.e., the map \(w\mapsto p(\cdot,\cdot|w)\) from parameter to model is one-to-one and 2) its Fisher information matrix \(I(w)\) is positive definite for arbitrary \(w\in W\). We call a model **strictly singular** if it is not regular. The term singular will refer to either regular or strictly singular models. See Figure 1 for an example of a strictly singular model with two truth settings. This figure illustrates an important lesson: for strictly singular models, even when the true parameter set \(W_{0}:=\{w:K(w)=0\}\) does not contain singularities, the posterior distribution is still far from Gaussian.

The following theorem from Watanabe (2009), adapted for notational consistency, gives precise conditions for the existence of **resolution maps**, algebraic-geometrical transformations which enables \(K(w)\) to be locally written as a
monomial_, i.e., a product of powers of variables such as in the right-hand-side of (4). The result is itself based on Hironaka's resolution of singularities, a celebrated result in modern algebraic geometry.

To prepare, let \(W_{\epsilon}=\{w\in W:K(w)\leq\epsilon\}\) for some small positive constant \(\epsilon\) and \(W_{\epsilon}^{(R)}\) be some real open set such that \(W_{\epsilon}\subset W_{\epsilon}^{(R)}\). The theorem below will make use of the multi-index notation: for a given \(\xi=(\xi_{1},\ldots,\xi_{d})\in\mathbb{R}^{d}\), define \(w^{\mathbf{k}}:=w_{1}^{k_{1}}\cdots w_{d}^{k_{d}}\) where the multi-index \(\mathbf{k}=(k_{1},\ldots,k_{d})\) with each \(k_{j}\) a nonnegative integer. Due to space constraints, Fundamental Conditions I and II required below are stated and discussed in Appendix A.

**Theorem 2.1** (Theorem 6.5 of Watanabe [2009]).: _Suppose the model-truth-prior triplet \((p,p_{0},\varphi)\) satisfies Fundamental Conditions I and II with \(s=2\). We can find a real analytic manifold \(M^{(R)}\) and a proper and real analytic map \(g:M^{(R)}\to W_{\epsilon}^{(R)}\) such that_

1. \(M=g^{-1}(W_{\epsilon})\) _is covered by a finite set_ \(M=\cup_{\alpha}M_{\alpha}\) _where_ \(M_{\alpha}=[0,b]^{d}\)_._
2. _In each_ \(M_{\alpha}\)_,_ \[K(g(\xi))=\xi^{\mathbf{2k}}=\xi_{1}^{2k_{1}}\cdots\xi_{d}^{2k_{d}},\] (4) _where_ \(k_{j}\in\mathbb{N},j=1,\ldots,d\) _are such that not all_ \(k_{j}\) _are zero._
3. _There exists a_ \(C^{\infty}\) _function_ \(b(\xi)\) _such that_ \[\varphi(g(\xi))|g^{\prime}(\xi)|=\xi^{\mathbf{h}}b(\xi)=\xi_{1}^{h_{1}}\cdots\xi_{ d}^{h_{d}}b(\xi),\] (5) _where_ \(h_{j}\in\mathbb{N},j=1,\ldots,d\)_,_ \(|g^{\prime}(\xi)|\) _is the absolute value of the determinant of the Jacobian and_ \(b(\xi)>c>0\) _for_ \(\xi\in[0,b]^{d}\)_._

In Theorem 2.1 we have suppressed the dependency on the manifold chart index \(\alpha\), but the reader should keep in mind that the maps \(g\) and the multi-indices are all indexed by \(\alpha\). It is also important to recognize that none of these said quantities are unique for a given triplet \((p,p_{0},\varphi)\).

A crucial quantity that appears in SLT is a rational number in \((0,d/2]\) known as the **real log canonical threshold** (RLCT). Let \(\{M_{\alpha}:\alpha\}\) be as in Theorem 2.1 and define

\[\lambda_{j}=\frac{h_{j}+1}{2k_{j}},j=1,\ldots,d\]

where \(h_{j}\) and \(k_{j}\) are the entries of the multi-indices \(\mathbf{h}\) and \(\mathbf{k}\) in a local coordinate \(M_{\alpha}\). When \(k_{j}=0\), \(\lambda_{j}\) is taken to be infinity.

Uniquely associated to a triplet \((p,p_{0},\varphi)\) are its real log canonical threshold (RLCT) and its multiplicity defined, respectively, as

\[\lambda=\min_{\alpha}\min_{j\in 1,\ldots,d}\lambda_{j},\quad m=\max_{\alpha} \#\{j:\lambda_{j}=\lambda\}. \tag{6}\]

Let \(\{\alpha^{*}\}\) be the set of those local coordinates in which both the \(\min\) and \(\max\) in (6) are attained. Watanabe [2009] calls this set the **essential coordinates** and the corresponding collection \(\{M_{\alpha}\}\) the **essential charts**.

If \(\{w:K(w)=0,\varphi(w)>0\}\) is not the empty set, 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]. In fact, (twice the) RLCT may be regarded as the effective degrees of freedom in strictly singular models [Wei et al., 2022]. The RLCT also shows up in important asymptotic results, see (10) and (11).

Henceforth, to make clear that the RLCT and multiplicity are invariants of the model-truth-prior triplet, we shall write \(\lambda(p,p_{0},\varphi)\) and \(m(p,p_{0},\varphi)\) to mark this dependence. In Appendix B, we recall a simple toy network, a two-parameter \(\tanh\) network, where the resolution map, the RLCT, and the multiplicity can be calculated explicitly.

### Posterior distribution in singular models

The posterior distribution in strictly singular models is decidedly not Gaussian. The correct asymptotic form can be derived using SLT. For a particular manifold chart index \(\alpha\), let us apply the transformation \(g_{\alpha}(\xi)=w\) and rewrite the posterior distribution in the new coordinate \(\xi\),

\[p(\xi|\mathcal{D}_{n})=\frac{\exp(-nK_{n}(g_{\alpha}(\xi)))\varphi(g_{\alpha}( \xi))|g^{\prime}_{\alpha}(\xi)|}{\bar{Z}(n)}, \tag{7}\]
with

\[K_{n}(w)=\frac{1}{n}\sum_{i=1}^{n}\log\frac{p_{0}(y_{i}|x_{i})}{p(y_{i}|x_{i},w)}\]

denoting the sample average log likelihood ratio. Note \(K_{n}(w)\) is the empirical counterpart to \(K(w)\).

By (cheekily) substituting (4) (5) into \(p(\xi|\mathcal{D}_{n})\) in (7), we obtain that the posterior distribution for large \(n\), in the chart \(M_{\alpha}\), is described as a so-called **standard form**Watanabe (2018):

\[\exp(-n\xi_{1}^{2k_{1}}\xi_{2}^{2k_{2}}\cdots\xi_{d}^{2k_{d}})|\xi_{1}^{h_{1}} \cdots\xi_{d}^{h_{d}}|b(\xi).\]

In other words, the posterior distribution over the parameters of a singular model can be transformed into a mixture of standard forms, asymptotically. In Figure 1 we display the singular posterior density contour plot for a toy 2D \(\tanh\)-neural network in two settings of the true distribution.

## 3 The Bayes posterior predictive distribution

Let the generalization error of some predictive distribution \(\hat{p}_{n}(y|x)\), estimated from a training set \(\mathcal{D}_{n}\), be measured using the KL divergence:

\[G_{n}(\hat{p}_{n}(y|x)):=\mathrm{KL}(p_{0}(y|x)p(x)\parallel\hat{p}_{n}(y|x)p (x)) \tag{8}\]

In the machine learning community, this goes by another name: \(G_{n}(\cdot)\) is, up to a constant and a sign flip, the population counterpart to the commonly reported test log-likelihood, aka the predictive log-likelihood or **test log-predictive density**. This can be seen by writing

\[\hat{G}_{n}=-\frac{1}{n^{\prime}}\sum_{(x,y)\in\mathcal{D}_{n^{\prime}}}(\log p _{0}(y|x)-\log\hat{p}_{n}(y|x)) \tag{9}\]

where \(\mathcal{D}_{n^{\prime}}\) is an independent dataset.1 According to Theorems 1.2 and 7.2 in Watanabe (2009), we have, for the Bayes posterior predictive distribution (1),

\[\mathbb{E}G_{n}(p(y|x,\mathcal{D}_{n}))=\lambda(p,p_{0},\varphi)/n+o(1/n) \tag{10}\]

where the expectation is taken with respect to \(\mathcal{D}_{n}\). We will call the left hand side (10) the expected **Bayes generalization error**. This can be contrasted to the expected generalization error of MLE (and similarly of MAP), which Theorem 6.4 of Watanabe (2009) shows to be \(\mathbb{E}G_{n}(p(y|x,\hat{w}_{\text{\emph{m}le}}))=S/n+o(1/n)\) where \(S\), the maximum of a Gaussian process, can be much larger than \(\lambda(p,p_{0},\varphi)\). The situation is markedly different for regular models, where differences between the three estimators become negligible in the large-\(n\) regime.

We briefly outline the derivation of (10) as it will inform the narrative on the VGE in the next section. First, for the normalized Bayes free energy, under the Fundamental Conditions I and II discussed in A, it was proven in (Watanabe, 2009, Main Theorem 6.2) that the following asymptotic expansion holds

\[\bar{F}(n)=\lambda(p,p_{0},\varphi)\log n+(m-1)\log\log n+O_{P}(1). \tag{11}\]

The result in (10) is then proven using the above expansion together with the well known relationship between the Bayes generalization error and the (normalized) Bayes free energy (Watanabe, 2009, Theorem 1.2):

\[\mathbb{E}G_{n}(p(y|x,\mathcal{D}_{n}))=\mathbb{E}\bar{F}(n+1)-\mathbb{E}\bar {F}(n). \tag{12}\]

where on the right-hand side, the first expectation is with respect to dataset \(\mathcal{D}_{n+1}\) and the second \(\mathcal{D}_{n}\). Due to this relationship, the Bayes free energy shares the _same_ coefficient as the Bayes generalization error.

## 4 A tale of two variational coefficients

Most applications of variational inference in BNNs labor under the following implicit assumptions: 1) optimizers of the variational objective in (2) have good induced predictive distributions, and 2) two variational families can be compared according to the performance of their induced predictive distributions. A look at the experimental sections of various works on variational BNNs reveal that these assumptions underlie standard practice Blundell et al. (2015), Rezende and Mohamed (2015), Louizos and Welling (2016, 2017), Osawa et al. (2019), Swiatkowski et al. (2020). We shall see in this section that these two assumptions do not always hold.

Let us associate to a variational family \(\mathcal{Q}\) its **normalized minimum variational free energy (MVFE)**,

\[\bar{F}^{*}_{vb}(n):=\min_{q\in\mathcal{Q}}\bar{F}_{vb}(n).\]
Asymptotics for the MVFE have so far been addressed on a case-by-case basis for certain models and certain variational families, e.g., Gaussian mean-field variational families for reduced rank regression Nakajima and Watanabe (2007), nonnegative matrix factorization Kohijima and Watanabe (2017), Hayashi (2020), normal mixture model Watanabe and Watanabe (2006), hidden Markov model Hosino et al. (2005). In all the cited instances above, the asymptotic expansion of the **average normalized MVFE** takes the form

\[\mathbb{E}\bar{F}_{vb}^{*}(n)=\lambda_{\text{\sf rfe}}\log n+o(\log n) \tag{13}\]

where the expectation is taken over datasets \(\mathcal{D}_{n}\). Note that \(\lambda_{\text{\sf rfe}}\geq\lambda(p,p_{0},\varphi)\) necessarily Nakajima and Watanabe (2007). Because the **variational approximation gap**,

\[\mathcal{G}:=\bar{F}_{vb}^{*}(n)-\bar{F}(n), \tag{14}\]

is the difference of the (normalized) MVFE and the (normalized) Bayes free energy, the gap boils down to the difference between two coefficients:

\[\mathcal{G}\approx(\lambda_{\text{\sf rfe}}-\lambda(p,p_{0},\varphi))\log n.\]

Now, under some natural conditions3, the VGE admits the asymptotic expansion,

Footnote 3: The predictive distribution should be consistent as \(n\) goes to infinity, see the discussion in Chapter 13 of Nakajima et al. (2019)

\[\mathbb{E}G_{n}(p_{vb}(y|x,\mathcal{D}_{n}))=\lambda_{\text{\sf rge}}/n+o(1/n). \tag{15}\]

Importantly, \(\lambda_{\text{\sf rge}}\neq\lambda_{\text{\sf rfe}}\) in general, e.g., Nakajima and Watanabe (2007). This is in contrast to the Bayesian posterior predictive distribution in (1), where the coefficient of the leading \(O(1/n)\) term is precisely the RLCT, \(\lambda(p,p_{0},\varphi)\). That \(\lambda_{\text{\sf rge}}\neq\lambda_{\text{\sf rfe}}\) results from the fact that the relationship (12) is not valid when a variational approximation to the posterior is employed.

In Figure 1(a), we illustrate the three possible configurations of the coefficients \(\lambda(p,p_{0},\varphi),\lambda_{\text{\sf rfe}},\lambda_{\text{\sf rge}}\) for a given variational family \(\mathcal{Q}\) and a model-truth-prior triplet. When \(\lambda_{\text{\sf rfe}}>\lambda_{\text{\sf rge}}\), we call the setting **favorable** since minimizing the VFE offers control over the VGE. When \(\lambda_{\text{\sf rfe}}<\lambda_{\text{\sf rge}}\), we call this **unfavorable** since achieving even a small variational approximation gap could result in an induced predictive distribution with high generalization error. The distribution of favorable versus unfavorable settings in practice is unclear, as the exact relationship between \(\lambda_{\text{\sf rfe}}\) and \(\lambda_{\text{\sf rge}}\) has been derived in a limited number of works. The results in Nakajima and Watanabe (2007) on linear neural networks, aka reduced rank regression, show there are both favorable and unfavorable settings depending on the input and output dimension, the number of hidden units, and a rank measurement on the truth.

Note that even in favorable settings, we must be careful when comparing two variational families \(\mathcal{Q}_{1}\) and \(\mathcal{Q}_{2}\). Figure 1(b) illustrates a scenario where the family \(\mathcal{Q}_{1}\) incurs a smaller variational approximation gap than \(\mathcal{Q}_{2}\), but the induced predictive distribution of \(\mathcal{Q}_{1}\) has \(\lambda_{\text{\sf rge}}\) higher than that of \(\mathcal{Q}_{2}\). This shows that comparing different variational approximations by their test log predictive density is fraught with potential misinterpretations. In order to control the downstream predictive performance, it is thus important to find a variational family with a small approximation gap, so that we can inherit (and sometimes even beat!) the predictive advantages of the exact Bayes posterior predictive distribution (1), i.e., achieve \(\lambda_{\text{\sf rge}}<\lambda(p,p_{0},\varphi)\).

## 5 Related work

Although the perspective on offer here - that the discrepancy between test log predictive density and the variational objective _amounts to the relationship between two variational coefficients_ - is novel, we are not the first to point out this general phenomenon in variational inference Yao et al. (2018), Huggins et al. (2020), Deshpande et al. (2022), Dhaka et al. (2020). This phenomenon is also documented in the specific setting of variational inference for BNNs Heek (2018), Yao et al. (2019), Krishnan and Tickoo (2020), Foong et al. (2020). For instance, Foong et al. (2020) demonstrated in experiments that optimizing the ELBO may not lead to accurate predictive means or variances.

Another area of active research in variational BNNs is the design of the variational family itself. For the large part, the mean-field family of fully factorized Gaussian distributions is still predominant in the general practice of variational inference (Graves, 2011; Blundell et al., 2015; Hernandez-Lobato et al., 2016; Li and Turner, 2016; Khan et al., 2018; Sun et al., 2019). The mean-field assumption is mostly adopted for computational ease, though the limitations are well known (MacKay, 1992; Coker et al., 2022). Moving beyond mean-field Gaussian, we can find works that make use of more realistic covariance structures (Louizos and Welling, 2016; Zhang et al., 2018) or more expressive approximating families, e.g., via normalizing flows (Louizos and Welling, 2017; Papamakarios et al., 2021).

Finally, we note there have been a few recent works that recognize the non-identifiability of deep learning models Moore (2016), Pourzanjani et al. (2017), Kurle et al. (2022). These works however seem to treat the non-identifiability as an issue to be fixed.
## 6 Methodology

To achieve a good variational approximation, conventional wisdom says to make \(\mathcal{Q}\) as "expressive" as possible. We will approach the design of the variational family in a more principled manner using SLT. To this end, we rely on recent work in Bhattacharya et al. (2020) which leveraged SLT to produce an _idealized_ variational family as follows. Let \(\mathcal{Q}_{0}\) be a family consisting of generalized gamma distributions in \(\mathbb{R}^{d}\):

\[\mathcal{Q}_{0}=\{q_{0}(\xi|\boldsymbol{\lambda},\boldsymbol{k},\boldsymbol{ \beta})=\prod_{j=1}^{d}q_{0}^{j}(\xi_{j}|\lambda_{j},k_{j},\beta_{j})\} \tag{16}\]

where

\[q_{0}^{j}(\xi_{j}|\lambda_{j},k_{j},\beta_{j})\propto\xi_{j}^{2k_{j}\lambda_{j }-1}\exp(-\beta_{j}\xi_{j}^{2k_{j}})1_{[0,1]}(\xi_{j})\]

for \(\boldsymbol{\lambda}\in\mathbb{R}_{>0}^{d},\boldsymbol{k}\in\mathbb{R}_{>0}^{ d},\boldsymbol{\beta}\in(0,\infty)^{d}\). **Henceforth, let \(g:=g_{\alpha}\) where \(\alpha\) is such that \(M_{\alpha}\) is an essential chart.** In other words, we are fixing a resolution map \(g\), working in a fixed essential chart domain, and a coordinate \(\xi\) on that domain that makes \(K(g(\xi))\) a monomial as a function from \(\mathbb{R}^{d}\to\mathbb{R}^{d}\). The idealized variational family of Bhattacharya et al. (2020) is given as the pushforward of base distributions \(q_{0}\in\mathcal{Q}_{0}\) by said map \(g\):

\[\mathcal{Q}=\{g\sharp q_{0}:q_{0}\in\mathcal{Q}_{0}\}. \tag{17}\]

We refer to this as an _idealized_ variational family for the simple fact that the resolution map \(g\), though its existence is guaranteed, is almost never tractable except in the simplest model-truth-prior triplets. Also note that although the family \(\mathcal{Q}_{0}\) is mean-field, (17) is _not_.

To study the variational approximation gap incurred by the idealized family (17), we will first introduce some definitions to help us rewrite the gap \(\mathcal{G}\) in notation that is consistent with Bhattacharya et al. (2020). Define

\[\Psi_{n}(q_{0})=-\mathbb{E}_{q_{0}}nK_{n}(g(\xi))-\mathrm{KL}(q_{0}(\xi)\parallel \varphi(g(\xi))|g^{\prime}(\xi)|) \tag{18}\]

See Appendix C for the derivation that the variational approximation gap in (14) is equivalent to

\[\mathcal{G}=\log\bar{Z}(n)-\sup_{q_{0}\in\mathcal{Q}_{0}}\Psi_{n}(q_{0}). \tag{19}\]

Following Bhattacharya et al. (2020), we consider the deterministic approximation gap corresponding to (19). This is accomplished by replacing \(K_{n}\) with \(K\), leading to

\[\Psi(q_{0}):=-\mathbb{E}_{q_{0}}nK(g(\xi))-\mathrm{KL}(q_{0}(\xi)\parallel \varphi(g(\xi))|g^{\prime}(\xi)|). \tag{20}\]

Figure 2: We show in these schematics that evaluating variational approximations to BNNs according to their induced predictive distribution is fraught with potential misinterpretations.
and

\[\bar{Z}_{K}(n):=\int_{W}e^{-nK(w)}\varphi(w)\,dw.\]

For our theoretical investigation, we shall concern ourselves with the _deterministic_ variational approximation gap,

\[\mathcal{G}_{K}:=\log\bar{Z}_{K}(n)-\sup_{q_{0}\in\mathcal{Q}_{0}}\Psi(q_{0}). \tag{21}\]

Techniques for generalizing the main result Theorem 6.1 which concerns \(\mathcal{G}_{K}\) to the stochastic world can be found in Plummer (2021, Section 5.3.3).

We will appeal to large-\(n\) asymptotics to study the behavior of (21). Note that the study and deployment of BNNs is no stranger to large-\(n\) asymptotics, both in early MacKay (1992) and recent Ritter et al. (2018) works. We proceed under this tradition, but deviate from the crude (and incorrect) Laplace approximation that is often employed and instead use the correct asymptotics provided by SLT.

### Model evidence in singular models

To study the gap in (21), we begin by examining the asymptotic behavior of \(\bar{Z}_{K}(n)\). When the model is regular, we need not bother with SLT and may find to leading order, \(\bar{Z}_{K}(n)=\varphi(w_{0})\sqrt{\frac{(2\pi)^{d}}{\det H(w_{0})}}n^{-d/2}\) via the Laplace approximation. This approximation, however, is egregiously inappropriate for strictly singular models, in particular neural networks Wei et al. (2022). Nonetheless, perhaps due to a sense that no tractable alternatives exist, the Laplace approximation is seeing a resurgence of application in Bayesian deep learning Ritter et al. (2018), Immer et al. (2021).

For strictly singular models, the quantities \(Z(n),\bar{Z}(n)\) and \(\bar{Z}_{K}(n)\) manifest as singular integrals, i.e., integrals of the form \(\int_{W}e^{-nf(w)}\varphi(w)\,dw\) where \(W\subset\mathbb{R}^{d}\) is a compact semi-analytic subset, and \(f\) and \(\varphi\) are real analytic functions. The behavior of a singular integral depends critically on the zeros of \(f\). According to Theorem 6.7 in Watanabe (2009), we find to leading order:

\[\bar{Z}_{K}(n)=C(p,p_{0},\varphi)n^{-\lambda(p,p_{0},\varphi)}(\log n)^{m(p,p_ {0},\varphi)-1}, \tag{22}\]

where \(C(p,p_{0},\varphi)\) is a constant independent of \(n\) that we shall call the **leading coefficient** following the terminology of Lin (2011). Note that since \(\lambda(p,p_{0},\varphi)=d/2\) and \(m(p,p_{0},\varphi)=1\) in regular models, (22) is a true generalization of the Laplace approximation, holding for both regular and strictly singular models.

### Bounding \(\mathcal{G}_{K}\)

We show in Lemma D.2 in Appendix D, that for large \(n\), the following bound holds

\[\sup_{q_{0}\in\mathcal{Q}_{0}}\Psi(q_{0})\geq-\lambda(p,p_{0},\varphi)\log n+C \tag{23}\]

where \(C\) is the constant free of \(n\) in Lemma D.2. This result is in the same spirit as (Bhattacharya et al., 2020, Theorem 3.1), except that we have improved on the tightness of their lower bound, which in turn allows us to devise better initialization of the variational parameters. With Lemma D.2, we are now in a position to characterize the (deterministic) variational approximation gap, \(\mathcal{G}_{K}\).

**Theorem 6.1** (Deterministic variational approximation gap).: _Suppose the model-truth-prior triplet \((p,p_{0},\varphi)\) is such that Theorem 2.1 holds. Let \(g=g_{\alpha}\) where \(\alpha\) is such that \(M_{\alpha}\) is an essential chart. On this essential chart, write the local RLCTs \(\tilde{\lambda}_{j}=\frac{\tilde{h}_{j}+1}{2\tilde{k}_{j}},j=1,\ldots,d\) in descending order so that \(\tilde{\lambda}_{1}\) is the RLCT of the triplet \((p,p_{0},\varphi)\), i.e., \(\tilde{\lambda}_{1}=\lambda(p,p_{0},\varphi)\). If the multiplicity of the triplet is 1, we have, for \(n\) large, \(\mathcal{G}_{K}\leq\log C(p,p_{0},\varphi)-C+o(1),\) where the constant \(C\) is as given in Lemma D.2._

All that is needed for the proof of Theorem 6.1 is to put together the lower bound in Lemma D.2 with the fact that \(\bar{Z}_{K}(n)\) admits the asymptotic expansion in (22). Even when \(m(p,p_{0},\varphi)\neq 1\), there may be finite \(n\) situations when the two terms \((m(p,p_{0},\varphi)-1)\log\log n\) and \(\log C(p,p_{0},\varphi)-C\) are comparable. In such settings, the idealized variational family \(\mathcal{Q}\) in (17) could still perform well.

### Learning to desingularize

In the preceding section, we studied the deterministic variational approximation gap of an idealized variational family. Although Hironaka proved the existence of a resolution map and showed that it can be found by recursive blow up, known algorithms for finding such resolutions, other than a few exceptional cases (such as those for toric resolutions),
have complexity that vastly exceed existing computational capabilities. Thus we are precluded from directly applying the idealized variational family.

This leads us to consider _learning_ the resolution map \(g\) using an invertible architecture \(G_{\theta}\) resulting in the variational family

\[\hat{\mathcal{Q}}=\{G_{\theta}\sharp q_{0}(\boldsymbol{\lambda},\boldsymbol{k}, \boldsymbol{\beta}):\boldsymbol{\beta}=(n,\beta_{2},\ldots,\beta_{d})\}. \tag{24}\]

If the network is expressive enough, we can hope that \(g\in\{G_{\theta}:\theta\}\), which would lead \(\hat{\mathcal{Q}}\) to enjoy the theoretical guarantee provided in Theorem 6.1. Note in (24) the first coordinate of \(\boldsymbol{\beta}\) has been set to the sample size \(n\). The proof of Lemma D.2 reveals why we do so. Specifically, it is shown that the following parameters in \(q_{0}\) can achieve \(\Psi(q_{0})=-\lambda(p,p_{0},\varphi)\log n+C\):

\[\lambda_{1}=\lambda(p,p_{0},\varphi),\quad k_{1}=\tilde{k}_{1},\quad\beta_{1}=n\]

where \(\tilde{k}_{1}\) is as in Theorem 6.1. Note that \(\lambda(p,p_{0},\varphi)\) and \(\tilde{k}_{1}\) are unknown, but \(n\) is certainly known.

It might be readily apparent at this point that we have in \(\hat{\mathcal{Q}}\) a standard normalizing flow, albeit with the base distribution given by the generalized gamma distribution. To ease the computational cost, we fix the variational parameters \(\boldsymbol{\lambda},\boldsymbol{k},\boldsymbol{\beta}_{|-1|}\) and absorb the learning of their optimal values into the invertible transformation \(G_{\theta}\). Note that this is in line with standard practice, whereby normalizing flows adopt parameter-less base distributions.

To summarize, recognizing that the variational approximation gap can be theoretically studied using SLT allowed for the design of a principled variational family which incurs a variational approximation gap that is independent of sample size \(n\), to leading order. To the best of our knowledge, no existing works on normalizing flows for BNNs theoretically address the variational approximation gap. Furthermore, our results offer a new perspective on the benefits of using normalizing flows for variational inference in BNNs.

## 7 Experiments

In the following set of experiments4, we will isolate and examine the effect of the base distribution. Specifically, we compare the _generalized gamma base distribution_ to the commonly-adopted _Gaussian base distribution_, holding the architecture of \(G_{\theta}\) fixed when we do so. At the outset, we expect that when \(G_{\theta}\) is expressive enough, the effect of the base distribution will be small. However, when \(G_{\theta}\) is more limited (and thus less computationally expensive), we conjecture the generalized gamma base distribution can "pick up the slack" and outperform the Gaussian base distribution.

Footnote 4: The code to reproduce our results is available at [https://github.com/suswei/BNN_via_SLT](https://github.com/suswei/BNN_via_SLT).

In line with our earlier discussion, the parameters of the base distributions are frozen throughout training, see Appendix E for the initialization used. The invertible network \(G_{\theta}\) is implemented as a sequence of affine coupling transformations. We denote by base_numcouplingpairs_numhidden the variational family that results from pushing forward the base distribution through \(G_{\theta}\) with the said configuration, see Appendix E for a complete description of the implementation. We consider a total of four different expressivity levels of \(G_{\theta}\) from least to most: \(2\_4\), \(2\_1\)6, \(4\_4\), \(4\_16\).

The expression for the ELBO objective corresponding to each of the base distributions is given in (29) and (30) in Appendix E. Details of the training procedure such as epochs, learning rate, and optimizer are also given there. Let \(\hat{q}^{*}\)

\begin{table}
\begin{tabular}{l c c c c c} \hline \hline model & \(H\) & \(\dim_{p_{0}}\) & \(\lambda(p_{0},p_{0},\varphi)\) & \(\dim_{p_{0}}\) & \(\dim_{p_{0}}\) \\ \hline fifth & 3 & 42 & - & 13 & 1 \\  & 7 & 98 & - & 13 & 1 \\  & 16 & 22 & - & 13 & 1 \\  & 40 & 560 & 13 & 1 \\  & 2 & 144 & 5.0 & 1 & 2 \\  & 7 & 119 & 35.0 & 10 & 7 \\  & 10 & 202 & 65.0 & 13 & 10 \\  & 16 & 560 & 152.0 & 19 & 16 \\ tanh & 15 & 30 & - & 1 & 1 \\  & 50 & 100 & - & 1 & 1 \\  & 115 & 20 & - & 1 & 1 \\  & 200 & 500 & - & 1 & 1 \\ tanh (zero mean) & 15 & 30 & 1.93 & 1 & 1 \\  & 50 & 100 & 3.53 & 1 & 1 \\  & 115 & 220 & 5.36 & 1 & 1 \\  & 280 & 560 & 8.36 & 1 & 1 \\ \hline \hline \end{tabular}
\end{table}
Table 1: The various model-truth-prior triplets considered in experiments. The truth is realizable. The prior over network weights is standard Gaussian. The RLCT is only known in some of the cases.
be the variational distribution obtained at the end of training. Comparison of the base distributions, and hence the two different normalizing flows, will be made according to normalized MVFE, \(\bar{F}_{vb}^{*}(n)\), and VGE, \(G_{n}(p_{vb}(y|x,\mathcal{D}_{n})\). (For both, the lower, the better.) We will also estimate the coefficients \(\lambda_{\text{ref}}\) in (13) and \(\lambda_{\text{spe}}\) in (15), see Appendix E.

We consider four model-truth-prior triplets, summarized in Table 1, in which the truth is always realizable. In all four triplets, the prior over the neural network weights is chosen to be the standard Gaussian following conventional practice in BNNs Neal (1996), Bishop (2006). Note, priors for BNNs are notoriously difficult to design and is an area under active research Sun et al. (2019), Nalisnick et al. (2021).

### Results

Due to space constraints, we only show a subset of the results in Figure 3; complete results can be found in Appendix E. In the first column of Figure 3, we plot \(\log n\) versus the normalized MVFE. First, we observe that when \(G_{\theta}\) is not very expressive, the generalized gamma resoundingly outperforms the Gaussian base distribution for the reduced rank and ReLU experiments across all values of \(H\) in terms of achieving lower MVFE. (This can be better seen in Figure 8 in Appendix E.) On the other hand, as conjectured, when \(G_{\theta}\) is most expressive at the 4_16 configuration, the distinction in MVFE between the base distributions is still discernible but less dramatic, see Figure 10. Interestingly, for the \(\tanh\) triplet, the Gaussian base distribution sometimes achieves lower MVFE depending on the configuration of \(G_{\theta}\).

In the second column of Figure 3, we plot \(1/n\) versus the VGE. The results empirically verify the issues we highlighted in Section 4. In terms of VGE, the generalized gamma is not uniformly better than the Gaussian base distribution for the ReLU experiment, contrary to what the corresponding MVFE plots suggest. Only for the reduced rank experiment do we see one-to-one correspondence between MVFE and VGE. Note that the VGE fit is particularly poor for the Gaussian 2_4 and 2_16 configurations because these variational approximations are themselves poor. Next, note the scenario in Figure 2b is borne out by some of the \(\tanh\) experiments. Take for instance \(\tanh\) at \(H=115\) for the 2_4 configuration. Judging by MVFE alone the generalized gamma base is worse than Gaussian base, but the corresponding VGE curves show the opposite, see (3,3) subplot in Figures 8 and 9.

Figure 3: MVFE versus \(\log n\) is displayed in the first column and VGE versus \(1/n\) is displayed in the second. Each row corresponds to a different model-truth-prior triplet. Line color indicates the expressiveness of the network \(G_{\theta}\), darker being more expressive. Error bars represent mean, min and max over 30 draws of the training set \(\mathcal{D}_{n}\). The dashed line is the least squares fit with \(\lambda_{\text{ref}}\) and \(\lambda_{\text{spe}}\) coefficients and their \(R^{2}\) values displayed in legend.
## 8 Discussion

We conclude by discussing some of the limitations of the current work. On the empirical front, the reader may have noticed that our experiments did not involve truly deep BNNs. Strictly speaking this is not a limitation of the proposed method but rather a limitation of the scalability of normalizing flows for approximating deep BNNs. We expect the proposed methodology to benefit from orthogonal research advances in normalizing flow architectures.

On the theoretical side, it may be of interest to flush out the magnitude of \(\log C(p,p_{0},\varphi)-C\) in Theorem 6.1. The general expression for \(C(p,p_{0},\varphi)\), although known in special cases [11, Corollary 5.9], has complex dependency on \(K(w)\) and the prior. However, we do expect that the leading coefficient can be bounded with some effort. Relatedly, it is important to recognize that Theorem 6.1 only concerns the variational approximation gap of the idealized family in (17). Deriving an analogous result for the Gaussian base distribution would make for interesting future work.

We are optimistic that natural conditions on the model-truth-prior triplet and the variational family should allow for general statements about MVFE asymptotic expansions. Further efforts into studying the asymptotics of the MVFE will also advance knowledge of the relationship between \(\lambda_{\text{vfe}}\) and \(\lambda_{\text{vge}}\). In its place, our results here show that it is all the more important to pay attention to the variational approximation gap if we wish to have useful downstream predictions.

## Acknowledgements

We thank Daniel Murfet for helpful discussions. SW was supported by the ARC Discovery Early Career Researcher Award (DE200101253). This material is also based on work that is partially funded by an unrestricted gift from Google.