On the detrimental effect of invariances in the likelihood for variational inference
Authors: Richard Kurle
Ralf Herbrich
Tim Januschowski
Bernie Wang
Jan Gasthaus
Venue: NeurIPS 2022 Accept (2022)
Keywords: Bayesian Neural Networks
Variational Bayes
Variational Inference
Invariance
Symmetry
posterior collapse

Abstract: Variational Bayesian posterior inference often requires simplifying approximations such as mean-field parametrisation to ensure tractability. However, prior work has associated the variational mean-field approximation for Bayesian neural networks with underfitting in the case of small datasets or large model sizes. In this work, we show that invariances in the likelihood function of over-parametrised models contribute to this phenomenon because these invariances complicate the structure of the posterior by introducing discrete and/or continuous modes which cannot be well approximated by Gaussian mean-field distributions. In particular, we show that the mean-field approximation has an additional gap in the evidence lower bound compared to a purpose-built posterior that takes into account the known invariances. Importantly, this invariance gap is not constant; it vanishes as the approximation reverts to the prior. We proceed by first considering translation invariances in a linear model with a single data point in detail. We show that, while the true posterior can be constructed from a mean-field parametrisation, this is achieved only if the objective function takes into account the invariance gap. Then, we transfer our analysis of the linear model to neural networks. Our analysis provides a framework for future work to explore solutions to the invariance problem.


On the detrimental effect of invariances in the likelihood for variational inference
Richard Kurle AWS AI Labs
Ralf Herbrich Plastoe Institut ralf. www.richhhpi.ccoo
Tim Januschowski f Zalando SE
kurler@amazon.com com
tim. j januschowski Gzal
Yuyang Wang
Jan Gasthaus
AWS AI Labs
AWS AI Labs
yuyavang@amazon.com com
gasthaus@amazoon.oo
Abstract
Variational Bayesian posterior inference often requires simplifying approximations such as mean-field parametrisation ensure tractability. However, prior work has associated the variational mean-field approximation for Bayesian neural networks with underfitting in the case of small datasets large model sizes. In this work, we show that invariances in the likelihood function of over- over-paranmee parametrised models contribute to this phenomenon because these invariances complicate the structure of the posterior by introducing discrete and/or continuous modes which cannot be well approximated by Gaussian mean-field distributions. In partieular, we show that the mean-field approximation has additional gap in the evidence lower bound compared to a purpose-buill posterior that takes into account the known invariances. Importantly, this invariance gap is not constant; it vanishes the approximation reverts to the prior. We proceed by first considering translation invariances linear model with single data point in detail. We show that, while the true posterior can be constructed from a mean-field parametrisation, this is achieved only if the objective function takes into account the invariance gap. Then, transfer our analysis of the linear model to neural networks. Our analysis provides framework for future work to explore solutions to the invariance problem.
Introduction
Bayesian neural networks (BNNs) have several appealing advantages compared to deterministic neural networks (NN) such as improving generalization (38), capturing epistemic uncertainty [[7], Unfortunately. reaping these theoretical benefits has far been impeded (37). In partieular, several practical issues with variational Bayesian inference
and providing framework for continual learning methods 18 23
are the de-facto standard technique for scaling inference in BNNs large datasets [15 2 have been identified be partially responsible for a performance gap compared to deterministic NNs [6]. The most common variational approximation of the posterior is product of independent Normal distributions, commonly referred to the mean-field approximation. It has been observed that mean-field variational BNNs (VBNNs) suffer from severe underfitting for large models and small dataset sizes [13]. Recent work [5] showed that under certain assumptions such as an odd Lipschitz activation function and a finite dataset with bounded likelihood, the optimal mean-field approximation collapses the prior as the NN width increases, ignoring the data.
"Correspondence kurler warlerramazon.co Work done while at AWS Al Labs.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

The goal of this work is to shed light on why the mean-field variational approximation collapses and provide new angle for future works to address this shortcoming, To this end, we study invariances in the likelihood function of overparametrised models and their detrimental effect on variational inference. For instance, is well known that NNs exhibit several invariances with respect to their parameters, including node permutation invariance [1), sign-flip invariance (in the case of odd activation functions) [29 31, scaling invariance (in the case piece-wise linear activations) (26), and as we will show in Sec. 5 the parameter space of BNNs additionally has subspaces with translation invariance.
Invariance in the likelihood function does not necessarily pose challenge if maximum likelihood point estimation through stochastic gradient descent is used, because convergence to any of the equivalent optima (resulting from the invariance) suffices for prediction. However, these invariances are detrimental to the variational Bayesian approach. As first step to show this, we isolate the impact of the invariances in Sec. 3 To this end, we construct both a mean-field as well as invariance- abiding approximation which explicitly models the invariances by integrating over all transformations that leave the likelihood invariant. Notably, both aforementioned posterior approximations are constructed from the same (mean-field field) likelihood approximation. We then prove that, under the conditions outlined in Sec. 3 the mean-field approximation induces the same posterior predictive distribution as the invariance-abiding approximation. However, we also prove that the ELBO objective corresponding to these two approximations differs by the KL divergence between both posterior approximations; we refer to this difference as invariance gap. Importantly, the gap vanishes if both distributions are identical, which is the case if the mecan-field approximation reverts to the prior.
We then demonstrate the detrimental effect of invariances in the likelihood function of an overparametrised Bayesian linear regression model in Sec. This model is purposely selected as the canonical model exhibiting (only) translation invariance. We provide a detailed analysis of this model, including tractable solution for the invariance-abiding approximation. It turns out that the optimal parameters w.r.t. the ELBO objective with the invariance-abiding distribution result in the true posterior. In contrast, posterior approximations with parameters that are optimal w.r.t. the mean-field ELBO revert to the prior as the number of dimensions increases.
Finally, we transfer our analysis of the linear model to VBNNs in Sec. 5] showing that subspaces in BNNs exist that exhibit translation invariance. Combined with a previous analysis of the node permutation invariance (cf. Kurle et al. [99 or Ap the supplementary material), this provides the basis for future work to approximate the (generally intractable) invariance gap and thereby optimise for tighter and favorable ELBO objective. We start in Sec. by introducing the basic concepts of VBNNs and recent results on which we build our contribution.
2 Background
2.1 Variational Bayesian Neural Networks
Neural network functional model. Deep NNs are layered models that progressively transform their inputs in each layer. More formally, an L-layer NN computes algebraically
f(x) = hr. (wI (wI. 21,4 h1 (wI,x)
zr,1 = hu (wT,2=1)>
where the by R + R are monotonic transfer functions that introduce non-linearities. Such a network has n1 + n22 +... + nı,, many nodes 21,i. where indexes the layer-wise vectors z,. These nodes are each a weighted linear combination of either the input vector x (first hidden layer) or the value the hidden units z, (all other hidden layers and the output (x)). We denote all learnable parameters of the model by the stacked weight vectors of each layer and node, w =
wT..n.!".
Variational Bayesian treatment. If the weights and biases w are treated as random variables with a prior distribution pww). then the posterior p(w (DD) (nn-idd by the dataset D through the likelihood function f(w; := p(D w) that is defined via f(x)=is is referred to as Bayesian neural network (BNN). The variational Bayesian method approximates the posterior by distribution go (w) x p(w [ D) with variational parameters 0, casting inference as optimization problem
To w) argmin KT. [40 11 1P((11))]
go E Q
2

---

where KL [ge || p(- [) Bw-go Im go (w) In l [D] and Q is a family of distributions over w.
The optimization of [ achieved by maximising lower bound the (log) model evidence ((D).
LELBO (go, D) p(D) KL [ge p| [)] Ew [nn f(w; D) p(w)
4o (w)
(2)
ELL (g0. D) KL 100 l1p],
where E.L. (go, := [In <(w; D) is the expected log-likelihood.
Definition (Mean-field variational BNN), A mean-field variational BNN (VBNN) is the minimizer of ( where is the family of Gaussians with diagonal covariance matrix.
2.2 Data-related bound on the KL divergence
The term KL [go || p|
(2) admits data-dependent upper bound that naturally occurs as a consequence of the finite information provided by a finite dataset in the presence of noise. This result has been shown in (5) and we recall the result for a Gaussian likelihood with homogeneous noise (for other likelihood functions, we refer to App. F in (5)).
Assume a VBNN with Gaussian observation noise defined through y = [x) +G, with e NO ag)) an isotropic Gaussian prior p(w) = (w; 0, I). and a mean-field variational approximation go go(w) (w) N N(w; (w; m, Diag(v)). Define the NN output variance of (x(*)) Vwap [f(x("); w)] for some fixed input x()), where L indicates the last layer. Then,
N o (x(1)) + (y(r))2
KL [g || & ELL (q": D) ELL (p. D)
20
where q² is hypothetical approximation that predicts the data perfectly up to the noise variance of (see App. [ for details). Prior work (5) has used the data-related bound prove that the predictive distribution of mean-field BNNs converges to the prior predictive distribution if the network width large and the activation function odd, ultimately resulting in posterior collapse the prior. In contrast, we address the question why the mean-field approximation cannot approximate the posterior by relating this data-related bound to invariances in neural network functions the following section.
3 Invariance-abiding variational approximation
We develop a framework that enables modelling the invariances in the likelihood and understanding their impact on the ELBO objective. To achieve this, we approximate the likelihood by a Gaussian function with variational parameters and marginalise over all transformations to which the true likelihood is invariant. By taking the product with prior, we construct an invariance-abiding posterior approximation fmix which can be related to mean-field approximation Zo that does not model invariances. We then describe conditions under which both approximations yield an identical posterior predictive, while the KL regularisation term of go is lower bounded by the KL of qmix-
Variational likelihood and posterior approximations. Non-identifiability implies that the posterior does not concentrate on a single set of parameters irrespective of the dataset size, because the same likelihood is assigned to different parameter values [12]. We model this invariance through transformations t(, r) to which the likelihood f(w; ; D) = p(D) w) is invariant via variables
Vr ~ prr)(r): E(t(w, r); D)) = f(w; D).
4)
From 4 it follows that the likelihood is also invariant w.r.t. the marginalisation
é (w; D) = Er-p [t (t(,rr):)))
We assume that each of the equivalent parametrisations w' +(w, r) the likelihood has the same probability a priori. For discontinuous transformations such as node permutations (see App. E p(r) a uniform distribution over discrete variables indexing these transformations. For continuous translation invariance (see Sec. 4), we model the uniform distribution variance. is conceivable that the likelihood can constructed by marginalising over these transformations of simpler function lo(w: D) such that l(w; D) Er-p (t((,, D)). For
Gaussian with infinite
3

---

instance, permutation invariance induces factorial number of discontinuous modes that are each equivalent (cf. [19 App. E and as we show in Sec. 4 the full covariance Gaussian likelihood and posterior of an over-parametrised Bayesian linear regression model can be constructed from product of independent Gaussians. Curiously, it may sufficient approximate lo while taking into account the known invariances. To this end, we define Gaussian variational likelihood approximation go 0) xe fo(w, D). with variational parameters O. From this single mode approximation, we model an invariance-aiidin likelihood approximation using the same parametrisation through
Imix (w; 0) := Er-p [go (f( w, r); 0)] Re Er-p [oo((w, r), D)] = ((w;
(5)
We consider mean-field Gaussians for the prior p( (w) and likelihood approximation go (w; f), where (m, }} are means and variances of 9o- We then define mean-field posterior as the product
qo (: 0) == Z p p(w): go w: 0),
Zo = p(w) gu(w; 0)dw.
(6a)
Similarly, we define an invariance-abiding posterior as the product prior and invariant likelihood:
amix (w: 0) Zmix p(w). Imix (w: 0), Zmix = =f p ((W): 9mix (w: 0)dw,
(6b)
where Zo and Zmix are normalisation constants. While go (w; 0) is a mean-field approximation, the invariance-abiding approximation gmix (w; 0) is mixture with infinite continuous or finite discrete modes depending on the type of invariance. We will describe continuous translation invariance in Sec. for the discrete node permutation invariance, see App. of the supplementary material.
Posterior predictive equivalence. Next, we discuss conditions under which we can construct approximations gmix 9mix(w; (w; 0) and go(w; 0) that yield the same predictive distribution. We first write the density fmix as an expectation of product densities,
gmix (: 0) Zii p(w): / p(r) go (o r); 0)dr / p(r) p(w) go((((,,))) dr.
Zmix
Then, we assume that there exists a mapping r' = (()) such that
Vr p(r) p(w)- go(t(w, r): 0) = p(t(w, ((r)) (o ((x)) 0).
The condition in 7 allows to exploit the invariance property of the likelihood (approximation) also for each product density qo (f((w, ((r)); 0) x p(t(w, ((())
((((ww(())) 0), because then
(r( p(r) qo (w, (()): 0)dr,
Gmix (w 0) = Zmix
where the normalisation constants are Z, f p(w)- go (tW,, (p())) 0)dw, and Zmix is defined in (6b) We show that the condition in holds for translation and permutation invariance in App.
The second condition is that all invariance transformations t(., r) must be volume- teeeeeprreeeigg...m i.e.
Öt(w, r)
Vr p(r): det
Ow
Although (7) and (8) are fairly restricting, common invariances such translation and permutation invariance fulfil these conditions for Gaussian priors and likelihood approximations (see App. C With 7 and (8,, we show the posterior predictive equivalence (see Lemma in App- C
E Ew-qqaxx [In w)] = Ew--q (In p(D w)]
Invariance gap. If the conditions in 7 and 8 met, also show (see Lemma 2 in App. C)
LELBO (go, D) LELBO (quix, D) = KL 4oo p] KL [qmix || p] = KL [go [mix].
(10)
Interestingly, the gap between the two respective ELBO objectives is given exactly by the relative entropy between the mean-field and invariane--abdiin approximation; we therefore refer to it as the invariance gap. We make the following observation about the detrimental effect of the invariances: when maximising the standard ELBO LELBO (go. D) instead of the tighter objective LELBO (Amix. D) w.r.t. the parameters of the variational likelihood approximation gu (: 0) (used construct both qu and (mixx), we see that the former objective favors solutions where qu and quix coincide. This suboptimal solution is obtained if the mean-field posterior qo 4o(w) and thus also the invariance-abiding posterior ymix (w) revert to the prior. This is the case if 5o(w) is uniform, because then gmix (w) is uniform as well, and because both posteriors are constructed as the product of prior and likelihood approximation (cf. (a)))

---

Implication of data-related bound. Since KL Igo || is upper bounded by the best- and worst-case ELL described in Sec. 2.2 the invariance gap is also bounded (using (T) and (3):
KL 900 [mix] I KL (qo p] S ELL (q", D) ELL (p. D).
(11)
Consequently, this bound puts a constraint on the achievable solutions to the maximisation of the ELBO objective w.r.t. the variational parameters of the mean-field approximation go(w; 0). As discussed above, one specific parametrisation for which the invariance gap vanishes is the mean- field posterior collapse, qo (w: 0) x p(w). However, to show that the invariance gap not only admits posterior collapse potential parametrisation but indeed incentivises it, we next tackle the question how the invariance gap behaves for non-collapsed approximations. One particularly relevant variational parametrisation is the optimal solution w.r.t. the ELBO objective with the invariance- abiding approximation. In order to tackle this question, we now consider a simple model for which the relevant distributions and the invariance gap can be computed exactly.
4 Translation invariance in linear models
We study an over-parametrised Bayesian linear regression model as the canonical model that exhibits translation invariance. The model serves a useful tool to understand the detrimental effect of translation invariance since all interesting quantities can be computed analytically, incl. the mean-field
and evrariance-abiding distribution defined in (6) and the invariance gap from (10). We show how to lift results from this canonical model to the more general case of NNs in Sec.
Likelihood model. Consider a linear model with K latent variables w [W],,... wx] assume that we have N observations of variables y given dependent inputs x. To simplify the setting, we will assume that x = 1 (see App. D for the more general case). We further assume that the observation depends only on the inner product with the inputs and additive Gaussian noise. That is, K 11' + ~ N (0, Note that any change in w which leaves the sum of the elements unaffected does not change the likelihood. In general, we can model this translation invariance using K dimensional vector A E IR K-1 and observing that
I and
= -1 [-=[-1-]
I
1 ww= = 1 (w+
B B:
(12)
since 1 IBA This over-parametrised model exhibits translation invariance I(w, A) = w BB.
Prior. We assume a Gaussian prior p (w) = N (w; u, with diagonal covariance Σ Diag (a²). where we consider a² K o 1 proportional to the number of dimensions K, such that the predictive variance V to model prior over parameters that induces the same prior over functions for different K.
R Σn wy + e] is constant w.r.t. K. Another way to view this is
Posterior.
The posterior of this Gaussian linear model can be computed as (see App. D.3
N 11 111 K20
m V E. Di
p(w y) N' (w; m VØ,) V
(13)
Ko²
As can be seen, the resulting posterior has full covariance with diagonal plus rank-1 matrix. We will now show that we can construct this posterior from the prior and mean-field Gaussian approximation of the likelihood by marginalising over all translations A ~ p(A) ) to which the likelihood is invariant.
4.1 Mcan-field parametrisation of the likelihood function
Next, we construct the mean-field and invariance-abiding posterior approximations defined in (6) from the prior defined above and the mean-field likelihood approximation with locations m and variances X. We model that the likelihood is translation invariant by computing the marginal from (5)
and considering p(A) = =N (A; O, 3"1) with - Do as the uniform distribution over all translations:
Jmix (w; 0) = f g0 (1( A); 0) (A) dA
lim / N (w;m + BA, Diag (A)) -N) NN (A ds
N (W; Miix; V Vmix).
8-000
5

---

Writing the uniform distribution the limiting case of Gaussian with infinite variance allows us to compute the integral analytically. As shown App. D the resulting function a multivariate Gaussian with degenerate rank-1 covariance matrix and the same location parameter as go go(w):
mix m,
(14)
1
However, the invariance-abiding posterior is non-degenerate, and it can be computed efficiently
1 (m an-H) 12) Diag (o*) ('
Umix (w; 0) -N w; /
o² (0²)
(15)
I 12
As can be seen, this covariance matrix is not diagonal; it has an additive rank-1 term that vanishes as - cC. Interestingly, the location parameter of Qmix is translated from the prior location μ along the direction of the prior variance vector σ². This is because Imix is uniform along the K - dimensions hyper-plane determined by its normal vector 1. Taking the Gaussian product then translates the o². Note also that the invariance-abiding posterior can be
location in the direction Diag(or)1 written in the same form as the true posterior in (33) (see App. D.3). The optimal mnvariance-aiiiin posterior is thus the true posterior. We visualise the two respective posterior approximations and the
corresponding likelihood in Fig. of App. with two different parametrisations (ef. Sec. 4.3)
4.2 Invariance gap
To quantify the detrimental effect of the translation invariance in the considered linear model we now compute the invariance gap in (TI). Note again that assume scaled standard normal prior with variance o² = Koō 1. We simplify the form the KL divergence by assuming that all variances and means take the same value in the likelihood (and thus also in the posterior) approximation, i.c.
1k: A X, my m, ax à, Ma = j..
This choice motivated by the fact that the posterior approximation is a function of the sum 1 only (cf.. (15)). and, hence, it makes no difference for 9mix. However, since the prior variance is proportional 1, the Gaussian likelihood resulting in the highest ELBO for the approximation go also has a variance vector proportional to 1. Using this assumption, the invariance gap is
K
KT. [go
[imix]
In
(16)
2
Ø2 +
This is a convex function in 82+8 with a minimum however not evident whether the gap has large magnitude compared to the rest of the regularisation term and over-regulariess in practice. In the next section, we therefore analyse this term at the optima for the mean-field and the evarrance-abiding approximations defined in (17). The corresponding invariance gaps are visualised in Fig. where it can be seen that the invariance gap grows linearly when using the optimal parameters of the invariance-abiding distribution and the unattainable data- related bound is reached quickly, thus preventing this optimum if (17b) is optimised instead of (17a).
1; it is minimised as Ā + co. is
4.3 Optimal mean-field and invariance-abiding parametrisations
To better understand the detrimental effect of the invariance gap, now analyse the ELL and KL terms of the ELBO objective for the mean-field and invariance-abidin variational approximations, respectively. We compare these terms for both distributions with the parameters optimized for both posterior distributions, giving 2 x combinations. We denote the optimal parameters w.r.t. the invariance-abiding approximation and the mean-field approximation, respectively, as
0* mix argmax LELBO (quix (+; DD)
(17a)
0
0 argmax LELRO (go (+; D).
(17b)
Perhaps the simplest way compute these optimal parameters is to notice that ymix in (15) can be written the same form as the true posterior and then simply read out the optimal parameters. This is shown in App. D.3 the resulting optimal variational parameters are
Go (w: 0 ) = -V (w; m Diag (Xiixx)), m
1 N
Σ y()) n=
1.
X
Ko²
1. (18)
N

---

175
1759
1500
1130
invatiance siddng param.
opt. Itarr anre-abiting param.
mean-lied parmm.
1003
apt meat-deil param.
delaveebeed bound
Saow mixted bound
63
505
dimensions
dimenslons
(a) N 10 observations.
(b) N 100 observations.
: O fmix) 9mix 0 ) and KL (-0(:: 0%) (:; 0.]l Prior variances are 0² K 1, where K are the dimensions; the noise variance is oz (2re) all observations y = 1 and x = 1 are identical. As K increases, the invariance gap vanishes in case of the optimal parameters 0r. In contrast, the gap for O Ouix.. which induces the true posterior predictive, grows linearly. As the data-related bound (cf. Sec. 2.2 can not be exceeded by the optimal parameters, 0g can not coincide with the optimal 0mix-
Figure 1: Invariance gap evaluated at different optima (cf. (1)). i.c. KL [qu(::
The optimal parameters for the mean-fiche posterior can also computed analytically, since the Gaussian mean-field distribution that minimises the KL divergence to the multivariate Gaussian true posterior is known 34 The resulting optimal parameters of the mean-field likelihood are
K20y
go (w 0,) = -N (w; m.: Diag (AG)), m m
X
1.
(19)
N
As can be seen, the two respective optimal parameters differ by the factor K the likelihood variance. We then construct the two respective posterior approximations qu and qmix with these two optimal parameters and compare the 2 x 2 combinations in Fig. 2 where we visualise the ELBO terms.
Note again that we model prior variances o² = Kaž 1 such that prior and posterior over functions are identical for any K. Due to this choice and since Gmix (w; mixx approximates the true posterior exactly, it does not suffer from over-paranin parametrisatiionn. This can be seen by the loss terms in Fig. 2 being constant in K. contrast, since Aj depends quadratically on K and の² only linearly, 9mix ; 0
collapses to the prior K - This can be seen e.g. by the shrinking KL regularisation (Fig. 2a
and ELL (in 2b) term, As consequence, and in line with Coker et al. [5], the posterior predictive variance of the optimal mean-field approximation reverts to the prior predictive variance (Fig. 2d).
We have shown that-in the simplified setting linear model with dependent observations- the consequence of not handling translation invariance is indeed posterior collapse as K - co, while modelling the invariance coincides with the true posterior. An important observation and key takeaway is that, while we need to optimise for (17a). the mean-field field posterior approximation go (o: 0 0mix) is sufficient for prediction. Indeed, for the linear model, we could compute the invariance gap exactly and thereby correct the ELBO objective for the invariances of this model. For more complex non- linear models such as BNNs, this section may serve as direction for approximating the invariance gap. To this end, the next section describes the lase translation invariance exhibited by BNNs.
Translation invariance in Bayesian neural networks
Let now go back to the general NN model in Sec. 2.1 In order describe the set of all invariances, note that for each node 21,j in layer I (or the output node y). there a subspace that keeps the value zi.j (or y) invariant when using the ionslaatoo-invarianne model in (2). because the value itself only depends on the actual activation - (or the input x).
More formally, the following equations model all translation invariances of NN at data point x:
ALLL:
f(x) = hr 21 = -u ((WL.1 - BB AL,1) 21-1)
(20)
Vj E (1....,,. m Aij: 21. = hy (wi,z 211) hy ((WI. + Bz.., A1s) 21-1)
(21)
Vj E (1...., m) A Aij: Fi. = -h (wH,,) = ((W1) + BxAis)"xx
(22)

---

tove
80
-25
abiing param
-50
-73
40
-10
Teat opt. near ance-abiing paran.
-25
Sext opt. mean-fisid param.
ept mean-fels parsmn
opt. livo fa asiling
dimensionss
(a) KL [
pl.
(b) Eg (D w))
1.0
-100
sece-sbiding Op opt morrts movttaceeentonn
Rooi opt. Inverlance meneeaeee-aaee stkiing parm.
0.4
Rosi ost, moun param.
-400
spt. mean-fiell param.
Ivvrtance abro
(c) LELBO (g. D) = Egg p w)] KL lp..
(d) V
w +
Figure 2: ELBO loss terms and predictive variance for different posterior approximations. Lines/dots posterior, respectively. Red/blue indicates the optimal parameters (cf. [17)) Prior variance is o² = K-1, and o 1/ (2re); we used N 10 identical observations with value y = 1 and inputs = 1. Notably, both posterior approximations yield the same predictive distribution as can be seen by the ELL (b) and predictive variance (d). The variance of the mean- mean-field approximation reverts to the prior variance K -
indicate the invarianee-abiding/meal
invariace-abidnn/mmenn
where the layer index ranges from 2, ,,....4 , L and BN is given by
I
B
In order to efficiently compute the invariance-abiding likelihood Jmix. note that it also decomposes over the NN layers and the associated parameters as follows: For the nı nodes in the first layer we have
Umix (wi.)) = Bx ( (
x" (m 14)
xT E)x (Ex) (Ex)")
(Ex),E
(23)
(V Σ)x
P(a1,) = Bx [Bwwggmedww
AA1 [x]]]]
(24)
where the outer expectation is taken over pDD) and pD)) is the empirical distribution over the training set DD and where m. μ and Σ are constrained to the parts of the overall weight vector that correspond to W1.j- Now, since the ewariance-abiding mechanism for nodes in layer only depends on the value - 1 of the hidden units in layer n 1 (see (20) and (1 ), we have
(m 11) ZJ (V
z (V E) EZ) (Em)))
Tunix (wi.)) = F.g N Wi.5
EZ) E
(25)
Σ)
P(a.)) = Eg
Eer-qmie(w..))
The (wTz)_1)]]
21-1)]
(26)
where again m. V. f and are constrained to the parts of the overall weight vector that correspond to Wiy and the expectation over is taken with respect to z P(zz-1).
Thus, the invariance-abidin approximation could e.g. be computed with a layer-by-layer iterative optimisation: First, given data x ~ p(D) and prior p(w), use (23) and point estimates for all weight vectors in later layers to optimise the mean m and diagonal covariance V of the weights in the first layer, i.e., [w1,l) 1, ..... 17 Then, we can use (4 to generate P samples 21,n, of the activations of the first layer under the (now fitted) optimal [mixx (mix([wi.5)) Next, for each of these samples 1,, we can use (25) and point estimates for all weight vectors in later layers to optimise the mean and diagonal covariance V of the weights the second layer and average them for the optimal gmix of
the weights in the second layer. Finally, we can use (26) to generate samples Z2, i for the second layer.

---

This process is repeated until the last layer, and iterated until Imix (w) converges. The complexity of this procedure no larger than optimising the weight matrices of a single layer, because all previous layers are fully characterised by the distribution of the latent activations 2 of the hidden layers.
6 Related work
Poor empirical performance of VBNNS has been reported in several previous works, e-g..
67 16
For instance, it has been shown that single-layer meen-fiell VBNNs with ReLU activations can not have large predictive uncertainty between regions of low uncertainty [10]. In contrast, Farquhar et al. [ argue that mean-field approximations are expressive enough for deep BNNs; they prove a universality result that the predictive distribution of mean-field approximations of BNNs with at least 2 layers of hidden units approximate any true posterior distribution over function values arbitrarily closely. This result seems in conflict with our work and [3]. who show that the predictive distribution of mean-field VBNNs reverts to the prior predictive distribution the network width increases. The discrepancy between these two results is resolved by noting that Farquhar et al. [8] only shows that the expressive power of the model is large enough but not that the approximate inference algorithms will converge to this solution. Our work sheds light on why the mean-field approximation fails to approximate expressive posteriors via the invariance gap in the ELBO. Our framework to model the invariances similar to and extends previous preliminary works 19 22]
Surprisingly, restricting the parametrisation of the variational posterior results in competitive or even better performance [30 21 77. For example, Dusenberry et al. 7 proposes a variational approximation only for rank-1 I factors, resulting in inference of lower-dimensional subspace. The outer product of these low-rank factors perturb a weight matrix that is treated deterministically through maximum a posteriori (MAP) estimation. This approach avoids the invariance problem associated with variational Bayesian inference since the MAP estimate is not impeded by this problem and the subspace of the variational weights do not possess the same invariances.
Other attempts have proposed to approximate the posterior predictive directly, thereby circumventing the non-identifiability issue [28 20 36 While these approaches have shown promising empirical performance, estimating the KL regularisation term in function space is more complicated and can even be ill-defined due regions of zero prior probability mass [4] In a similar vein, previous work also proposed to map the BNN prior to a Gaussian process (functional) prior 9 [32]. Another promising direction to circumvent invariance in layered models
deep kernel processes, in
which Gram matrices are progressively transformed by nonlinear kernel functions (. The Gram matrices, which are treated as the random variables, are invariant to permutations/rotationss of the weights/features.
More broadly, (non- Jidenifiiii of parameters and latent variables in probabilistic models is widely-studied topic both from frequentist (27) and Bayesian perspectives 6
25
I1 24). In a recent
example, Wang et al. (35) attributed posterior collapse in the context of variational auto-encoders to non-identifiability of latent variables.
7 Conclusion
We have associated the posterior collapse phenomenon of mean-field VBNNs with invariances in the likelihood function, in particular translation invariance. While the invariance does not affect the predictive distribution, the approximations of the posterior-which abide and do not abide the invariance-diffe in the KL regularisation term and consequently in the tightness of the ELBO objective. We proved that the objectives of the two approximations differ by the relative entropy (KL) between the standard mean-field field approximation and the invariance-abiding distribution. We related this to a data-related bound on the KL regularisation, which prevents fits for which the gap is large.
We studied over- par-parametrieed Bayesian linear regression as the canonical model that exhibits translation invariance and for which the relevant terms can be computed exactly. A detailed analysis of this model confirms our hypothesis that the invariance leads to a significant additional gap the ELBO objective compared to approximations that model the invariance. It is this very gap which prevents mean-field approximation to achieve the same fit as an invariance-abiding approximation and instead leads to collapse of the posterior variances to the prior variance.

---

While we can compute the invariance gap for the over weer-parametrised linear model, we have not yet identified a computationally efficient procedure for layered models or for other invariances. However, our work provides the mathematical tools to address over-regularisation due to invariances. Future work will focus on approximations of the invariance gap order to correct the ELBO objective for translation and permutation invariances in general neural network functions.
References
[1] L. Aitchison, A. Yang, and S. W. Ober. Deep kernel processes. In M. Meila and T. Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 130-140. PMLR, 18-24 Jul 2021.
[2] C. Blundell, J. Cornebise, K. Kavukcuoglu, and D. Wierstra. Weight uncertainty in neural networks. In Proceedings of the 32nd International Conference on International Conference on Machine Learning- Volume 37, ICML' 15, page 1613-1622. JMLR.org, 2015.
[3] Brea, B. Simsek, B. Illing, and W. Gerstner. Weight-space symmetry in deep networks gives rise to permutation saddles, connected by equal-loss valleys across the loss landscape. ArXiv, abs/1907 09071.07.0022 00777.0211.111 2019.
[4] D. R. Burt, S. W. Ober, A. Garriga-Alonso, and M. van der Wilk. Understanding variational inference in function-space. In Third Symposium on Advances in Approximate Bayesian Inference, 2021.
[5] B. Coker, W. P. Bruinsma, D. R. Burt, W. Pan, and F. Doshi Velez. Wide mean-field variational Bayesian neural networks ignore the data. In Proceedings of the 25th International Conference on Artificial Intelligence and Statistics (AISTATS), 2022.
[6] A. Dawid. Conditional independence in statistical theory. Journal of the Royal Statistical Society, Series B (Methodological), 41(1):1-31, 1979. ISSN 00359246.
[7] M. W. Dusenberry, G. Jerfel, Y. Wen, Y.-A. Ma, J. Snoek, K. Heller, B. Lakshminarayanan, and
D. Tran. Efficient and scalable Bayesian neural nets with rank-1 factors. In Proceedings of the
37th International Conference on Machine Learning, ICML'20. JMLR.org. 2020.
[8] S. Farquhar, L. Smith, and Y. Gal. Liberty or depth: Deep Bayesian neural nets do not need complex weight posterior approximations. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 4346 4357. Curran Associates, Inc., 2020.
[9] D. Flam-Shepherd. Mapping gaussian process priors to Bayesian neural networks. Bayesian Deep Learning NeurIPS workshop, 2017.
[10] A. Foong, D. Burt, Y. Li, and R. Turner. On the expressiveness of approximate inference in Bayesian neural networks. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 15897-15908. Curran Associates, Inc., 2020.
[11] A. E. Gelfand and S. K. Sahu. Identifiability, improper improper priors, and gibbs sampling for generalized linear models. Journal of the American Statistical Association, 94(445):247-25 253,, 1999.
[12] A. Gelman, D. Simpson, and M. Betancourt. The prior can often only be understood in the context of the likelihood. Entropy, 19(10), 2017. ISSN 1099-4300. doi: 10.3390/e1910053
[13] S. Ghosh, J. Yao, and Doshi Velez. Structured variational learning of Bayesian neural networks with horseshoe priors. J. Dy and A. Krause, editors, Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 1744-1753, PMLR, 10-15 Jul 2018
[14] R. Herbrich. Learning Kernel Classifiers: Theory and Algorithms. The MIT Press, 2nd edition edition, 2002.
[15] G. E. Hinton and D. van Camp. Keeping the neural networks simple by minimizing the description length of the weights. In Proceedings of the Sixth Annual Conference Computational Learning Theory, COLT '93, page 5-13. Association for Computing Machinery, 1993. ISBN 0897916115. doi: 10. 1t..14//168004168000
10

---

[16] P. Izmailov, S. Vikram. M. D. Hoffman, and A. G. G. Wilson. What are Bayesian neural network posteriors really like? In M. Meila and T. Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 4629-4640. PMLR, 18-24 Jul 2021.
[17] A. Kendall and Y. Gal, What uncertainties do we need in Bayesian deep learning for computer vision? In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
[18] R. Kurle, B. Cseke, A. Klushyn, P. van der Smagt, and S. Günnemann. Continual learning with Bayesian neural networks for non-stationary data. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
[19] R. Kurle, T. Januschowski, J. Gasthaus, and B. Wang. On symmetries in variational Bayesian neural nets. In Bayesian Deep Learning NeurIPS workshop. 2021.
[20] C. Ma, Y. Li, and J. M. Hernandez-Lobato Variational implicit processes. In K. Chaudhuri and R. Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 4222-4233, PMLR, 09-15 Jun 2019.
[21] A. Mishkin, F. Kunstner, D. Nielsen, M. W. Schmidt, and M. E. Khan. Slang: Fast structured covariance approximations for Bayesian deep learning with natural gradient. In NeurlPS, pages 6248-6258, 2018.
[22] D. A. Moore. Symmetrized variational inference. NIPS Workshop on Advances in Approximate Bayesian Inference, December 2016.
[23] C. V. Nguyen, Y. Li, D. Bui, and R. E. Turner. Variational continual learning. In International Conference on Learning Representations, 2018.
[24] R. Nishihara, T. P. Minka, and D. Tarlow. Detecting parameter symmetries in probabilistic models. arXiv: Machine Learning. 2013.
[25] D. J. Poirier. Revising beliefs in nonidentified models. Econometric Theory, 14(4):483-509
1998.
[26] A. Pourzanjani, R. M. Jiang, and L.. Petzold. Improving the identifiability of neural networks for Bayesian inference. In Bayesian Deep Learning NeurIPS workshop, 2017.
[27] B. L.. S. Prakasa Rao. Identifiability in Stochastic Models. Academic Press, Oxford, 1992.
[28] S. Sun, G. Zhang. J. Shi, and R. B. Grosse. Functional variational Bayesian neural networks. ArXiv, abs/1903.05579, 2019.
[29] H. Sussmann. Uniqueness of the weights for minimal feedforward nets with given input- output map. Neural Networks, 5(4):589-593, 1992. ISSN 0893-6080. doi: http:://oi.orr//I 1016/50893-6088005)88000
[30] J. Swiatkowski, K. Roth, B. S. Veeling, L. Tran, J. V. Dillon, J. Snoek, S. Mandt, T. Salimans, R. Jenatton, and S. Nowozin. The k-tied normal distribution: A compact parameterization of gaussian mean field posteriors in Bayesian neural networks. In Proceedings of the 37th International Conference on Machine Learning, ICML'20. JMLR.org. 2020.
[31] M. Tomezak, S. Swaroop, and R. Turner. Efficient low rank gaussian variational inference for neural networks. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 4610-4622. Curran Associates, Inc., 2020.
[32] B.-H. Tran, S. Rossi, D. Milios, and M. Filippone. All you need is a good functional prior for Bayesian deep learning. Journal of Machine Learning Research, 23(74):1-56, 2022.
[33] B. Trippe and R. Turner. Overpruning in variational Bayesian neural networks. arXiv preprint arXiv:
2018.
[34] M. J. Wainwright, M. I. Jordan, et al. Graphical models, exponential families, and variational inference. Foundations and Trends@ in Machine Learning, 1(1-2):1-305, 2008
[35] Y. Wang, D. Blei, and J. P. Cunningham. Posterior collapse and latent variable non-identifiability. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan, editors, Advances in
11

---

Neural Information Processing Systems, volume 34, pages 5443-5455, Curran Associates, Inc., 2021.
[36] Z. Wang. T. Ren, J. Zhu, and B. Zhang. Function space particle optimization for Bayesian neural networks. In International Conference on Learning Representations, 2019.
[37] Wenzel, K. Roth, B. Veeling. J. Swiatkowski, L. Tran, S. Mandt, J. Snoek, Salimans, R. Jenatton, and S. Nowozin. How good is the Bayes posterior in deep neural networks really? H. D. III and A. Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 10248-10259, PMLR, 13-18 Jul 2020.
[38] A. G. Wilson and Izmailov. Bayesian deep learning and a probabilistic perspective of generalization. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 4697-4708. Curran Associates, Inc., 2020.
Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] We point out limitations and future work throughout the paper.
(c) Did you discuss any potential negative societal impacts of your work? [No] This is purely theoretical work without direct negative societal impacts.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes]
(b) Did you include complete proofs of all theoretical results? [Yes]
3. If you eNper No experiments
4. If you using existing assets (e.g., code, data, models) or curating/releasing new assets...
[No] No existing assets used.
5. If you used crowdsourcing or conducted research with human subjects... [No] No crowdsourcinghhuman subjects.
12


Reviews:
[{"review_id": "beSp3c6iGG", "reviewer": "Reviewer_BBQR", "summary": "This paper evaluates the impact of the influence of the invariances in the likelihood function on the ELBO, when performing Variational Bayes. The authors propose to enforce posteriors to be invariant to a few transformations (permutations of parameters, translations). Compared to the usual posteriors, the posteriors obtained this way achieve (of course) the same performance in terms of likelihood, but lead to a much larger KL term in the ELBO bound. That is, changing a posterior $q\\_0$ to an invariant posterior $q\\_{\\mathrm{mix}}$ increases the distance of the posterior to the prior $p$. \n\nThis observation can be useful to explain and fix the \"posterior-to-prior collapse\" effect (see [4], \"Wide Mean-Field Variational Bayesian Neural Networks Ignore the Data\", Coker et al., 2021). According to [4], when performing Variational Bayes (VB) on a Neural Network (NN) whose number of neurons tends to infinity, the variational posterior (i.e., the distribution optimizing the ELBO) tends to the prior. The authors of the present work claim that their study will be useful in the future to understand and overcome this effect, since they are able to affect the KL term of the ELBO without changing its likelihood term.", "questions": "-", "limitations": "-", "rating": 4, "confidence": 2, "soundness": 2, "presentation": 1, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n\nThe impact of the invariances of a NN on the optimization process and the best achievable validation loss is a very hot topic, with a wide range of works spanning many aspects of the NNs, for instance: implicit bias of the SGD [A], regularization effect of overparameterization [B].\n\nMoreover, this work is in line with [4], which sheds light on inconsistencies in VB applied to NNs.\n\n## Weaknesses\n\nThe major weakness of the paper is its readability. \n\n### Readability\n\nThis paper is written as a succession of statements and sketches of proofs, without highlighting anything. Providing technical details in the main part of the paper is a good idea. But, at least, the authors should be clearer about the nature of each of their paragraphs. \n\nFor instance, definitions may be put into \"definition\" environments with \"Definition X\" written in bold. Same with claims, which may be put into \"Theorem\" or \"Proposition\" environments, and sketches of proofs, which may be put into \"Proof\" environments.\n\nThe term \"mean-field\" is not well defined by the authors. For me, \"mean-field variational Bayes\" means \"VB where the posterior distribution is a product of distributions over each variable\". That is, independence between variables in the posterior. I did not manage to check whether I was right or not by reading the paper.\n\nBesides, I was not able to fully understand the paragraph lines 108-128. In particular, the function $g\\_0(w, \\theta)$ \"approximating\" $l(w, \\mathcal{D})$ looks like a way to introduce the \"variationalized\" model, but there is not enough information and precise definition to be sure of that. Besides, there is no definition of $l(w, \\mathcal{D})$.\n\nSeveral mathematical objects are not well or explicitly defined, which decreases the readability of the paper. For instance:\n * Equ. (7): $Z\\_0(\\mathbf{r})$ is not defined, even in the appendix;\n * Equ. (7): what is $f$?\n\nNotes: \n * denoting a variance by \"$\\sigma$\": \"$\\sigma^2$\" must be used instead;\n * headers of LaTeX \"paragraphs\" should not end with \":\", but with \".\" or nothing;\n * what is the signification of the \"!\" over the \"=\" (see Equ. (7), (8) and many equations in the appendix)?\n * line 166: \"models\" -> \"model\";\n * line 116: \"see 4\" -> \"see Section 4\"\n\n### Correctness\n\nSee appendix, Equations (46)-(50):\n * $\\mathrm{d}\\mathbf{w}$ and $\\mathrm{d}\\mathbf{r}$ must be swapped, to be consistent with the integration order;\n * the authors must use parentheses when using $\\ln$ on several factors: prefer $\\ln(ab)$ instead of $\\ln a b$;\n * the use of Jensen's inequality-equality between (47) and (48) is very obscure and needs an explanation; if $g\\_{\\mathrm{mix}}(\\cdots)$ is constant in $\\mathbf{r}$, there is no need to use Jensen (and we recover the authors' result), if it is not constant in $\\mathbf{r}$, how do the authors justify the \"Jensen's equality\" they used?\n\nAbout the \"infinite variance Gaussian distribution\" (see lines 115-116 and Equation line 191): there is no theoretic justification or discussion about this choice. The main question to answer is: would we get the same result with any sequence of distributions tending to an improper \"uniform distribution\" over $\\mathbb{R}$? \n\n## References\n\n[A] \"In Search of the Real Inductive Bias: On the Role of Implicit Regularization in Deep Learning\", Neyshabur et al., 2015\n\n[B] \"Stochastic complexity of Bayesian networks\", Yamazaki and Watanabe, 2002", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "YunMGcQOMJC", "reviewer": "Reviewer_ZF6G", "summary": "This paper posits that the main reason that overparameterized Bayesian neural networks tend to collapse to the prior distribution is attributed to the fact that there exist invariances in the likelihood with respect to model parameters. Conversely, the paper proposes that should these invariances be accounted for and handled properly then this behavior can be avoided. The paper develops an analysis framework to tackle this as well as investigates a case-study of accommodating translational invariance in a particular linear model and a general feed-forward Bayesian neural network.", "questions": "1. On line 92, what does the $k+1$ in $\\sigma_{k+1}$ indicate?\n2. On lines 199 and 200, you assume that \"all variances and means take the same value for the approximation and the prior\". Can you expand on why this is okay to make this assumption?\n3. To better my understanding, what exactly is holding back a practitioner from trying to find $\\theta^*_\\text{mix}$ (Eq. 19a) for say a generic Bayesian neural network? In other words, what is the immediate roadblock to optimizing for this objective that requires further research?\n4. When deriving the invariance gap in the Appendix, on Eq. 48 after line 470 Jensen's inequality is applied to move a logarithm outside the inner integral; however, the inequality is used as an equality. I have two concerns with this. The first being that the logarithm essentially \"skipped past\" a scalar constant (i.e., what happened: $\\int c \\log(x) dx \\rightarrow \\log \\int cx dx$), which I do not believe is applicable for Jensen's inequality. Second, Jensen's inequality is only an equality if the function being brought outside of the integral is affine (which $\\log$ is not) or that the remainder of the integrand is constant (which $Z_\\text{mix}g_\\text{mix}(t^{-1}(\\mathbf{w},f(\\mathbf{r}))d\\mathbf{w}$ is not w.r.t. $\\mathbf{w}$). Can you please expand on why this was a valid operation to do as I may have missed something. Should this not have been valid, how does this then effect the invariance gap and subsequent findings?", "limitations": "The authors have adequately addressed the work's limitations.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper proposes a novel hypothesis to the often cited phenomenon of a Bayesian neural network's posterior collapsing to the prior distribution. Backing this up a new bound is developed, the invariance gap, that serves to illustrate in what scenarios will a model collapse (resulting in a invariance gap of 0). I believe these findings to be significant and meaningful to the community. That being said, I found the work to be very dense and hard to follow at times. Particularly the paragraph titled \"Likelihood and Posterior approximations\" starting on line 108 was difficult and required rereading many times over to follow.\n\nAs for the work done on the translation invariant, the work largely seemed meaningful. While it would have been nice to see it actually applied in an experimental setting, I do recognize that this isn't in the intended scope of the paper. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "Nk61sEa3ZhF", "reviewer": "Reviewer_11Q5", "summary": "The paper provides a novel take to the long-standing problem of posterior collapse in training mean-field variational Bayesian neural networks (VBNNs). It is argued that the invariances inherent to the architures of the network lead to posterior distributions with complicated structures that are hard to capture with mean-field variational families. A quantitative measure, referred to as the \"invariance gap\" is proposed to study how this invariance hinders variational inference. The translational invariances in Bayesian linear regression and deep neural network models are studied under this framework.", "questions": "- It is not very clear to me how Figure 2 should be interpreted: basically, we have two sets of optimal parameters: $\\theta_0^*$ and $\\theta_{\\text{mix}}^*$. As dimensionality increase, the invariance gap $\\text{KL}[q_0 || q_{\\text{mix}}]$ dominates when you plug in $\\theta_{\\text{mix}}^*$ but not for $\\theta_{0}^*$. How does one arrive at the conclusion that invariance-gap incentivises posterior collapse? What does it mean for the gap for the mean-field approximation to vanish?\n\n- Does the proposed invariance-abiding variational family lead to less posterior collapse in practice? In the case where some but not all of the invariances are addressed with this mixture variational family, can we still see improvements nontheless?\n\n- We know from the KL gap that collapsing to the prior can be a local minimum that is easy to obtain. However, are there any insights on the training dynamics that lead to this local minimum, instead of learning one posterior mode that \"breaks the symmetry\" of the system? From experience, even for problems without this type of invariance, bad hyperparameters and/or initialization could still lead to posterior collapse.", "limitations": "The biggest limitation of this work is the lack of practical procedures to address the invariance problem in practice, which the authors addressed clearly in the conclusion section.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper presents a novel idea in a clear and mathematically precise manner. In-depth theoretical analysis is provided which supports the ideas presented. The background knowledge (Bayesian neural network and variational inference) is presented in a way that is easy to follow. I also like how the idea of the detriments of invariances are introduced and discussed point by point, gradually giving the reader a clearer picture of what is going on. \n\nHowever, there are a couple of weak points later on in the paper. The theoretical results sections (subsections 4.1-4.3) are a bit unnecessarily dense with technical details, with not enough exposition and discuss on the significance of these results. Subsection 4.3 in particular has a lot of formulae but very little discussion of the significance of these results. I wish that the theoretical results can be presented more concisely, giving more room for more insightful discussions, such as potential solutions to the invariance problem.\n\nOverall the paper is well written and the contributions seem solid, with some room for improvement in the theoretical sections.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "A4w58enT-50", "reviewer": "Reviewer_aga4", "summary": "This paper looks at why mean-field variational inference can collapse to the prior with little data / large models, motivated by observations of this effect with Bayesian neural networks in the field. They define an \"invariance-abiding\" posterior approximation that is a function of the mean-field approximation, and introduce conditions on its form so that the likelihood term is the same, and the KL-to-prior term can be analysed. They look at translation invariance in a Bayesian linear regression model, showing both mathematically and empirically that the mean-field version reverts to the prior as the number of parameters increases, while the invariance-abiding alternative does not. There is then an initial attempt to use similar analysis for neural networks. ", "questions": "I have provided detailed questions and suggestions in the Strengths and Weaknesses section above. \n\n---------------\nAfter rebuttal: \n\nI thank the authors for improving the writing and clarity. I have increased my score. I did not increase by more because I still have some missing intuition between the linear model and neural networks (again, I stress this is important because of how the paper currently sells this as why it is important). I do appreciate the attempts the authors made already. I hope the authors can close this gap in future work to validate the importance of their analysis in this work. ", "limitations": "No potential negative societal impacts that I can see. ", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Overall, I think this kind of analysis is extremely relevant and of significance to the Bayesian deep learning community. Although there has been some previous work in this area (eg Kurle et al and Coker et al), a lot more can still be done, and this is a big strength of this paper for me. \n\nAlthough the maths and approach for invariance-abiding posterior approximations for the linear model were overall nice, it is a toy example. In my mind, the key insight would be relating this maths to a less toy model such as a deep linear network or a full neural network with non-linearities. This key link is missing, as Section 4.3 does not provide any takeaways or insights. I would strongly consider raising my score if the authors could provide any intuitions or takeaways in Section 4.3. For example, can we say how the KL gap grows (intuitively)? Or are there any simplifications one can make to provide any such statement? Or anything else. Note that I am looking for any small but relevant statement for this paper (and any deeper analysis can be future work). I want something so that I know that this analysis may be useful in the deep learning case; right now it is still possible that this analysis is not useful at all because everything is just completely intractable and does not transfer over to neural networks. \n\nI would also consider changing my view if I am convinced that the analysis for the linear model is a strong contribution in itself, however, I currently do not view it as enough for a paper (and this is not how the paper is written either). \n\nI thought the Abstract and Introduction were written very well. However, I found the maths in the rest of the paper difficult to follow, and had to read it a few times before understanding most of it. It was very easy to confuse what was 'mix' and not, and although there was an attempt to provide some intuition before writing some equations, I think this can be looked over again and made simpler to parse. I think spending time on this would make the paper much more accessible and hence impactful to the community. It did not help that some things were not clearly defined / there were some typos: (i) Line 71, 'y' undefined, (ii) missing 'ln' in front of p(w|D) in line 80, (iii) line 90 should be $\\sigma_y^2$, (iv) missing close bracket in Equation after line 117, (v) what is 'f' in Equation 7?, (vi) presumably 'lambda' is the variance of g(w) -- this is in line 123 but the way it is written could also be that {m, lambda} are means and variances of the approximate q, and not of g, (vii) Line 191 Equation does not define beta, (viii) Equation 17, mu is undefined (presumably it is prior mean, so zero?), (ix) Equation 21 should be $m^*_{mix}$ not $q_{mix}$. \n\nFinally, the authors may also be interested in the paper, \"On the expressiveness of approximate inference in Bayesian neural networks\", Foong et al, 2021.\n\n(Note that I did not check the derivations/proofs in the supplementary material.) ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept