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)
Revised: Anonymous Author(s)

Original: AffiliationAddress email
Revised: 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 approximationssuch 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 hasassociated the variational mean-field approximation for Bayesian neural networkswith 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 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.

Original: In particular, we showthat the mean-field approximation has a gap in the evidence lower bound objectivecompared to a purpose-built posterior that takes into account the known invariances.
Revised: In particular, we showthat the mean-field approximation has a gap in the evidence lower bound objectivecompared to a purpose-built posterior that takes into account the known invariances.

Original: Importantly, the invariance gap is not constant; it vanishes as the approximationreverts to the prior.
Revised: Importantly, the invariance gap is not constant; it vanishes as the approximationreverts 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 ina 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 trueposterior can be constructed from a mean-field parametrisation, this is achievedonly if the objective function takes into account the invariance gap.

Original: Then, wetransfer our analysis of the linear model to neural networks.
Revised: Then, wetransfer 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 providesa 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]
Revised: Bayesian neural networks (BNNs) have several appealing advantages compared to deterministicneural networks (NN) such as improving generalization [35], capturing epistemic uncertainty [16]

Original: and providing a framework for continual learning methods [17, 21].
Revised: and providing a framework for continual learning methods [17, 21].

Original: Unfortunately, reaping thesetheoretical benefits has thus far been impeded [34].
Revised: Unfortunately, reaping thesetheoretical benefits has thus far been impeded [34].

Original: In particular, several practical issues withvariational Bayesian inference methods—which are the de-facto standard technique for scalinginference in BNNs to large datasets [14, 1]—have been identified to be partially responsible for thisperformance gap [15].
Revised: In particular, several practical issues withvariational Bayesian inference methods—which are the de-facto standard technique for scalinginference in BNNs to large datasets [14, 1]—have been identified to be partially responsible for thisperformance gap [15].

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

Original: It has beenobserved that mean-field variational BNNs (VBNNs) suffer from severe underfitting for large modelsand small dataset sizes [12].
Revised: It has beenobserved that mean-field variational BNNs (VBNNs) suffer from severe underfitting for large modelsand small dataset sizes [12].

Original: Recent work [4] showed that under certain assumptions such as an odd
Revised: Recent work [4] showed that under certain assumptions such as an odd

Original: Lipschitz activation function and a finite dataset with bounded likelihood, the optimal mean-fieldapproximation collapses to the prior as the NN width increases, ignoring the data .
Revised: Lipschitz activation function and a finite dataset with bounded likelihood, the optimal mean-fieldapproximation collapses to the prior as the NN width increases, ignoring the data .

Original: In this work, we study the invariances of the likelihood function and the resulting non-identifiability inoverparametrised models.
Revised: In this work, we study the invariances of the likelihood function and the resulting non-identifiability inoverparametrised models.

Original: We identify it as the primary reason why mean-field variational distributionscannot approximate the posterior, ultimately leading to the reversion to the prior analysed in [4].
Revised: We identify it as the primary reason why mean-field variational distributionscannot approximate the posterior, ultimately leading to the reversion to the prior analysed in [4].

Original: It is well known that NNs exhibit several invariances with respect to their parameters , includingnode permutation invariance and sign-flip invariance (in the case of odd activation functions) [26, 2],and scaling invariance in the case of piece-wise linear activations [23].
Revised: It is well known that NNs exhibit several invariances with respect to their parameters , includingnode permutation invariance and sign-flip invariance (in the case of odd activation functions) [26, 2],and scaling invariance in the case of piece-wise linear activations [23].

Original: Recent work showed thatthe permutation invariance introduces an additional gap in the evidence lower bound (ELBO) for a mean-field approximation compared to an (intractable) mixture that includes all modes [18].
Revised: Recent work showed thatthe permutation invariance introduces an additional gap in the evidence lower bound (ELBO) for a mean-field approximation compared to an (intractable) mixture that includes all modes [18].

Original: Weshow that the parameter space of BNNs additionally has subspaces with translation invariance .
Revised: Weshow that 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 likelihoodpoint estimation through stochastic gradient descent is used, because convergence to any of theequivalent optima (resulting from the invariance) suffices for prediction.

Original: However, we showthat the variational Bayesian approach is particularly impaired by invariances in the likelihood,since the resulting posterior has a complicated structure that can not be well approximated by themean-field approximation.
Revised: However, we showthat the variational Bayesian approach is particularly impaired by invariances in the likelihood,since the resulting posterior has a complicated structure that can not be well approximated by themean-field approximation.

Original: In order to isolate the impact of the invariances, we define a posteriorapproximation in Sec. 3 that explicitly models the invariances by integrating over all transformationsthat leave the likelihood invariant.
Revised: In order to isolate the impact of the invariances, we define a posteriorapproximation in Sec. 3 that explicitly models the invariances by integrating over all transformationsthat leave the likelihood invariant.

Original: We construct this invariance-abiding distribution from a meanfield likelihood approximation such that it induces an identical posterior predictive distributioncompared to an alternative mean-field posterior approximation that is agnostic to the invariances.
Revised: We construct this invariance-abiding distribution from a meanfield likelihood approximation such that it induces an identical posterior predictive distributioncompared to an alternative mean-field posterior approximation that is agnostic to the invariances.

Original: Both approximations are parametrised with the same variational parameters and result in the sameexpected log-likelihood, but differ in their modelling of the invariances, and, hence, the corresponding
Revised: Both approximations are parametrised with the same variational parameters and result in the sameexpected log-likelihood, but differ in their modelling of the invariances, and, hence, the corresponding

Original: ELBO objective differs only in the KL regularisation term.
Revised: ELBO objective differs only in the KL regularisation term.

Original: We refer to the difference between thesetwo KL terms as the invariance-gap .
Revised: We refer to the difference between thesetwo KL terms as the invariance-gap .

Original: While our constructed approximation is intractable in general, itserves as a useful tool for studying the invariance in simple models where it is indeed tractable, and itprovides a framework for future work to devise practical remedies.
Revised: While our constructed approximation is intractable in general, itserves as a useful tool for studying the invariance in simple models where it is indeed tractable, and itprovides a framework for future work to devise practical remedies.

Original: The main contribution of this work is a detailed analysis (Sec. 4) of a linear model that implementstranslation invariance and explains the reversion of the posterior to the prior as the result of optimizingthe ELBO that exhibits a non-trivial invariance-gap.
Revised: The main contribution of this work is a detailed analysis (Sec. 4) of a linear model that implementstranslation invariance and explains the reversion of the posterior to the prior as the result of optimizingthe ELBO that exhibits a non-trivial invariance-gap.

Original: This gap vanishes for approximations thatcollapse to the prior and grows linearly in the number of dimensions with translation invariance formean-field approximations that have the same posterior predictive distribution as the true posterior.
Revised: This gap vanishes for approximations thatcollapse to the prior and grows linearly in the number of dimensions with translation invariance formean-field approximations that have the same posterior predictive distribution as the true posterior.

Original: We transfer our analysis of the linear model to VBNNs in Sec. 5, showing that subspaces in BNNsexist that exhibit translation invariance.
Revised: We transfer our analysis of the linear model to VBNNs in Sec. 5, showing that subspaces in BNNsexist that exhibit translation invariance.

Original: We start by introducing the basic concepts and recent resultson which we build our contribution in Sec. 2.
Revised: We start by introducing the basic concepts and recent resultson which we build our contribution in Sec. 2.

Original:  Background2.1 Variational Bayesian Neural Networks
Revised:  Background2.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 transformtheir 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 which are each a weighted linear combination of either theinput vector x (first hidden layer) or the value of the hidden units z l (all other hidden layers and theoutput f ( x ) ).
Revised: Such a networkhas n 1 + n 2 + · · · + n L many nodes z l,i which are each a weighted linear combination of either theinput vector x (first hidden layer) or the value of the hidden units z l (all other hidden layers and theoutput f ( x ) ).

Original: We use w =
Revised: We use w =

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

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

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

Original: If the weights and biases w are treated as random variables witha prior distribution p ( w ) , then the posterior p ( w | D ) (induced by the dataset D through the likelihoodfunction ℓ ( w ; D ) :
Revised: If the weights and biases w are treated as random variables witha prior distribution p ( w ) , then the posterior p ( w | D ) (induced by the dataset D through the likelihoodfunction ℓ ( w ; D ) :

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

Original: The variational Bayesian method approximates the posterior by a distribution q θ ( 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: (1) where KL [ q θ || p ( · | D )] :=
Revised: (1) where KL [ q θ || p ( · | D )] :=

Original: E w ∼
Revised: E w ∼

Original: [Equation]
Revised: [Equation]

Original: θ [ln
Revised: θ [ln

Original: [Equation]
Revised: [Equation]

Original: θ ( w ) − ln p ( w |D )] and Q is a family of distributions over w .
Revised: θ ( 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 minimizerof (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 aconsequence 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 resulthas been shown in [4] and we recall the result for a Gaussian likelihood with homogeneous noise (forother 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 , σ y ) ,an isotropic Gaussian prior p ( w ) =
Revised: f w ( x )+ ϵ , with ϵ ∼ N (0 , σ 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 σ L (cid:0) x ( n ) (cid:1) :
Revised: Define the NN output variance σ 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 some fixedinput x ( n ) , where L indicates the last layer.
Revised: ( x ( n ) ; w ) (cid:3) for some fixedinput 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 σ y(see App.
Revised: where q ∗ is a hypothetical approximation that predicts the data perfectly up to the noise variance σ 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 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.

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: 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.

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 understandingtheir 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 Gaussianfunction with variational parameters and marginalise over all transformations to which the truelikelihood 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-abidingposterior approximation q mix which can be related to a mean-field approximation q 0 that does notmodel 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 identicalposterior 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 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].

Original: We model this invariance throughtransformations t ( · , r ) to which the likelihood ℓ ( w ; D ) =
Revised: We model this invariance throughtransformations 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 wrt.

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 sameprobability a priori.

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

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

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 thesetransformations 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: Forinstance, permutation invariance induces a factorial number of discontinuous modes that are eachequivalent (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 andposterior of an over-parametrised Bayesian linear regression model can be constructed from a productof 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 accountthe 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 approximationg 0 ( w ; θ ) ≈ ℓ 0 ( w , D ) , with variational parameters θ .

Original: From this single mode approximation, wemodel an invariance-abiding likelihood approximation from the same parametrisation through
Revised: From this single mode approximation, wemodel an invariance-abiding likelihood approximation from 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, theinvariance-abiding approximation q mix ( w ; θ ) is a mixture with infinite continuous or finite discretemodes 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 constructapproximations 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 writethe 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]

Original: This assumption allows us to exploit the invariance property of the likelihood (approximation) alsofor each product density q 0 ( t ( w , φ ( r )); θ )
Revised: This assumption allows us to exploit the invariance property of the likelihood (approximation) alsofor each product density q 0 ( t ( w , φ ( r )); θ )

Original: =
Revised: =

Original: p ( t ( w , φ ( r ))) · g 0 ( t
Revised: p ( t ( w , φ ( r ))) · g 0 ( t

Original: [Equation]
Revised: [Equation]

Original: w
Revised: w

Original: [Equation]
Revised: [Equation]

Original: φ ( r )); θ ) , because then
Revised: φ ( 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 definedin (6b).

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

Original: App.
Revised: App.

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

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: Although (7) and (8) are fairly restricting, common invariances such as translation and permutationinvariance fulfil this condition for Gaussian priors and likelihood approximations (see App. C).
Revised: Although (7) and (8) are fairly restricting, common invariances such as translation and permutationinvariance fulfil this condition 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] [Equation]
Revised: [Equation] [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 relativeentropy between the mean-field and invariance-abiding approximation; we therefore refer to it as theinvariance 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 )wrt.

Original: the parameters of the variational likelihood approximation g 0 ( w ; θ ) (which is used to constructboth q 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 ; θ ) (which is used to constructboth q 0 and q mix ), we see that the former objective favors solutions where q 0 and q mix coincide.

Original: Onesuch suboptimal solution is when the likelihood is uninformative (uniform), and, thus, both q 0 andq mix revert to the prior.
Revised: Onesuch suboptimal solution is when the likelihood is uninformative (uniform), and, thus, both q 0 andq mix revert to the prior.

Original: In Sec. 4, we study a simple model that exhibits translation invariance andshow that this additional gap is significant.
Revised: In Sec. 4, we study a simple model that exhibits translation invariance andshow that this additional gap is significant.

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 ofthe ELBO objective wrt.
Revised: Consequently, this bound puts a constraint on the achievable solutions to the maximisation ofthe ELBO objective wrt.

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: One particular parametrization that guarantees this behaviour is the mean-field posterior collapse,q 0 ( w ; θ )
Revised: One particular parametrization that guarantees this behaviour is the mean-field posterior collapse,q 0 ( w ; θ )

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

Original: However, in order to show that the invariance-gap not only admits posterior collapse as a potential parametrization but indeed incentivises posterior collapse, wenext tackle the question how the invariance-gap grows for non-collapsed approximations.
Revised: However, in order to show that the invariance-gap not only admits posterior collapse as a potential parametrization but indeed incentivises posterior collapse, wenext tackle the question how the invariance-gap grows for non-collapsed approximations.

Original: We observe in the subsequent sections that the invariance-gap increases with the model sizefor non-collapsed approximations and the invariances we consider.
Revised: We observe in the subsequent sections that the invariance-gap increases with the model sizefor non-collapsed approximations and the invariances we consider.

Original: Examples of the invarianceinclude permutation invariance (interchangeability of neurons in the same layer of a BNN) in thesupplementary material.
Revised: Examples of the invarianceinclude permutation invariance (interchangeability of neurons in the same layer of a BNN) in thesupplementary material.

Original: We show that KL [ q 0 || q mix ] is between zero and the logarithm of thenumber of discrete permutation-induced modes, which grows factorial with the number of neuronsper layer.
Revised: We show that KL [ q 0 || q mix ] is between zero and the logarithm of thenumber of discrete permutation-induced modes, which grows factorial with the number of neuronsper layer.

Original: Another example and the focus of the main body of this work is continuous translationinvariance.
Revised: Another example and the focus of the main body of this work is continuous translationinvariance.

Original: In Sec. 4, we analyse translation invariance in over-parametrised linear models, show thatthe gap grows linear with the number of dimensions (weights) for parameter values θ ∗ mix that areoptimal in terms of the posterior predictive and the ELBO optimized wrt.
Revised: In Sec. 4, we analyse translation invariance in over-parametrised linear models, show thatthe gap grows linear with the number of dimensions (weights) for parameter values θ ∗ mix that areoptimal in terms of the posterior predictive and the ELBO optimized wrt.

Original: q mix ( w , θ ) .
Revised: q mix ( w , θ ) .

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 exhibitstranslation 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 oftranslation invariance, since all interesting quantities can be computed analytically, incl.

Original: the meanfield and invariance-abiding distribution defined in (6), and the invariance gap from (10).
Revised: the meanfield and invariance-abiding distribution defined in (6), and the invariance gap from (10).

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

Original: We start byspecifiying the canonical model in the following.
Revised: We start byspecifiying the canonical model in the following.

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 andassume 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 in the supplementary material for the more general case).
Revised: D in the supplementary material for the more general case).

Original: Wefurther assume that the observation depends only on the inner product with the inputs and additive
Revised: Wefurther assume that the observation depends only on the inner product with the inputs and additive

Original: Gaussian noise.
Revised: Gaussian noise.

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

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

Original: In general, we can model thistranslation invariance using a K − 1 dimensional vector ∆ ∈ R K − 1 and observing that
Revised: In general, we can model thistranslation 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 ( σ ) ,where we consider σ = K · σ 0 · 1 proportional to the number of dimension K such that the predictivevariance V (cid:2) 1 K (cid:80)
Revised: Diag ( σ ) ,where we consider σ = K · σ 0 · 1 proportional to the number of dimension K such that the predictivevariance V (cid:2) 1 K (cid:80)

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

Original: Another way to view this is to model a priorover parameters that induces the same prior over functions for different K .
Revised: Another way to view this is to model a priorover 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)

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 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.

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 andvariances λ .

Original: We model that the likelihood is translation invariant by computing the marginal from (5)and consider p (∆) =
Revised: We model that the likelihood is translation invariant by computing the marginal from (5)and consider 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 allowsus compute the integral analytically.
Revised: Writing the uniform distribution as the limiting case of a Gaussian with infinite variance allowsus 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; it can be computed efficiently as

Original: [Equation]
Revised: [Equation]

Original: As can be seen, this covariance matrix is not diagonal, but has an additive rank-1 term that vanishes asλ → ∞ .
Revised: As can be seen, this covariance matrix is not diagonal, but 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 σ .
Revised: Interestingly, the location parameter of q mix is translated from the prior location µ along thedirection of the prior variance vector σ .

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 ( σ ) 1 = σ .
Revised: 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 ( σ ) 1 = σ .

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 bewritten 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-abidingposterior is thus the true posterior.

Original: We visualise the two respective posterior approximations and thecorresponding likelihood with two different parametrisations (cf.
Revised: We visualise the two respective posterior approximations and thecorresponding likelihood with two different parametrisations (cf.

Original: Sec. 4.3) in Fig.
Revised: Sec. 4.3) in Fig.

Original: 3, App.
Revised: 3, App.

Original: D.
Revised: D.

Original: 4.2 Invariance-gap
Revised: 4.2 Invariance-gap

Original: In order to quantify the detrimental effect of the translation invariance in the considered linear modelwe will now compute the invariance-gap in (11).
Revised: In order to quantify the detrimental effect of the translation invariance in the considered linear modelwe will now compute the invariance-gap in (11).

Original: Note again that we assume a scaled standard normalprior with variance σ = K · σ 0 · 1 .
Revised: Note again that we assume a scaled standard normalprior with variance σ = K · σ 0 · 1 .

Original: In order to simplify the form of the KL divergence, we thenassume that all variances and means take the same value for in the likelihood (and thus also in theposterior) approximation, i.e.
Revised: In order to simplify the form of the KL divergence, we thenassume that all variances and means take the same value for in the likelihood (and thus also in theposterior) 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 ; 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: (15)), and, hence, it makes no difference for q mix ; 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 .

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: ˆ σ +ˆ
Revised: ˆ σ +ˆ

Original: λ ˆ λ with a minimum at ϕ = 1 ; it is minimised as ˆ λ → ∞ .
Revised: λ ˆ λ with a minimum at ϕ = 1 ; it is minimised as ˆ λ → ∞ .

Original: It ishowever not evident from the above whether the gap has a large magnitude compared to the rest ofthe regularisation term and over-regularises in practice.
Revised: It ishowever not evident from the above whether the gap has a large magnitude compared to the rest ofthe regularisation term and over-regularises in practice.

Original: In the next section, we therefore analysethis term at the optima for the mean-field and the invariance-abiding approximation defined in (17),respectively.
Revised: In the next section, we therefore analysethis term at the optima for the mean-field and the invariance-abiding approximation defined in (17),respectively.

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

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 KLterms 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 forboth posterior distributions, giving 2 × 2 combinations.

Original: We denote the optimal parameters wrt.
Revised: We denote the optimal parameters wrt.

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

Original: [Equation]
Revised: [Equation]

Original: 17a
Revised: 17a

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 bewritten in the same form as the true posterior and then simply read out the optimal parameters.

Original: Thisis shown in App.
Revised: Thisis 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 (see [31]),since the Gaussian mean-field distribution that minimises the KL divergence to the multivariate
Revised: The optimal parameters for the mean-field posterior can also be computed analytically (see [31]),since the Gaussian mean-field distribution that minimises the KL divergence to the multivariate

Original: Gaussian true posterior is known.
Revised: Gaussian true posterior is known.

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

Original: [Equation] [Equation]
Revised: [Equation] [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 optimalparameters 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 the prior variance σ = Kσ 0 · 1 such that the prior and posterior overfunctions are identical for any K .
Revised: Note again that we model the prior variance σ = Kσ 0 · 1 such that the prior and posterior overfunctions are identical for any K .

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

Original: This can be seen by the constant (in K )loss terms in Fig.
Revised: This can be seen by the constant (in K )loss terms in Fig.

Original: In contrast, since λ ∗ 0 depends quadratically on K and σ only linearly, q mix ( w ; θ ∗ 0 )collapses to the prior as K → ∞ .
Revised: In contrast, since λ ∗ 0 depends quadratically on K and σ 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.
Revised: This can be seen e.g. by the shrinking KL regularisation (Fig.

Original: 2a)and ELL (in 2b) term.
Revised: 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 predictivevariance 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—theconsequence of not handling translation invariance is indeed posterior collapse as K → ∞ , whilemodelling 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 takeawayis that while we need to optimise for (17a), the mean-field posterior approximation q 0 ( w ; θ ∗ mix ) issufficient 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 exactlyand thereby correct the ELBO objective for the invariances of this model class.

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 invariancegap.

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 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 ).

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 σ = K · 1 , and σ y = 1 / (2 π e) ; we used N = 10 identical observations with value y = 1 and inputs x = 1 .
Revised: Prior variance is σ = K · 1 , and σ 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 decomposesover the NN layers and the associated parameters as follows:
Revised: In order to efficiently compute the invariance-abiding likelihood g mix , note that it also decomposesover the NN layers and the associated parameters as follows:

Original: For the n 1 nodes in the first layer we have
Revised: For the n 1 nodes in the first layer we have

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: where m , V , µ and Σ are constrained to the parts of the overall weight vector that correspond tow 1 ,j .
Revised: where m , V , µ and Σ are constrained to the parts of the overall weight vector that correspond tow 1 ,j .

Original: Now, since the invariance-abiding mechanism for nodes in layer l only depends on the valuez l − 1 of the hidden units in layer n − 1 (see (20) and (21)), we have
Revised: Now, since the invariance-abiding mechanism for nodes in layer l only depends on the valuez l − 1 of the hidden units in layer n − 1 (see (20) and (21)), we have

Original: [Equation]
Revised: [Equation]

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 correspondto w l,j and the expectation over z is taken with respect to z ∼ P ( z l − 1 ) .

Original: Thus, in a practical application, the invariance-abidingmean-field approximation could e.g. becomputed with a layer-by-layer iterative optimization: First, giventhe data x anda prior p ( w ) , weuse (23) and point estimates for all weight vectors in later layers to optimize the mean m and diagonalcovariance V of the weights associated with the first layer, i.e., { w 1 ,j | j = 1 , . . . , n 1 } .
Revised: Thus, in a practical application, the invariance-abiding approximation could e.g. be computed with alayer-by-layer iterative optimization: First, given data x and prior p ( w ) , use (23) and point estimatesfor all weight vectors in later layers to optimize the mean m and diagonal covariance V of the weightsin the first layer, i.e., { w 1 ,j | j = 1 , . . . , n 1 } .

Original: Then, wecan 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 } ) .
Revised: Then, we can use (24) to generate P samples z 1 ,n ofthe activations of the first layer under the (now fitted) optimal q mix ( { w 1 ,j } ) .

Original: Next, for each of these samples z 1 ,n we  use (25) and point estimates for allweight vectors in later layers to optimize the mean m and diagonal covariance V of the weightsassociated with the second layer and average them for the optimal q mix of the weights associated with the second layer.
Revised: Next, for each of thesesamples z 1 ,n we can use (25) and point estimates for all weight vectors in later layers to optimizethe mean m and diagonal covariance V of the weights in the second layer and average them for theoptimal 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 processcan be repeated until the last layer and the whole procedure is iterated until the whole distributionq mix ( w ) does not change in onefull optimization loop.
Revised: This process is repeated until the last layer and the whole procedure can be iterateduntil the q mix ( w ) does not change in one optimization loop.

Original: The complexity of this procedure is no largerthan the optimization of the weight matrices of a single layer because all previous layers are fullycharacterized by the distribution of the latent activations z l of the hidden layers.
Revised: The complexity of this procedure wouldbe no larger than the optimization of the weight matrices of a single layer because all previous layersare fully characterized 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., [34, 15, 30].

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 cannot have large predictive uncertainty between regions of low uncertainty [9];

Original: In contrast, Farquharet al.
Revised: In contrast, Farquharet 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 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.

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 predictivedistribution of mean-field VBNNs reverts to the prior predictive distribution as the network widthincreases.

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 approximateinference 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-fieldapproximation fails to approximate expressive posteriors via the invariance-gap in the ELBO.

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 circumventingthe non-identifiability issue [25, 19, 33].

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 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].

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 workalso proposed to map the BNN prior to a Gaussian process (functional) prior [8, 29].

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 oreven better performance [27, 20, 28, 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 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.

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 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.

Original: More broadly,the (non-)identifiability of parameters and latent variables in probabilistic models is awidely-studied topicboth from frequentist [24] and Bayesian perspectives [5, 22, 10].
Revised: 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].

Original: In a recentexample, Wang et al. [32] attribute posterior collapse inthe context of variational auto-encoders tonon-identifiability of latent variables.
Revised: Recently, Wang et al.[32] 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 inthe 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 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.

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 relatedthis 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 exhibitstranslation 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 analysisof 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 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.

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 yetidentified 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: Future work will focus on approximations of the invariance gap in order to correct the ELBO objective fortranslation and permutation invariances in general neural network functions.
Revised: 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.

Original: Checklist
Revised: Checklist

Original: For all authors...
Revised: For all authors...

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

Original: contributions and scope?
Revised: contributions and scope?

Original: [Yes]
Revised: [Yes]

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

Original: [Yes] We point out limitations andfuture work throughout the paper.
Revised: [Yes] We point out limitations and future work throughout the paper.

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

Original: [No] This ispurely theoretical work without direct negative societal impacts.
Revised: [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?
Revised: (d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?

Original: [Yes]
Revised: [Yes]

Original: If you are including theoretical results...
Revised: If you are including theoretical results...

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

Original: [Yes]
Revised: [Yes]

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

Original: [Yes]
Revised: [Yes]3.

Original: If you ran experiments...
Revised: If you ran experiments...

Original: [No] No experiments
Revised: [No] No experiments

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

Original: [No] No existing assets used.
Revised: [No] No existing assets used.

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

Original: [No] Nocrowdsourcing/human subjects.
Revised: [No] No crowdsourcing/human subjects.