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: Abstract
Revised: Abstract

Original: Variational Bayesian posterior inference often requires simplifying approximations such 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 has associated the variational mean-field approximation for Bayesian neural networks with 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 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.
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 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.
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 approximation reverts 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 in a 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 true posterior can be constructed from a mean-field parametrisation, this is achieved only 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, we transfer 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 provides a 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 deterministic neural networks (NN) such as improving generalization [36], capturing epistemic uncertainty [16], and providing a framework for continual learning methods [17, 22].
Revised: Bayesian neural networks (BNNs) have several appealing advantages compared to deterministic neural networks (NN) such as improving generalization [38], capturing epistemic uncertainty [17], and providing a framework for continual learning methods [18, 23].

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

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 gap compared to deterministic NNs[15].
Revised: BNNs to large datasets [15, 2]—have been identified to be partially responsible for a performance gap compared to deterministic NNs[16].

Original: The most common variational approximation of the posterior is a product 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 for large 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 [13].

Original: 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 .
Revised: Recent work [5] showed that under certain assumptions such as an odd Lipschitz activation function and a finite dataset with bounded likelihood, the optimal mean-field approximation collapses 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 collapses and 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 study invariances in the likelihood function of overparametrised models and their detrimental effect on variational 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 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 .
Revised: For instance, it is well known that NNs exhibit several invariances with respect to their parameters , including node permutation invariance [19], sign-flip invariance (in the case of odd activation functions) [29, 3], scaling invariance (in the case of piece-wise linear activations) [26], 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 likelihood point estimation through stochastic gradient descent is used, because convergence to any of the equivalent 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 invariances are 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 impact of 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 transformations that 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 are constructed 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 the conditions outlined in Sec. 3, the mean-field approximation induces the same posterior predictive distribution 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 objective corresponding to these two approximations differs by the KL divergence between both posterior approximations; 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 both distributions 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 an overparametrised 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 selected as 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 of this 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 that the optimal parameters w.r.t.
Revised: It turns out that the optimal parameters w.r.t.

Original: the ELBO objective with the invariance-abiding distribution result in the 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 w.r.t.
Revised: In contrast, posterior approximations with parameters that are optimal w.r.t.

Original: the mean-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 subspaces in 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 node permutation 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: [19] or App.

Original: 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.
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 basic concepts 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:  Background 2.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 transform their 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 network has 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 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 ) ).
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 denote all learnable parameters of the model by the stacked weight vectors of each layer and node, 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 , 1 , . . .
Revised: [ w T 1 , 1 , . . .

Original: , w T L,n L ] T .
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 variables with a prior distribution p ( w ) , then the posterior p ( w | D ) —induced by the dataset D through the likelihood 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 Bayesian neural 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 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: 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 minimizer of (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 a consequence 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 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.
Revised: This result has been shown in [5] and we recall the result for a Gaussian likelihood with homogeneous noise (for other likelihood functions, we refer to App.

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

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 some fixed 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 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.
Revised: Prior work [5] 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: 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.
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 understanding their 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 Gaussian function with variational parameters and marginalise over all transformations to which the true likelihood 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-abiding posterior approximation q mix which can be related to a mean-field approximation q 0 that does not model 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 identical posterior 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 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].
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 [12].

Original: We model this invariance through transformations 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 w.r.t.
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 same probability 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 continuous translation invariance (see Sec. 4), we model the uniform distribution as a Gaussian with infinite variance.
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 these transformations 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: For instance, permutation invariance induces a factorial number of discontinuous modes that are each equivalent (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 and posterior of an over-parametrised Bayesian linear regression model can be constructed from a product of independent Gaussians.
Revised: [19], App. E); and as we show in Sec. 4, the full covariance Gaussian likelihood and posterior of an over-parametrised Bayesian linear regression model can be constructed from a product of independent Gaussians.

Original: Curiously, it may be sufficient to approximate ℓ 0 while taking into account the 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 approximation g 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, we model 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, the invariance-abiding approximation q mix ( w ; θ ) is a mixture with infinite continuous or finite discrete modes 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 construct approximations 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 write the 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]

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 )); θ ) = p ( t ( w
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 )); θ ) ∝ p ( t ( w

Original: [Equation]
Revised: [Equation]

Original: φ
Revised: φ

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

Original: ( t
Revised: ( t

Original: [Equation]
Revised: [Equation]

Original: w
Revised: w

Original: [Equation]
Revised: [Equation]

Original: φ
Revised: φ

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 defined in (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: 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).
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 relative entropy between the mean-field and invariance-abiding approximation; we therefore refer to it as the invariance 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 ) w.r.t.
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 both 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 ; θ ) (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: 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.
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 ) is uniform as well, and because both posteriors are constructed as the product of prior and likelihood approximation (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 w.r.t.
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: As discussed 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 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.
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 relevant variational parametrisation is the optimal solution w.r.t.
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 which the 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 exhibits translation 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 of translation 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-field and 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 to lift 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 and assume 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 the observation 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 elements unaffected 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 using a 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 dimensions K , such that the 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 w.r.t.
Revised: w k + ϵ (cid:3) is constant w.r.t.

Original: K .
Revised: K .

Original: Another way to view this is to model 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.
Revised: The posterior of this Gaussian linear model can be computed as (see App.

Original: D.3) p
Revised: 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 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.
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 and variances λ .
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 us to 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, and 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 the direction 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 dimensions hyper-plane determined by its normal vector 1 .
Revised: This is because g mix is uniform along the K − 1 dimensions hyper-plane determined by its normal vector 1 .

Original: Taking the Gaussian product then translates the location in the direction Diag ( σ 2 ) 1 = σ 2 .
Revised: Taking the Gaussian product then translates the location in the direction Diag ( σ 2 ) 1 = σ 2 .

Original: Note also that the invariance-abiding posterior can be written 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-abiding posterior 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 the corresponding 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 now compute 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 with variance σ 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 variances and 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 is proportional to 1 , the Gaussian likelihood resulting in the highest ELBO for the approximation qalso 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 is however not evident whether the gap has a large magnitude compared to the rest of the regularisation term 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 optima for 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 corresponding invariance 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 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).
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 KL terms 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 for both 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 w.r.t.
Revised: We denote the optimal parameters w.r.t.

Original: the invariance-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 be written 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: This is 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 true posterior is known [32].
Revised: Gaussian mean-field distribution that minimises the KL divergence to the multivariate Gaussian true posterior is known [34].

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 optimal parameters 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 functions are 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 posterior exactly, 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 being constant 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 predictive variance of the optimal mean-field approximation reverts to the prior predictive variance (Fig. 2d).
Revised: [5], 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—the consequence of not handling translation invariance is indeed posterior collapse as K → ∞ , while modelling 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 takeaway is that, while we need to optimise for (17a), the mean-field posterior approximation q 0 ( w ; θ ∗ mix ) is sufficient 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 exactly and thereby correct the ELBO objective for the invariances of this model.
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 invariance gap.
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 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 ).
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 decomposes over 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 decomposes over 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]
Revised: [Equation]

Original: 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 .
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 depends 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 correspond to 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 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 } .
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 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 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 ,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.
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, and iterated until q mix ( w ) converges.
Revised: This process is repeated until the last layer, and iterated until q mix ( w ) converges.

Original: 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.
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., [35, 15, 31].
Revised: Poor empirical performance of VBNNs has been reported in several previous works, e.g., [37, 16, 33].

Original: 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].
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 [10].

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

Original: [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.
Revised: [8] 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 predictive distribution of mean-field VBNNs reverts to the prior predictive distribution as the network width increases.
Revised: This result seems in conflict with our work and [5], 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 approximate inference algorithms will converge to this solution.
Revised: [8] only shows that the expressive power of the model is large enough but not that the approximate inference algorithms will converge to this solution.

Original: Our work sheds light on why the mean-field approximation 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.

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

Original: 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 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 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 or even better performance [28, 20, 29, 6].
Revised: Surprisingly, restricting the parametrisation of the variational posterior results in competitive or even better performance [30, 21, 31, 7].

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

Original: [6] proposes a variational approximation only for rank-1 factors, resulting in inference of a lower-dimensional subspace.
Revised: [7] proposes a variational approximation only for rank-1 factors, resulting in inference of a lower-dimensional subspace.

Original: The outer product of these low-rank factors perturb a weight matrix that is treated deterministically through maximum a posteriori (MAP) estimation.
Revised: 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 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.
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.

Revised: Other attempts have proposed to approximate the posterior predictive directly, thereby circumventing the non-identifiability issue [28, 20, 36].

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

Revised: In a similar vein, previous work also proposed to map the BNN prior to a Gaussian process (functional) prior [9, 32].

Revised: Another promising direction to circumvent invariance in layered models are deep kernel processes, in which Gram matrices are progressively transformed by nonlinear kernel functions [1].

Revised: The Gram matrices, which are treated as the random variables, are invariant to permutations/rotations of the weights/features.

Original: 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].
Revised: More broadly, (non-)identifiability of parameters and latent variables in probabilistic models is a widely-studied topic both from frequentist [27] and Bayesian perspectives [6, 25, 11, 24].

Original: Recently, Wang et al.
Revised: In a recent example, Wang et al.

Original: [33] attributed posterior collapse in  variational auto-encoders to non-identifiability of latent variables.
Revised: [35] attributed posterior collapse in the context of 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 in the 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 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.
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 related this 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 exhibits translation 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 analysis of 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 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.
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 yet identified 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: 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. 
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. Checklist

Revised: For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?

Revised: [Yes] (b) Did you describe the limitations of your work?

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

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

Revised: [No] This is purely theoretical work without direct negative societal impacts.

Revised: (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?

Revised: [Yes] 2.

Revised: If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?

Revised: [Yes] (b) Did you include complete proofs of all theoretical results?

Revised: [Yes] 3.

Revised: If you ran experiments...

Revised: [No] No experiments 4.

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

Revised: [No] No existing assets used.

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

Revised: [No] No crowdsourcing/human subjects.