Original: On the detrimental effect of invariances in the likelihood for variational inference
Revised: On the detrimental effect of invariances in the likelihood for variational inference

Original: Anonymous Author(s)

Original: AffiliationAddress email

Original: Abstract
Revised: Abstract

Original: Variational Bayesian posterior inference often requires simplifying approximationssuch as mean-field parametrisation to ensure tractability.
Revised: Variational Bayesian posterior inference often requires simplifying approximations such as mean-field parametrisation to ensure tractability.

Original: However, prior work hasassociated the variational mean-field approximation for Bayesian neural networkswith underfitting in the case of small datasets or large model sizes.
Revised: 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.

Original: In this work,we show that invariances in the likelihood function of over-parametrised modelscontribute to this phenomenon because these invariances complicate the structureof the posterior by introducing discrete and/or continuous modes which cannot bewell approximated by Gaussian mean-field distributions.
Revised: 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.

Original: In particular, we show thatthe mean-field approximation has an additional gap in the evidence lower boundcompared to a purpose-built posterior that takes into account the known invariances.
Revised: 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.

Original: Importantly, this invariance gap is not constant; it vanishes as the approximationreverts to the prior.
Revised: Importantly, this invariance gap is not constant; it vanishes as the approximation reverts to the prior.

Original: We proceed by first considering translation invariances ina linear model with a single data point in detail.
Revised: We proceed by first considering translation invariances in a linear model with a single data point in detail.

Original: We show that, while the trueposterior can be constructed from a mean-field parametrisation, this is achievedonly if the objective function takes into account the invariance gap.
Revised: 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.

Original: Then, wetransfer our analysis of the linear model to neural networks.
Revised: Then, we transfer our analysis of the linear model to neural networks.

Original: Our analysis providesa framework for future work to explore solutions to the invariance problem.
Revised: Our analysis provides a framework for future work to explore solutions to the invariance problem.

Original:  Introduction
Revised:  Introduction

Original: Bayesian neural networks (BNNs) have several appealing advantages compared to deterministicneural networks (NN) such as improving generalization [35], capturing epistemic uncertainty [16] and providing a framework for continual learning methods [17, 21].
Revised: Bayesian neural networks (BNNs) have several appealing advantages compared to deterministic neural networks (NN) such as improving generalization [36], capturing epistemic uncertainty [16], and providing a framework for continual learning methods [17, 22].

Original: Unfortunately, reaping thesetheoretical benefits has so far been impeded [34].
Revised: Unfortunately, reaping these theoretical benefits has so far been impeded [35].

Original: In particular, several practical issues with variational
Revised: In particular, several practical issues with variational

Original: Bayesian inference methods—which are the de-facto standard technique for scaling inference in
Revised: Bayesian inference methods—which are the de-facto standard technique for scaling inference in

Original: BNNs to large datasets [14, 1]—have been identified to be partially responsible for a performance gapcompared to deterministic NNs[15].
Revised: BNNs to large datasets [14, 1]—have been identified to be partially responsible for a performance gap compared to deterministic NNs[15].

Original: The most common variational approximation of the posterior is aproduct of independent Normal distributions, commonly referred to as the mean-field approximation.
Revised: The most common variational approximation of the posterior is a product of independent Normal distributions, commonly referred to as the mean-field approximation.

Original: It has been observed that mean-field variational BNNs (VBNNs) suffer from severe underfitting forlarge models and small dataset sizes [12].
Revised: It has been observed that mean-field variational BNNs (VBNNs) suffer from severe underfitting for large models and small dataset sizes [12].

Original: Recent work [4] showed that under certain assumptionssuch as an odd Lipschitz activation function and a finite dataset with bounded likelihood, the optimalmean-field approximation collapses to the prior as the NN width increases, ignoring the data .
Revised: Recent work [4] 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 to the prior as the NN width increases, ignoring the data .

Original: The goal of this work is to shed light on why the mean-field variational approximation collapsesand to provide a new angle for future works to address this shortcoming.
Revised: The goal of this work is to shed light on why the mean-field variational approximation collapses and to provide a new angle for future works to address this shortcoming.

Original: To this end, we studyinvariances in the likelihood function of overparametrised models and their detrimental effect onvariational inference.
Revised: To this end, we study invariances in the likelihood function of overparametrised models and their detrimental effect on variational inference.

Original: For instance, it is well known that NNs exhibit several invariances with respectto their parameters , including node permutation invariance [18], sign-flip invariance (in the caseof odd activation functions) [26, 2], scaling invariance (in the case of piece-wise linear activations) [23], and as we will show in Sec. 5, the parameter space of BNNs additionally has subspaces withtranslation invariance .
Revised: For instance, it is well known that NNs exhibit several invariances with respect to their parameters , including node permutation invariance [18], sign-flip invariance (in the case of odd activation functions) [27, 2], scaling invariance (in the case of piece-wise linear activations) [24], and as we will show in Sec. 5, the parameter space of BNNs additionally has subspaces with translation invariance .

Original: Invariance in the likelihood function does not necessarily pose a challenge if maximum likelihoodpoint estimation through stochastic gradient descent is used, because convergence to any of theequivalent optima (resulting from the invariance) suffices for prediction.
Revised: Invariance in the likelihood function does not necessarily pose a 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.

Original: However, these invariancesare detrimental to the variational Bayesian approach.
Revised: However, these invariances are detrimental to the variational Bayesian approach.

Original: As a first step to show this, we isolate the impactof the invariances in Sec. 3.
Revised: As a first step to show this, we isolate the impact of the invariances in Sec. 3.

Original: To this end, we construct both a mean-field as well as an invarianceabiding approximation which explicitly models the invariances by integrating over all transformationsthat leave the likelihood invariant.
Revised: To this end, we construct both a mean-field as well as an invarianceabiding approximation which explicitly models the invariances by integrating over all transformations that leave the likelihood invariant.

Original: Notably, both aforementioned posterior approximations areconstructed from the same (mean-field) likelihood approximation.
Revised: Notably, both aforementioned posterior approximations are constructed from the same (mean-field) likelihood approximation.

Original: We then prove that, under theconditions outlined in Sec. 3, the mean-field approximation induces the same posterior predictivedistribution as the invariance-abiding approximation.
Revised: 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.

Original: However, we also prove that the ELBO objectivecorresponding to these two approximations differs by the KL divergence between both posteriorapproximations; we refer to this difference as invariance-gap .
Revised: 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 .

Original: Importantly, the gap vanishes if bothdistributions are identical, which is the case if the mean-field approximation reverts to the prior.
Revised: Importantly, the gap vanishes if both distributions are identical, which is the case if the mean-field approximation reverts to the prior.

Original: We then demonstrate the detrimental effect of invariances in the likelihood function of anoverparametrised Bayesian linear regression model in Sec. 4.
Revised: We then demonstrate the detrimental effect of invariances in the likelihood function of an overparametrised Bayesian linear regression model in Sec. 4.

Original: This model is purposely selectedas the canonical model exhibiting (only) translation invariance .
Revised: This model is purposely selected as the canonical model exhibiting (only) translation invariance .

Original: We provide a detailed analysis ofthis model, including a tractable solution for the invariance-abiding approximation.
Revised: We provide a detailed analysis of this model, including a tractable solution for the invariance-abiding approximation.

Original: It turns out thatthe optimal parameters wrt.
Revised: It turns out that the optimal parameters w.r.t.

Original: the ELBO objective with the invariance-abiding distribution result inthe true posterior.
Revised: the ELBO objective with the invariance-abiding distribution result in the true posterior.

Original: In contrast, posterior approximations with parameters that are optimal wrt.
Revised: In contrast, posterior approximations with parameters that are optimal w.r.t.

Original: themean-field ELBO revert to the prior as the number of dimensions increases.
Revised: the mean-field ELBO revert to the prior as the number of dimensions increases.

Original: Finally, we transfer our analysis of the linear model to VBNNs in Sec. 5, showing that subspacesin BNNs exist that exhibit translation invariance.
Revised: Finally, we transfer our analysis of the linear model to VBNNs in Sec. 5, showing that subspaces in BNNs exist that exhibit translation invariance.

Original: Combined with a previous analysis of the nodepermutation invariance (cf.
Revised: Combined with a previous analysis of the node permutation invariance (cf.

Original: Kurle et al.
Revised: Kurle et al.

Original: [18] or App.
Revised: [18] or App.

Original: E of the supplementary material), this providesthe basis for future work to approximate the (generally intractable) invariance-gap and therebyoptimise for a tighter and favorable ELBO objective.
Revised: E of the supplementary material), this provides the basis for future work to approximate the (generally intractable) invariance gap and thereby optimise for a tighter and favorable ELBO objective.

Original: We start in Sec. 2 by introducing the basicconcepts of VBNNs and recent results on which we build our contribution.
Revised: We start in Sec. 2 by introducing the basic concepts of VBNNs and recent results on which we build our contribution.

Original:  Background2.1 Variational Bayesian Neural Networks
Revised:  Background 2.1 Variational Bayesian Neural Networks

Original: Neural network functional model.
Revised: Neural network functional model.

Original: Deep NNs are layered models that progressively transformtheir inputs in each layer.
Revised: Deep NNs are layered models that progressively transform their inputs in each layer.

Original: More formally, an L –layer NN computes algebraically
Revised: More formally, an L –layer NN computes algebraically

Original: [Equation]
Revised: [Equation]

Original: where the h l : R → R are monotonic transfer functions that introduce non-linearities.
Revised: where the h l : R → R are monotonic transfer functions that introduce non-linearities.

Original: Such a networkhas n 1 + n 2 + · · · + n L many nodes z l,i , where i indexes the layer-wise vectors z l .
Revised: Such a network has n 1 + n 2 + · · · + n L many nodes z l,i , where i indexes the layer-wise vectors z l .

Original: These nodes areeach a weighted linear combination of either the input vector x (first hidden layer) or the value of thehidden units z l (all other hidden layers and the output f ( x ) ).
Revised: These nodes are each a weighted linear combination of either the input vector x (first hidden layer) or the value of the hidden units z l (all other hidden layers and the output f ( x ) ).

Original: We use w =
Revised: We denote all learnable parameters of the model by the stacked weight vectors of each layer and node, w =

Original: [ w T 1 ,  . . .
Revised: [ w T 1 , 1 , . . .

Original: , w T n 1 + ··· + n L ] T todenote all learnable parameters of the model.
Revised: , w T L,n L ] T .

Original: Variational Bayesian treatment.
Revised: Variational Bayesian treatment.

Original: If the weights and biases w are treated as random variableswith a prior distribution p ( w ) , then the posterior p ( w | D ) —induced by the dataset D through thelikelihood function ℓ ( w ; D ) :
Revised: If the weights and biases w are treated as random variables with a prior distribution p ( w ) , then the posterior p ( w | D ) —induced by the dataset D through the likelihood function ℓ ( w ; D ) :

Original: = p ( D | w ) that is defined via f ( x ) —is referred to as a Bayesianneural network (BNN).
Revised: = p ( D | w ) that is defined via f ( x ) —is referred to as a Bayesian neural network (BNN).

Original: The variational Bayesian method approximates the posterior by a distributionq θ ( w ) ≈ p ( w | D ) with variational parameters θ , casting inference as an optimization problem
Revised: The variational Bayesian method approximates the posterior by a distribution q θ ( w ) ≈ p ( w | D ) with variational parameters θ , casting inference as an optimization problem

Original: [Equation]
Revised: [Equation]

Original: where KL [ q θ || p ( · | D )] :=
Revised: where KL [ q θ || p ( · | D )] :=

Original: E w ∼ q θ
Revised: E w ∼ q θ

Original: [ln q θ ( w ) − ln p ( w |D )] and Q is a family of distributions over w .
Revised: [ln q θ ( w ) − ln p ( w |D )] and Q is a family of distributions over w .

Original: The optimization of (1) is achieved by maximising a lower bound to the (log) model evidence ln p ( D ) ,
Revised: The optimization of (1) is achieved by maximising a lower bound to the (log) model evidence ln p ( D ) ,

Original: [Equation]
Revised: [Equation]

Original: where ELL ( q θ , D ) :
Revised: where ELL ( q θ , D ) :

Original: =
Revised: =

Original: E w ∼ q θ
Revised: E w ∼ q θ

Original: [ln ℓ ( w ; D )] is the expected log-likelihood.
Revised: [ln ℓ ( w ; D )] is the expected log-likelihood.

Original: Definition 1 (Mean-field variational BNN) .
Revised: Definition 1 (Mean-field variational BNN) .

Original: A mean-field variational BNN (VBNN) is the minimizerof (1) where Q is the family of Gaussians with diagonal covariance matrix.
Revised: A mean-field variational BNN (VBNN) is the minimizer of (1) where Q is the family of Gaussians with diagonal covariance matrix.

Original: 2.2 Data-related bound on the KL divergence
Revised: 2.2 Data-related bound on the KL divergence

Original: The term KL [ q θ || p ] in (2) admits a data-dependent upper bound that naturally occurs as aconsequence of the finite information provided by a finite dataset in the presence of noise.
Revised: The term KL [ q θ || p ] in (2) admits a data-dependent upper bound that naturally occurs as a consequence of the finite information provided by a finite dataset in the presence of noise.

Original: This resulthas been shown in [4] and we recall the result for a Gaussian likelihood with homogeneous noise (forother likelihood functions, we refer to App.
Revised: This result has been shown in [4] and we recall the result for a Gaussian likelihood with homogeneous noise (for other likelihood functions, we refer to App.

Original: F in [4]).
Revised: F in [4]).

Original: Assume a VBNN with Gaussian observation noise defined through y =
Revised: Assume a VBNN with Gaussian observation noise defined through y =

Original: f w ( x )+ ϵ , with ϵ ∼ N (0 , σ 2 y ) ,an isotropic Gaussian prior p ( w ) =
Revised: f w ( x )+ ϵ , with ϵ ∼ N (0 , σ 2 y ) , an isotropic Gaussian prior p ( w ) =

Original: N ( w ; 0 , I ) , and a mean-field variational approximation q θ ( w ) =
Revised: N ( w ; 0 , I ) , and a mean-field variational approximation q θ ( w ) =

Original: N ( w ; m , Diag( v )) .
Revised: N ( w ; m , Diag( v )) .

Original: Define the NN output variance σ 2 L (cid:0) x ( n ) (cid:1) :
Revised: Define the NN output variance σ 2 L (cid:0) x ( n ) (cid:1) :

Original: = V w ∼ p (cid:2) f
Revised: = V w ∼ p (cid:2) f

Original: ( x ( n ) ; w ) (cid:3) for somefixed input x ( n ) , where L indicates the last layer.
Revised: ( x ( n ) ; w ) (cid:3) for some fixed input x ( n ) , where L indicates the last layer.

Original: Then,
Revised: Then,

Original: [Equation]
Revised: [Equation]

Original: where q ∗ is a hypothetical approximation that predicts the data perfectly up to the noise variance σ 2 y(see App.
Revised: where q ∗ is a hypothetical approximation that predicts the data perfectly up to the noise variance σ 2 y (see App.

Original: F for details).
Revised: F for details).

Original: Prior work [4] has used the data-related bound to prove that the predictivedistribution of mean-field BNNs converges to the prior predictive distribution if the network widthis large and the activation function is odd, ultimately resulting in posterior collapse to the prior.
Revised: Prior work [4] has used the data-related bound to prove that the predictive distribution of mean-field BNNs converges to the prior predictive distribution if the network width is large and the activation function is odd, ultimately resulting in posterior collapse to the prior.

Original: Incontrast, we address the question why the mean-field approximation cannot approximate the posteriorby relating this data-related bound to invariances in neural network functions in the following section.
Revised: 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 in the following section.

Original:  Invariance-abiding variational approximation
Revised:  Invariance-abiding variational approximation

Original: We develop a framework that enables modelling the invariances in the likelihood and understandingtheir impact on the ELBO objective.
Revised: We develop a framework that enables modelling the invariances in the likelihood and understanding their impact on the ELBO objective.

Original: To achieve this, we approximate the likelihood by a Gaussianfunction with variational parameters and marginalise over all transformations to which the truelikelihood is invariant.
Revised: 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.

Original: By taking the product with the prior, we construct an invariance-abidingposterior approximation q mix which can be related to a mean-field approximation q 0 that does notmodel invariances.
Revised: By taking the product with the prior, we construct an invariance-abiding posterior approximation q mix which can be related to a mean-field approximation q 0 that does not model invariances.

Original: We then describe conditions under which both approximations yield an identicalposterior predictive, while the KL regularisation term of q 0 is lower bounded by the KL of q mix .
Revised: We then describe conditions under which both approximations yield an identical posterior predictive, while the KL regularisation term of q 0 is lower bounded by the KL of q mix .

Original: Variational likelihood and posterior approximations.
Revised: Variational likelihood and posterior approximations.

Original: Non-identifiability implies that theposterior does not concentrate on a single set of parameters irrespective of the dataset size, becausethe same likelihood is assigned to different parameter values [11].
Revised: 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 [11].

Original: We model this invariance throughtransformations t ( · , r ) to which the likelihood ℓ ( w ; D ) =
Revised: We model this invariance through transformations t ( · , r ) to which the likelihood ℓ ( w ; D ) =

Original: p ( D | w ) is invariant via variables r :
Revised: p ( D | w ) is invariant via variables r :

Original: [Equation]
Revised: [Equation]

Original: From (4) it follows that the likelihood is also invariant wrt.
Revised: From (4) it follows that the likelihood is also invariant w.r.t.

Original: the marginalisation
Revised: the marginalisation

Original: [Equation]
Revised: [Equation]

Original: We assume that each of the equivalent parametrisations w ′ =
Revised: We assume that each of the equivalent parametrisations w ′ =

Original: t ( w , r ) of the likelihood has the sameprobability a priori.
Revised: t ( w , r ) of the likelihood has the same probability a priori.

Original: For discontinuous transformations such as node permutations (see App. E), p ( r )is a uniform distribution over discrete variables indexing these transformations.
Revised: For discontinuous transformations such as node permutations (see App. E), p ( r ) is a uniform distribution over discrete variables indexing these transformations.

Original: For the continuoustranslation invariance (see Sec. 4), we model the uniform distribution as a Gaussian with infinitevariance.
Revised: For the continuous translation invariance (see Sec. 4), we model the uniform distribution as a Gaussian with infinite variance.

Original: It is conceivable that the likelihood can be constructed by marginalising over thesetransformations of a simpler function ℓ 0 ( w ; D ) such that ℓ ( w ; D ) =
Revised: It is conceivable that the likelihood can be constructed by marginalising over these transformations of a simpler function ℓ 0 ( w ; D ) such that ℓ ( w ; D ) =

Original: E r ∼ p
Revised: E r ∼ p

Original: [ ℓ 0 ( t ( w , r ); D )] .
Revised: [ ℓ 0 ( t ( w , r ); D )] .

Original: Forinstance, permutation invariance induces a factorial number of discontinuous modes that are eachequivalent (cf.
Revised: For instance, permutation invariance induces a factorial number of discontinuous modes that are each equivalent (cf.

Original: [18], App. E); and as we show in Sec. 4, the full covariance Gaussian likelihood andposterior of an over-parametrised Bayesian linear regression model can be constructed from a productof independent Gaussians.
Revised: [18], 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 a product of independent Gaussians.

Original: Curiously, it may be sufficient to approximate ℓ 0 while taking into accountthe known invariances.
Revised: Curiously, it may be sufficient to approximate ℓ 0 while taking into account the known invariances.

Original: To this end, we define a Gaussian variational likelihood approximationg 0 ( w ; θ ) ≈ ℓ 0 ( w , D ) , with variational parameters θ .
Revised: To this end, we define a Gaussian variational likelihood approximation g 0 ( w ; θ ) ≈ ℓ 0 ( w , D ) , with variational parameters θ .

Original: From this single mode approximation, wemodel an invariance-abiding likelihood approximation using the same parametrisation through
Revised: From this single mode approximation, we model an invariance-abiding likelihood approximation using the same parametrisation through

Original: [Equation]
Revised: [Equation]

Original: We consider mean-field Gaussians for the prior p ( w ) and likelihood approximation g 0 ( w ; θ ) , whereθ = { m , λ } are means and variances of g 0 .
Revised: We consider mean-field Gaussians for the prior p ( w ) and likelihood approximation g 0 ( w ; θ ) , where θ = { m , λ } are means and variances of g 0 .

Original: We then define a mean-field posterior as the product
Revised: We then define a mean-field posterior as the product

Original: [Equation]
Revised: [Equation]

Original: Similarly, we define an invariance-abiding posterior as the product of prior and invariant likelihood:
Revised: Similarly, we define an invariance-abiding posterior as the product of prior and invariant likelihood:

Original: [Equation]
Revised: [Equation]

Original: where Z 0 and Z mix are normalisation constants.
Revised: where Z 0 and Z mix are normalisation constants.

Original: While q 0 ( w ; θ ) is a mean-field approximation, theinvariance-abiding approximation q mix ( w ; θ ) is a mixture with infinite continuous or finite discretemodes depending on the type of invariance.
Revised: While q 0 ( w ; θ ) is a mean-field approximation, the invariance-abiding approximation q mix ( w ; θ ) is a mixture with infinite continuous or finite discrete modes depending on the type of invariance.

Original: We will describe continuous translation invariance in
Revised: We will describe continuous translation invariance in

Original: Sec. 4; for the discrete node permutation invariance, see App.
Revised: Sec. 4; for the discrete node permutation invariance, see App.

Original: E of the supplementary material.
Revised: E of the supplementary material.

Original: Posterior predictive equivalence.
Revised: Posterior predictive equivalence.

Original: Next, we discuss conditions under which we can constructapproximations q mix ( w ; θ ) and q 0 ( w ; θ ) that yield the same predictive distribution.
Revised: Next, we discuss conditions under which we can construct approximations q mix ( w ; θ ) and q 0 ( w ; θ ) that yield the same predictive distribution.

Original: We first writethe density q mix as an expectation of product densities,
Revised: We first write the density q mix as an expectation of product densities,

Original: [Equation]
Revised: [Equation]

Original: Then, we assume that there exists a mapping r ′ =
Revised: Then, we assume that there exists a mapping r ′ =

Original: φ ( r ) such that
Revised: φ ( r ) such that

Original: [Equation]
Revised: [Equation]

Revised: The condition in (7) allows us to exploit the invariance property of the likelihood (approximation) also for each product density q 0 ( t ( w , φ ( r )); θ )

Revised: = p ( t ( w

Original: [Equation]
Revised: [Equation]

Revised: φ

Revised: [Equation]

Revised: [Equation]

Revised: ( t

Original: [Equation]
Revised: [Equation]

Original: The condition in (7) allows us to exploit the invariance property of the likelihood (approximation)also for each product density q 0 ( t ( w, φ ( r )); θ )
Revised: w

Original: =
Revised: [Equation]

Original: p ( t ( w , φ ( r ))) · g 0 ( t ( w , φ ( r )); θ ) , because then
Revised: φ [Equation] r )); θ ) , because then

Original: [Equation]
Revised: [Equation]

Original: where the normalisation constants are Z 0 ( r )
Revised: where the normalisation constants are Z 0 ( r )

Original: = (cid:82) p ( w ) · g 0 ( t ( w , φ ( r )); θ ) d w , and Z mix is definedin (6b).
Revised: = (cid:82) p ( w ) · g 0 ( t ( w , φ ( r )); θ ) d w , and Z mix is defined in (6b).

Original: We show that the condition in (7) holds for translation and permutation invariance in App.
Revised: We show that the condition in (7) holds for translation and permutation invariance in App.

Original: C.
Revised: C.

Original: The second condition is that all invariance transformations t ( · , r ) must be volume-preserving, i.e.
Revised: The second condition is that all invariance transformations t ( · , r ) must be volume-preserving, i.e.

Original: [Equation]
Revised: [Equation]

Original: [Equation] Although (7) and (8) are fairly restricting, common invariances such as translation and permutationinvariance fulfil these conditions for Gaussian priors and likelihood approximations (see App. C).
Revised: Although (7) and (8) are fairly restricting, common invariances such as translation and permutation invariance fulfil these conditions for Gaussian priors and likelihood approximations (see App. C).

Original: With (7) and (8), we can show the posterior predictive equivalence (see Lemma 1 in App. C)
Revised: With (7) and (8), we can show the posterior predictive equivalence (see Lemma 1 in App. C)

Original: [Equation]
Revised: [Equation]

Original: Invariance gap.
Revised: Invariance gap.

Original: If the conditions in (7) and (8) are met, we also show (see Lemma 2 in App.
Revised: If the conditions in (7) and (8) are met, we also show (see Lemma 2 in App.

Original: C)
Revised: C)

Original: [Equation]
Revised: [Equation]

Original: Interestingly, the gap between the two respective ELBO objectives is given exactly by the relativeentropy between the mean-field and invariance-abiding approximation; we therefore refer to it as theinvariance gap .
Revised: Interestingly, the gap between the two respective ELBO objectives is given exactly by the relative entropy between the mean-field and invariance-abiding approximation; we therefore refer to it as the invariance gap .

Original: We make the following observation about the detrimental effect of the invariances:when maximising the standard ELBO L ELBO ( q 0 , D ) instead of the tighter objective L ELBO ( q mix , D )wrt.
Revised: We make the following observation about the detrimental effect of the invariances: when maximising the standard ELBO L ELBO ( q 0 , D ) instead of the tighter objective L ELBO ( q mix , D ) w.r.t.

Original: the parameters of the variational likelihood approximation g 0 ( w ; θ ) (used to construct bothq 0 and q mix ), we see that the former objective favors solutions where q 0 and q mix coincide.
Revised: the parameters of the variational likelihood approximation g 0 ( w ; θ ) (used to construct both q 0 and q mix ), we see that the former objective favors solutions where q 0 and q mix coincide.

Original: Thissuboptimal solution is obtained if the mean-field posterior q 0 ( w ) and thus also the invariance-abidingposterior q mix ( w ) revert to the prior.
Revised: This suboptimal solution is obtained if the mean-field posterior q 0 ( w ) and thus also the invariance-abiding posterior q mix ( w ) revert to the prior.

Original: This is the case if g 0 ( w ) is uniform, because then g mix ( w ) isuniform as well, and because both posteriors are constructed as the product of prior and likelihoodapproximation (cf.
Revised: This is the case if g 0 ( w ) is uniform, because then g mix ( w ) is uniform as well, and because both posteriors are constructed as the product of prior and likelihood approximation (cf.

Original: (6a)).
Revised: (6a)).

Original: Implication of data-related bound.
Revised: Implication of data-related bound.

Original: Since KL [ q 0 || p ] is upper bounded by the best- and worst-case
Revised: Since KL [ q 0 || p ] is upper bounded by the best- and worst-case

Original: ELL as described in Sec. 2.2, the invariance gap is also bounded (using (10) and (3)):
Revised: ELL as described in Sec. 2.2, the invariance gap is also bounded (using (10) and (3)):

Original: [Equation]
Revised: [Equation]

Original: Consequently, this bound puts a constraint on the achievable solutions to the maximisation of the
Revised: Consequently, this bound puts a constraint on the achievable solutions to the maximisation of the

Original: ELBO objective wrt.
Revised: ELBO objective w.r.t.

Original: the variational parameters of the mean-field approximation q 0 ( w ; θ ) .
Revised: the variational parameters of the mean-field approximation q 0 ( w ; θ ) .

Original: Asdiscussed above, one specific parametrisation for which the invariance-gap vanishes is the meanfield posterior collapse, q 0 ( w ; θ )
Revised: As discussed above, one specific parametrisation for which the invariance gap vanishes is the meanfield posterior collapse, q 0 ( w ; θ )

Original: ≈ p ( w ) .
Revised: ≈ p ( w ) .

Original: However, to show that the invariance-gap not onlyadmits posterior collapse as a potential parametrisation but indeed incentivises it, we next tackle thequestion how the invariance-gap behaves for non-collapsed approximations.
Revised: However, to show that the invariance gap not only admits posterior collapse as a potential parametrisation but indeed incentivises it, we next tackle the question how the invariance gap behaves for non-collapsed approximations.

Original: One particularly relevantvariational parametrisation is the optimal solution wrt.
Revised: One particularly relevant variational parametrisation is the optimal solution w.r.t.

Original: the ELBO objective with the invarianceabiding approximation.
Revised: the ELBO objective with the invarianceabiding approximation.

Original: In order to tackle this question, we now consider a simple model for whichthe relevant distributions and the invariance-gap can be computed exactly.
Revised: 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.

Original:  Translation invariance in linear models
Revised:  Translation invariance in linear models

Original: We study an over-parametrised Bayesian linear regression model as the canonical model that exhibitstranslation invariance.
Revised: We study an over-parametrised Bayesian linear regression model as the canonical model that exhibits translation invariance.

Original: The model serves as a useful tool to understand the detrimental effect oftranslation invariance since all interesting quantities can be computed analytically, incl.
Revised: The model serves as a useful tool to understand the detrimental effect of translation invariance since all interesting quantities can be computed analytically, incl.

Original: the mean-fieldand invariance-abiding distribution defined in (6) and the invariance gap from (10).
Revised: the mean-field and invariance-abiding distribution defined in (6) and the invariance gap from (10).

Original: We show how tolift results from this canonical model to the more general case of NNs in Sec. 5.
Revised: We show how to lift results from this canonical model to the more general case of NNs in Sec. 5.

Original: Likelihood model.
Revised: Likelihood model.

Original: Consider a linear model with K latent variables w =
Revised: Consider a linear model with K latent variables w =

Original: [ w 1 , . .
Revised: [ w 1 , . .

Original: . , w K ] T andassume that we have N observations of variables y given dependent inputs x .
Revised: . , w K ] T and assume that we have N observations of variables y given dependent inputs x .

Original: To simplify the setting,we will assume that x = 1 (see App.
Revised: To simplify the setting, we will assume that x = 1 (see App.

Original: D for the more general case).
Revised: D for the more general case).

Original: We further assume that theobservation depends only on the inner product with the inputs and additive Gaussian noise.
Revised: We further assume that the observation depends only on the inner product with the inputs and additive Gaussian noise.

Original: That is,y = 1 K 1 T w + ϵ, ϵ ∼ N (cid:0) 0 , σ 2 y (cid:1) .
Revised: That is, y = 1 K 1 T w + ϵ, ϵ ∼ N (cid:0) 0 , σ 2 y (cid:1) .

Original: Note that any change in w which leaves the sum of the elementsunaffected does not change the likelihood.
Revised: Note that any change in w which leaves the sum of the elements unaffected does not change the likelihood.

Original: In general, we can model this translation invariance usinga K − 1 dimensional vector ∆ ∈ R K − 1 and observing that
Revised: In general, we can model this translation invariance using a K − 1 dimensional vector ∆ ∈ R K − 1 and observing that

Original: [Equation]
Revised: [Equation]

Original: since 1 T B∆ = 0 .
Revised: since 1 T B∆ = 0 .

Original: This over-parametrised model exhibits translation invariance t ( w , ∆) =
Revised: This over-parametrised model exhibits translation invariance t ( w , ∆) =

Original: w − B ∆ .
Revised: w − B ∆ .

Original: Prior.
Revised: Prior.

Original: We assume a Gaussian prior p ( w ) =
Revised: We assume a Gaussian prior p ( w ) =

Original: N ( w ; µ , Σ ) with diagonal covariance Σ =
Revised: N ( w ; µ , Σ ) with diagonal covariance Σ =

Original: Diag (cid:0) σ 2 (cid:1) , where we consider σ 2 = K · σ 20 · 1 proportional to the number of dimension K  such thatthe predictive variance V (cid:2) 1 K (cid:80) k
Revised: Diag (cid:0) σ 2 (cid:1) , where we consider σ 2 = K · σ 20 · 1 proportional to the number of dimensions K , such that the predictive variance V (cid:2) 1 K (cid:80) k

Original: w k + ϵ (cid:3) is constant for different K .
Revised: w k + ϵ (cid:3) is constant w.r.t. K .

Original: Another way to view this is tomodel a prior over parameters that induces the same prior over functions for different K .
Revised: Another way to view this is to model a prior over parameters that induces the same prior over functions for different K .

Original: Posterior.
Revised: Posterior.

Original: The posterior of this Gaussian linear model can be computed as (see App. D.3) 
Revised: The posterior of this Gaussian linear model can be computed as (see App. D.3) p

Original: [Equation]
Revised: [Equation]

Original: As can be seen, the resulting posterior has full covariance with a diagonal plus rank-1 matrix.
Revised: As can be seen, the resulting posterior has full covariance with a diagonal plus rank-1 matrix.

Original: We willnow show that we can construct this posterior from the prior and a mean-field Gaussian approximationof the likelihood by marginalising over all translations ∆ ∼ p (∆) to which the likelihood is invariant.
Revised: We will now show that we can construct this posterior from the prior and a mean-field Gaussian approximation of the likelihood by marginalising over all translations ∆ ∼ p (∆) to which the likelihood is invariant.

Original: 4.1 Mean-field parametrisation of the likelihood function
Revised: 4.1 Mean-field parametrisation of the likelihood function

Original: 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 andvariances λ .
Revised: 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 λ .

Original: We model that the likelihood is translation invariant by computing the marginal from (5)and considering p (∆) =
Revised: We model that the likelihood is translation invariant by computing the marginal from (5) and considering p (∆) =

Original: N (cid:0) ∆ ; 0 , β 2 I (cid:1) with β → ∞ as the uniform distribution over all translations:
Revised: N (cid:0) ∆ ; 0 , β 2 I (cid:1) with β → ∞ as the uniform distribution over all translations:

Original: [Equation]
Revised: [Equation]

Original: Writing the uniform distribution as the limiting case of a Gaussian with infinite variance allows usto compute the integral analytically.
Revised: Writing the uniform distribution as the limiting case of a Gaussian with infinite variance allows us to compute the integral analytically.

Original: As shown in App.
Revised: As shown in App.

Original: D, the resulting function is a multivariate
Revised: D, the resulting function is a multivariate

Original: Gaussian with a degenerate rank-1 covariance matrix and the same location parameter as g 0 ( w ) :
Revised: Gaussian with a degenerate rank-1 covariance matrix and the same location parameter as g 0 ( w ) :

Original: [Equation]
Revised: [Equation]

Original: However, the invariance-abiding posterior is non-degenerate; it can be computed efficiently as
Revised: However, the invariance-abiding posterior is non-degenerate, and it can be computed efficiently as

Original: [Equation]
Revised: [Equation]

Original: As can be seen, this covariance matrix is not diagonal; it has an additive rank-1 term that vanishes asλ → ∞ .
Revised: As can be seen, this covariance matrix is not diagonal; it has an additive rank-1 term that vanishes as λ → ∞ .

Original: Interestingly, the location parameter of q mix is translated from the prior location µ along thedirection of the prior variance vector σ 2 .
Revised: Interestingly, the location parameter of q mix is translated from the prior location µ along the direction of the prior variance vector σ 2 .

Original: This is because g mix is uniform along the K − 1 dimensionshyper-plane determined by its normal vector 1 ; taking the Gaussian product then translates thelocation in the direction Diag ( σ 2 ) 1 = σ 2 .
Revised: This is because g mix is uniform along the K − 1 dimensions hyper-plane determined by its normal vector 1 . Taking the Gaussian product then translates the location in the direction Diag ( σ 2 ) 1 = σ 2 .

Original: Note also that the invariance-abiding posterior can bewritten in the same form as the true posterior in (13) (see App. D.3).
Revised: Note also that the invariance-abiding posterior can be written in the same form as the true posterior in (13) (see App. D.3).

Original: The optimal invariance-abidingposterior is thus the true posterior.
Revised: The optimal invariance-abiding posterior is thus the true posterior.

Original: We visualise the two respective posterior approximations and thecorresponding likelihood in Fig.
Revised: We visualise the two respective posterior approximations and the corresponding likelihood in Fig.

Original: 3 of App.
Revised: 3 of App.

Original: D with two different parametrisations (cf.
Revised: D with two different parametrisations (cf.

Original: Sec. 4.3).
Revised: Sec. 4.3).

Original: 4.2 Invariance-gap
Revised: 4.2 Invariance gap

Original: To quantify the detrimental effect of the translation invariance in the considered linear model we nowcompute the invariance-gap in (11).
Revised: To quantify the detrimental effect of the translation invariance in the considered linear model we now compute the invariance gap in (11).

Original: Note again that we assume a scaled standard normal prior withvariance σ 2 = Kσ 20 · 1 .
Revised: Note again that we assume a scaled standard normal prior with variance σ 2 = Kσ 20 · 1 .

Original: We simplify the form of the KL divergence by assuming that all variancesand means take the same value in the likelihood (and thus also in the posterior) approximation, i.e.
Revised: We simplify the form of 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.e.

Original: [Equation]
Revised: [Equation]

Original: This choice is motivated by the fact that the posterior approximation is a function of the sum 1 T λonly (cf.
Revised: This choice is motivated by the fact that the posterior approximation is a function of the sum 1 T λ only (cf.

Original: (15)), and, hence, it makes no difference for q mix .
Revised: (15)), and, hence, it makes no difference for q mix .

Original: However, since the prior variance isproportional to 1 , the Gaussian likelihood resulting in the highest ELBO for the approximation q 0also has a variance vector proportional to 1 .
Revised: However, since the prior variance is proportional to 1 , the Gaussian likelihood resulting in the highest ELBO for the approximation qalso has a variance vector proportional to 1 .

Original: Using this assumption, the invariance gap is
Revised: Using this assumption, the invariance gap is

Original: [Equation]
Revised: [Equation]

Original: This is a convex function in ϕ =
Revised: This is a convex function in ϕ =

Original: ˆ σ 2 +ˆ λ ˆ λ with a minimum at ϕ = 

Revised: ˆ σ 2 +ˆ λ ˆ λ with a minimum at ϕ =

Original: ; it is minimised as ˆ λ → ∞ .
Revised: ; it is minimised as ˆ λ → ∞ .

Original: It ishowever not evident whether the gap has a large magnitude compared to the rest of the regularisationterm and over-regularises in practice.
Revised: It is however not evident whether the gap has a large magnitude compared to the rest of the regularisation term and over-regularises in practice.

Original: In the next section, we therefore analyse this term at the optimafor the mean-field and the invariance-abiding approximations defined in (17).
Revised: In the next section, we therefore analyse this term at the optima for the mean-field and the invariance-abiding approximations defined in (17).

Original: The correspondinginvariance gaps are visualised in Fig.
Revised: The corresponding invariance gaps are visualised in Fig.

Original: 1 where it can be seen that the invariance-gap grows linearlywhen using the optimal parameters of the invariance-abiding distribution and the unattainable datarelated bound is reached quickly, thus preventing this optimum if (17b) is optimised instead of (17a).
Revised: 1 where it can be seen that the invariance gap grows linearly when using the optimal parameters of the invariance-abiding distribution and the unattainable datarelated bound is reached quickly, thus preventing this optimum if (17b) is optimised instead of (17a).

Original: 4.3 Optimal mean-field and invariance-abiding parametrisations
Revised: 4.3 Optimal mean-field and invariance-abiding parametrisations

Original: To better understand the detrimental effect of the invariance gap, we now analyse the ELL and KLterms of the ELBO objective for the mean-field and invariance-abiding variational approximations,respectively.
Revised: To better understand the detrimental effect of the invariance gap, we now analyse the ELL and KL terms of the ELBO objective for the mean-field and invariance-abiding variational approximations, respectively.

Original: We compare these terms for both distributions with the parameters optimized forboth posterior distributions, giving 2 × 2 combinations.
Revised: We compare these terms for both distributions with the parameters optimized for both posterior distributions, giving 2 × 2 combinations.

Original: We denote the optimal parameters wrt.
Revised: We denote the optimal parameters w.r.t.

Original: theinvariance-abiding approximation and the mean-field approximation, respectively, as
Revised: the invariance-abiding approximation and the mean-field approximation, respectively, as

Original: [Equation]
Revised: [Equation]

Original: Perhaps the simplest way compute these optimal parameters is to notice that q mix in (15) can bewritten in the same form as the true posterior and then simply read out the optimal parameters.
Revised: Perhaps the simplest way compute these optimal parameters is to notice that q mix in (15) can be written in the same form as the true posterior and then simply read out the optimal parameters.

Original: Thisis shown in App.
Revised: This is shown in App.

Original: D.3; the resulting optimal variational parameters are
Revised: D.3; the resulting optimal variational parameters are

Original: [Equation]
Revised: [Equation]

Original: [Figure]
Revised: [Figure]

Original: The optimal parameters for the mean-field posterior can also be computed analytically since the
Revised: The optimal parameters for the mean-field posterior can also be computed analytically, since the

Original: Gaussian mean-field distribution that minimises the KL divergence to the multivariate Gaussian trueposterior is known [31].
Revised: Gaussian mean-field distribution that minimises the KL divergence to the multivariate Gaussian true posterior is known [32].

Original: The resulting optimal parameters of the mean-field likelihood are
Revised: The resulting optimal parameters of the mean-field likelihood are

Original: [Equation]
Revised: [Equation]

Original: As can be seen, the two respective optimal parameters differ by the factor K in the likelihood variance.
Revised: As can be seen, the two respective optimal parameters differ by the factor K in the likelihood variance.

Original: We then construct the two respective posterior approximations q 0 and q mix with these two optimalparameters and compare the 2 × 2 combinations in Fig.
Revised: We then construct the two respective posterior approximations q 0 and q mix with these two optimal parameters and compare the 2 × 2 combinations in Fig.

Original: 2 where we visualise the ELBO terms.
Revised: 2, where we visualise the ELBO terms.

Original: Note again that we model prior variances σ 2 = Kσ 20 · 1 such that prior and posterior over functionsare identical for any K .
Revised: Note again that we model prior variances σ 2 = Kσ 20 · 1 such that prior and posterior over functions are identical for any K .

Original: Due to this choice and since q mix ( w ; θ ∗ mix ) approximates the true posteriorexactly it does not suffer from over-parametrisation.
Revised: Due to this choice and since q mix ( w ; θ ∗ mix ) approximates the true posterior exactly, it does not suffer from over-parametrisation.

Original: This can be seen by the loss terms in Fig.
Revised: This can be seen by the loss terms in Fig.

Original: 2 beingconstant in K .
Revised: 2 being constant in K .

Original: In contrast, since λ ∗ 0 depends quadratically on K and σ 2 only linearly, q mix ( w ; θ ∗ 0 )collapses to the prior as K → ∞ .
Revised: In contrast, since λ ∗ 0 depends quadratically on K and σ 2 only linearly, q mix ( w ; θ ∗ 0 ) collapses to the prior as K → ∞ .

Original: This can be seen e.g. by the shrinking KL regularisation (Fig. 2a)and ELL (in 2b) term.
Revised: This can be seen e.g. by the shrinking KL regularisation (Fig. 2a) and ELL (in 2b) term.

Original: As a consequence, and in line with Coker et al.
Revised: As a consequence, and in line with Coker et al.

Original: [4], the posterior predictivevariance of the optimal mean-field approximation reverts to the prior predictive variance (Fig. 2d).
Revised: [4], the posterior predictive variance of the optimal mean-field approximation reverts to the prior predictive variance (Fig. 2d).

Original: We have shown that—in the simplified setting of a linear model with dependent observations—theconsequence of not handling translation invariance is indeed posterior collapse as K → ∞ , whilemodelling the invariance coincides with the true posterior.
Revised: We have shown that—in the simplified setting of a linear model with dependent observations—the consequence of not handling translation invariance is indeed posterior collapse as K → ∞ , while modelling the invariance coincides with the true posterior.

Original: An important observation and key takeawayis that while we need to optimise for (17a), the mean-field posterior approximation q 0 ( w ; θ ∗ mix ) issufficient for prediction.
Revised: An important observation and key takeaway is that, while we need to optimise for (17a), the mean-field posterior approximation q 0 ( w ; θ ∗ mix ) is sufficient for prediction.

Original: Indeed, for the linear model, we could compute the invariance-gap exactlyand thereby correct the ELBO objective for the invariances of this model class.
Revised: Indeed, for the linear model, we could compute the invariance gap exactly and thereby correct the ELBO objective for the invariances of this model.

Original: For more complex nonlinear models such as BNNs, this section may serve as a direction for approximating the invariancegap.
Revised: For more complex nonlinear models such as BNNs, this section may serve as a direction for approximating the invariance gap.

Original: To this end, the next section describes the layer-wise translation invariance exhibited by BNNs.
Revised: To this end, the next section describes the layer-wise translation invariance exhibited by BNNs.

Original:  Translation invariance in Bayesian neural networks
Revised:  Translation invariance in Bayesian neural networks

Original: Let us now go back to the general NN model in Sec. 2.1.
Revised: Let us now go back to the general NN model in Sec. 2.1.

Original: In order to describe the set of all invariances,note that for each node z l,j in layer l (or the output node y ), there is a subspace that keeps the valuez l,j (or y ) invariant when using the translation-invariance model in (12) because the value itself onlydepends on the actual activation z l − 1 (or the input x ).
Revised: In order to describe the set of all invariances, note that for each node z l,j in layer l (or the output node y ), there is a subspace that keeps the value z l,j (or y ) invariant when using the translation-invariance model in (12), because the value itself only depends on the actual activation z l − 1 (or the input x ).

Original: More formally, the following equations model all translation invariances of a NN at a data point x :
Revised: More formally, the following equations model all translation invariances of a NN at a data point x :

Original: [Equation]
Revised: [Equation]

Original: [Figure]
Revised: [Figure]

Original: Prior variance is σ 2 = K · 1 , and σ 2 y = 1 / (2 π e) ; we used N = 10 identical observations with value y = 1 and inputs x = 1 .
Revised: Prior variance is σ 2 = K · 1 , and σ 2 y = 1 / (2 π e) ; we used N = 10 identical observations with value y = 1 and inputs x = 1 .

Original: Notably, both posterior approximations yield the same predictive distribution as can be seen by the ELL (b) and predictive variance (d).
Revised: Notably, both posterior approximations yield the same predictive distribution as can be seen by the ELL (b) and predictive variance (d).

Original: The variance of the mean-field approximation reverts to the prior variance as K → ∞ .
Revised: The variance of the mean-field approximation reverts to the prior variance as K → ∞ .

Original: where the layer index l ranges from 2 , . .
Revised: where the layer index l ranges from 2 , . .

Original: , L − 1 and B z is given by
Revised: , L − 1 and B z is given by

Original: [Equation]
Revised: [Equation]

Original: In order to efficiently compute the invariance-abiding likelihood g mix , note that it also decomposes258 over the NN layers and the associated parameters as follows: For the n 1 nodes in the first layer we have
Revised: In order to efficiently compute the invariance-abiding likelihood g mix , note that it also decomposes over the NN layers and the associated parameters as follows: For the n 1 nodes in the first layer we have

Original: [Equation]
Revised: [Equation]

Original: where the outer expectation is taken over x ∼ p ( D ) and p ( D ) is the empirical distribution over thetraining set D , and where m , V , µ and Σ are constrained to the parts of the overall weight vector that261 correspond to w 1 ,j .
Revised: where the outer expectation is taken over x ∼ p ( D ) and p ( D ) is the empirical distribution over the training set D , and where m , V , µ and Σ are constrained to the parts of the overall weight vector that correspond to w 1 ,j .

Original: Now, since the invariance-abiding mechanism for nodes in layer l only depends262 on the value
Revised: Now, since the invariance-abiding mechanism for nodes in layer l only depends on the value

Original: z l − 1 of the hidden units in layer n − 1 (see (20) and (21)), we have
Revised: z l − 1 of the hidden units in layer n − 1 (see (20) and (21)), we have

Original: [Equation]
Revised: [Equation]

Original: where again m , V , µ and Σ are constrained to the parts of the overall weight vector that correspondto w l,j and the expectation over z is taken with respect to z ∼ P ( z l − 1 ) .
Revised: where again m , V , µ and Σ are constrained to the parts of the overall weight vector that correspond to w l,j and the expectation over z is taken with respect to z ∼ P ( z l − 1 ) .

Original: Thus, the invariance-abiding approximation could e.g. be computed with a layer-by-layer iterativeoptimization: First, given data x ∼ p ( D ) and prior p ( w ) , use (23) and point estimates for all weightvectors in later layers to optimize the mean m and diagonal covariance V of the weights in the firstlayer, i.e., { w 1 ,j | j = 1 , . . . , n 1 } .
Revised: Thus, the invariance-abiding 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., { w 1 ,j | j = 1 , . . . , n 1 } .

Original: Then, we can use (24) to generate P samples z 1 ,n of the activationsof the first layer under the (now fitted) optimal q mix ( { w 1 ,j } ) .
Revised: Then, we can use (24) to generate P samples z 1 ,n of the activations of the first layer under the (now fitted) optimal q mix ( { w 1 ,j } ) .

Original: Next, for each of these samples z 1 ,nwe can use (25) and point estimates for all weight vectors in later layers to optimize the mean mand diagonal covariance V of the weights in the second layer and average them for the optimal q mixof the weights in the second layer.
Revised: Next, for each of these samples z 1 ,n , we can use (25) and point estimates for all weight vectors in later layers to optimise the mean m and diagonal covariance V of the weights in the second layer and average them for the optimal q mix of the weights in the second layer.

Original: Finally, we can use (26) to generate samples z 2 ,i for the second layer.
Revised: Finally, we can use (26) to generate samples z 2 ,i for the second layer.

Original: This process is repeated until the last layer andthe whole procedure can be iterated until theq mix ( w ) does not change in one optimization loop.
Revised: This process is repeated until the last layer, and iterated until q mix ( w ) converges.

Original: The complexity of this procedure would be nolarger than the optimization of the weight matrices of a single layer because all previous layers arefully characterized by the distribution of the latent activations z l of the hidden layers.
Revised: The complexity of this procedure is 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 z l of the hidden layers.

Original:  Related work
Revised:  Related work

Original: Poor empirical performance of VBNNs has been reported in several previous works, e.g., [34, 15, 30].
Revised: Poor empirical performance of VBNNs has been reported in several previous works, e.g., [35, 15, 31].

Original: For instance, it has been shown that single-layer mean-field VBNNs with ReLU activations cannot have large predictive uncertainty between regions of low uncertainty [9].
Revised: For instance, it has been shown that single-layer mean-field VBNNs with ReLU activations can not have large predictive uncertainty between regions of low uncertainty [9].

Original: In contrast, Farquharet al.
Revised: In contrast, Farquhar et al.

Original: [7] argue that mean-field approximations are expressive enough for deep BNNs; they provea universality result that the predictive distribution of mean-field approximations of BNNs with atleast 2 layers of hidden units can approximate any true posterior distribution over function valuesarbitrarily closely.
Revised: [7] 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 can approximate any true posterior distribution over function values arbitrarily closely.

Original: This result seems in conflict with our work and [4] who show that the predictivedistribution of mean-field VBNNs reverts to the prior predictive distribution as the network widthincreases.
Revised: This result seems in conflict with our work and [4], who show that the predictive distribution of mean-field VBNNs reverts to the prior predictive distribution as the network width increases.

Original: The discrepancy between these two results is resolved by noting that Farquhar et al.
Revised: The discrepancy between these two results is resolved by noting that Farquhar et al.

Original: [7]only shows that the expressive power of the model is large enough but not that the approximateinference algorithms will converge to this solution.
Revised: [7] only shows that the expressive power of the model is large enough but not that the approximate inference algorithms will converge to this solution.

Original: Our work sheds light on why the mean-fieldapproximation fails to approximate expressive posteriors via the invariance-gap in the ELBO.
Revised: Our work sheds light on why the mean-field approximation fails to approximate expressive posteriors via the invariance gap in the ELBO.

Revised: Our framework to model the invariances is similar to and extends previous preliminary works [18, 21].

Original: Other attempts have proposed to approximate the posterior predictive directly, thereby circumventingthe non-identifiability issue [25, 19, 33].
Revised: Other attempts have proposed to approximate the posterior predictive directly, thereby circumventing the non-identifiability issue [26, 19, 34].

Original: While these approaches have shown promising empiricalperformance, estimating the KL regularisation term in function space is more complicated and caneven be ill-defined due to regions of zero prior probability mass [3].
Revised: 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 to regions of zero prior probability mass [3].

Original: In a similar vein, previous workalso proposed to map the BNN prior to a Gaussian process (functional) prior [8, 29].
Revised: In a similar vein, previous work also proposed to map the BNN prior to a Gaussian process (functional) prior [8, 30].

Original: Surprisingly, restricting the parametrisation of the variational posterior results in competitive oreven better performance [27, 20, 28, 6].
Revised: Surprisingly, restricting the parametrisation of the variational posterior results in competitive or even better performance [28, 20, 29, 6].

Original: For example, Dusenberry et al.
Revised: For example, Dusenberry et al.

Original: [6] proposes a variationalapproximation only for rank-1 factors, resulting in inference of a lower-dimensional subspace; theouter product of these low-rank factors perturb a weight matrix that is treated deterministicallythrough maximum a posteriori (MAP) estimation.
Revised: [6] proposes a variational approximation only for rank-1 factors, resulting in inference of a 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.

Original: This approach avoids the invariance problemassociated with variational Bayesian inference since the MAP estimate is not impeded by this problemand the subspace of the variational weights do not possess the same invariances.
Revised: 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.

Original: More broadly, (non-)identifiability of parameters and latent variables in probabilistic models is awidely-studied topic from frequentist [24] and Bayesian perspectives [5, 22, 10].
Revised: More broadly, (non-)identifiability of parameters and latent variables in probabilistic models is a widely-studied topic from frequentist [25] and Bayesian perspectives [5, 23, 10].

Original: Recently, Wang et al.[32] attributed posterior collapse in variational auto-encoders to non-identifiability of latent variables.
Revised: Recently, Wang et al. [33] attributed posterior collapse in variational auto-encoders to non-identifiability of latent variables.

Original:  Conclusion
Revised:  Conclusion

Original: We have associated the posterior collapse phenomenon of mean-field VBNNs with invariances inthe likelihood function, in particular translation invariance.
Revised: We have associated the posterior collapse phenomenon of mean-field VBNNs with invariances in the likelihood function, in particular translation invariance.

Original: While the invariance does not affectthe predictive distribution, the approximations of the posterior which abide and do not abide theinvariance differ in the KL regularisation term and consequently in the tightness of the ELBOobjective.
Revised: While the invariance does not affect the predictive distribution, the approximations of the posterior—which abide and do not abide the invariance—differ in the KL regularisation term and consequently in the tightness of the ELBO objective.

Original: We proved that the objectives of the two approximations differ by the relative entropy (KL)between the standard mean-field approximation and the invariance-abiding distribution.
Revised: We proved that the objectives of the two approximations differ by the relative entropy (KL) between the standard mean-field approximation and the invariance-abiding distribution.

Original: We relatedthis to a data-related bound on the KL regularisation which prevents fits for which the gap is large.
Revised: We related this to a data-related bound on the KL regularisation, which prevents fits for which the gap is large.

Original: We studied over-parametrised Bayesian linear regression as the canonical model that exhibitstranslation invariance and for which the relevant terms can be computed exactly.
Revised: We studied over-parametrised Bayesian linear regression as the canonical model that exhibits translation invariance and for which the relevant terms can be computed exactly.

Original: A detailed analysisof this model confirms our hypothesis that the invariance leads to a significant additional gap in the
Revised: A detailed analysis of this model confirms our hypothesis that the invariance leads to a significant additional gap in the

Original: ELBO objective compared to approximations that model the invariance.
Revised: ELBO objective compared to approximations that model the invariance.

Original: It is this very gap whichprevents mean-field approximation to achieve the same fit as an invariance-abiding approximationand instead leads to a collapse of the posterior variances to the prior variance.
Revised: It is this very gap which prevents mean-field approximation to achieve the same fit as an invariance-abiding approximation and instead leads to a collapse of the posterior variances to the prior variance.

Original: While we can compute the invariance gap for the over-parametrised linear model, we have not yetidentified a computationally efficient procedure for layered models or for other invariances.
Revised: While we can compute the invariance gap for the over-parametrised linear model, we have not yet identified a computationally efficient procedure for layered models or for other invariances.

Original: However,our work provides the mathematical tools to address over-regularisation due to invariances.
Revised: However, our work provides the mathematical tools to address over-regularisation due to invariances.

Original: Futurework will focus on approximations of the invariance gap in order to correct the ELBO objective fortranslation and permutation invariances in general neural network functions.Checklist For all authors...
Revised: Future work will focus on approximations of the invariance gap in order to correct the ELBO objective for translation and permutation invariances in general neural network functions.

Original: (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s

Original: contributions and scope?

Original: [Yes]

Original: (b) Did you describe the limitations of your work?

Original: [Yes] We point out limitations and

Original: future work throughout the paper.

Original: (c) Did you discuss any potential negative societal impacts of your work?

Original: [No] This ispurely theoretical work without direct negative societal impacts.

Original: (d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?

Original: [Yes]

Original: If you are including theoretical results...

Original: (a) Did you state the full set of assumptions of all theoretical results?

Original: [Yes]

Original: (b) Did you include complete proofs of all theoretical results?

Original: [Yes]

Original: If you ran experiments...

Original: [No] No experiments

Original: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...

Original: [No] No existing assets used.

Original: If you used crowdsourcing or conducted research with human subjects...

Original: [No] Nocrowdsourcing/human subjects.