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

Original: Importantly, this invariance gap is not constant; it vanishes as the approximationreverts to the prior.
Revised: Importantly, this invariance gap is not constant; it vanishes as the 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 so far been impeded [34].
Revised: Unfortunately, reaping thesetheoretical benefits has so far been impeded [34].

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

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

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

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

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

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

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

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

Original: For instance, it is well known that NNs exhibit several invariances with respectto their parameters , including node permutation invariance [18], sign-flip invariance (in the caseof odd activation functions) [26, 2], scaling invariance (in the case of piece-wise linear activations)
Revised: For instance, it is well known that NNs exhibit several invariances with respectto their parameters , including node permutation invariance [18], sign-flip invariance (in the caseof odd activation functions) [26, 2], scaling invariance (in the case of piece-wise linear activations)

Original: [23], and as we will show in Sec. 5, the parameter space of BNNs additionally has subspaces withtranslation invariance .
Revised: [23], and as we will show in Sec. 5, the parameter space of BNNs additionally has subspaces withtranslation 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, these invariancesare detrimental to the variational Bayesian approach.
Revised: However, these invariancesare detrimental to the variational Bayesian approach.

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

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

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

Original: We then prove that, under theconditions outlined in Sec. 3, the mean-field approximation induces the same posterior predictivedistribution as the invariance-abiding approximation.
Revised: We then prove that, under theconditions outlined in Sec. 3, the mean-field approximation induces the same posterior predictivedistribution as the invariance-abiding approximation.

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

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

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

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

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

Original: It turns out thatthe optimal parameters wrt.
Revised: It turns out thatthe optimal parameters wrt.

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

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

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

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

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

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

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

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

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

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 , where i indexes the layer-wise vectors z l .
Revised: Such a networkhas n 1 + n 2 + · · · + n L many nodes z l,i , where i indexes the layer-wise vectors z l .

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

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

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

Original: , w T n 1 + ··· + n L ] T todenote all learnable parameters of the model.
Revised: , w T n 1 + ··· + n L ] T todenote 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 variableswith a prior distribution p ( w ) , then the posterior p ( w | D ) —induced by the dataset D through thelikelihood function ℓ ( w ; D ) :
Revised: If the weights and biases w are treated as random variableswith a prior distribution p ( w ) , then the posterior p ( w | D ) —induced by the dataset D through thelikelihood function ℓ ( w ; D ) :

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

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

Original: [Equation]
Revised: [Equation]

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

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

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

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

Original: [Equation]
Revised: [Equation]

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

Original: =
Revised: =

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

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

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

Original: A mean-field variational BNN (VBNN) is the minimizerof (1) where Q is the family of Gaussians with diagonal covariance matrix.
Revised: A mean-field variational BNN (VBNN) is the 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 , σ 2 y ) ,an isotropic Gaussian prior p ( w ) =
Revised: f w ( x )+ ϵ , with ϵ ∼ N (0 , σ 2 y ) ,an isotropic Gaussian prior p ( w ) =

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

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

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

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

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

Original: Then,
Revised: Then,

Original: [Equation]
Revised: [Equation]

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

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

Original: Prior work [4] has used the data-related bound to prove that the predictivedistribution of mean-field BNNs converges to the prior predictive distribution if the network widthis large and the activation function is odd, ultimately resulting in posterior collapse to the prior.
Revised: Prior work [4] has used the data-related bound to prove that the 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 indexing these transformations.
Revised: For discontinuous transformations such as node permutations (see App. E), p ( r )is a uniform distribution over discrete variables indexing these transformations.

Original: For the continuoustranslation invariance (see Sec. 4), we model the uniform distribution as a Gaussian with infinitevariance.
Revised: For the 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 using the same parametrisation through
Revised: From this single mode approximation, wemodel an invariance-abiding likelihood approximation using the same parametrisation through

Original: [Equation]
Revised: [Equation]

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

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

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

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

Original: While q 0 ( w ; θ ) is a mean-field approximation, theinvariance-abiding approximation q mix ( w ; θ ) is a mixture with infinite continuous or finite discretemodes depending on the type of invariance.
Revised: While q 0 ( w ; θ ) is a mean-field approximation, 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] [Equation]
Revised: [Equation] [Equation]

Original: [Equation]
Revised: [Equation]

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

Original: =
Revised: =

Original: p ( t ( w , φ ( r ))) · g 0 ( t ( w , φ ( r )); θ ) , because then
Revised: p ( t ( w , φ ( r ))) · g 0 ( t ( w , φ ( 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 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] [Equation]
Revised: [Equation] [Equation]

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

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

Original: [Equation]
Revised: [Equation]

Original: Invariance gap.
Revised: Invariance gap.

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

Original: C)
Revised: C)

Original: [Equation]
Revised: [Equation]

Original: Interestingly, the gap between the two respective ELBO objectives is given exactly by the relativeentropy between the mean-field and invariance-abiding approximation; we therefore refer to it as theinvariance gap .
Revised: Interestingly, the gap between the two respective ELBO objectives is given exactly by the 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 ; θ ) (used to construct bothq 0 and q mix ), we see that the former objective favors solutions where q 0 and q mix coincide.
Revised: the parameters of the variational likelihood approximation g 0 ( w ; θ ) (used to construct bothq 0 and q mix ), we see that the former objective favors solutions where q 0 and q mix coincide.

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

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

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

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

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

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

Original: [Equation]
Revised: [Equation]

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

Original: ELBO objective wrt.
Revised: ELBO objective 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: Asdiscussed above, one specific parametrisation for which the invariance-gap vanishes is the meanfield posterior collapse, q 0 ( w ; θ )
Revised: Asdiscussed above, one specific parametrisation for which the invariance-gap vanishes is the meanfield posterior collapse, q 0 ( w ; θ )

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

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

Original: One particularly relevantvariational parametrisation is the optimal solution wrt.
Revised: One particularly relevantvariational parametrisation is the optimal solution wrt.

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

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

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

Original: We study an over-parametrised Bayesian linear regression model as the canonical model that exhibitstranslation invariance.
Revised: We study an over-parametrised Bayesian linear regression model as the canonical model that 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 mean-fieldand invariance-abiding distribution defined in (6) and the invariance gap from (10).
Revised: the mean-fieldand invariance-abiding distribution defined in (6) and the invariance gap from (10).

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

Original: Likelihood model.
Revised: Likelihood model.

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

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

Original: . , w K ] T andassume that we have N observations of variables y given dependent inputs x .
Revised: . , w K ] T 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 for the more general case).
Revised: D for the more general case).

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

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

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

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

Original: [Equation]
Revised: [Equation]

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

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

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

Original: Prior.
Revised: Prior.

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

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

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

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

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

Original: Posterior.
Revised: Posterior.

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

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 considering p (∆) =
Revised: We model that the likelihood is translation invariant by computing the marginal from (5)and considering p (∆) =

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

Original: [Equation]
Revised: [Equation]

Original: Writing the uniform distribution as the limiting case of a Gaussian with infinite variance allows usto compute the integral analytically.
Revised: Writing the uniform distribution as the limiting case of a Gaussian with infinite variance allows usto 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; it has an additive rank-1 term that vanishes asλ → ∞ .
Revised: As can be seen, this covariance matrix is not diagonal; it has an additive rank-1 term that vanishes asλ → ∞ .

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

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

Original: Note also that the invariance-abiding posterior can bewritten in the same form as the true posterior in (13) (see App. D.3).
Revised: Note also that the invariance-abiding posterior can 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 in Fig.
Revised: We visualise the two respective posterior approximations and thecorresponding likelihood in Fig.

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

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

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

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

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

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

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

Original: [Equation]
Revised: [Equation]

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

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

Original: However, since the prior variance isproportional to 1 , the Gaussian likelihood resulting in the highest ELBO for the approximation q 0also has a variance vector proportional to 1 .
Revised: However, since the prior variance 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: ˆ σ 2 +ˆ λ ˆ λ with a minimum at ϕ = 

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

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

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

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

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

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

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: 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 since the
Revised: The optimal parameters for the mean-field posterior can also be computed analytically since the

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

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

Original: [Equation]
Revised: [Equation]

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

Original: We then construct the two respective posterior approximations q 0 and q mix with these two optimalparameters and compare the 2 × 2 combinations in Fig.
Revised: We then construct the two respective posterior approximations q 0 and q mix with these two 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 prior variances σ 2 = Kσ 20 · 1 such that prior and posterior over functionsare identical for any K .
Revised: Note again that we model prior variances σ 2 = Kσ 20 · 1 such that prior and posterior over functionsare identical for any K .

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

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

Original: 2 beingconstant in K .
Revised: 2 beingconstant 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.
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 σ 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 258 over the NN layers and the associated parameters as follows: For the n 1 nodes in the first layer we have
Revised: In order to efficiently compute the invariance-abiding likelihood g mix , note that it also decomposes 258 over the NN layers and the associated parameters as follows: For the n 1 nodes in the first layer we have

Original: [Equation]
Revised: [Equation]

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

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

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

Original: [Equation]
Revised: [Equation]

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

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

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

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

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 the whole procedure can be iterated until theq mix ( w ) does not change in one optimization loop.
Revised: This process is repeated until the last layer and the whole procedure can be iterated until theq mix ( w ) does not change in one optimization loop.

Original: The complexity of this procedure would be nolarger than the optimization of the weight matrices of a single layer because all previous layers arefully characterized by the distribution of the latent activations z l of the hidden layers.
Revised: The complexity of this procedure would be nolarger than the optimization of the weight matrices of a single layer because all previous layers arefully characterized by the distribution of the latent activations z l of the hidden layers.

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

Original: Recently, Wang et al.[32] attributed posterior collapse in variational auto-encoders to non-identifiability of latent variables.
Revised: Recently, Wang et al.[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: 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.
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’scontributions and scope?
Revised: (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s 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]2.
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]3.
Revised: [Yes]

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...[No] Nocrowdsourcing/human subjects. A Notation
Revised: If you used crowdsourcing or conducted research with human subjects...

Original: We denote matrices and vectors with bold upper- and lower-case letters, e.g. a and A .

Original: In order toaddress an element of a vector or matrix, we will use non-bold letters, e.g. x = [ x 1 , x 2 , . . .

Original: , x N ] T .

Original: We will use the superscript T to denote a transpose of a vector and | A | to denote the determinant ofthe matrix A .

Original: We use Diag ( a ) to denote the diagonal matrix constructed from a vector, and diag ( A )to denote the diagonal vector of a matrix.

Original: We use the notation N ( x ; m , V ) to denote the density ofthe Gaussian distribution given by

Original: [Equation]

Original: An alternative representation of the Gaussian distribution is in terms of its canonical parametersη =

Original: Σ − 1 µ and Λ =

Original: Σ − 1 :

Original: [Equation]

Original: Note that G ( x ; η 1 , Λ 1 ) · G ( x ; η 2 , Λ 2 ) ∝ G ( x ; η 1 + η 2 , Λ 1 + Λ 2 ) which renders this canonicalparameterisation very useful when considering the multiplication of Gaussian densities.

Original: Furthermore,we use E w ∼ p [ f ( w )] to denote the expectation of the function f ( · ) when w is drawn from thedistribution p .

Original: Also, we use V w ∼ p [ f ( w )] to denote the variance of the function f ( · ) when w isdrawn from the distribution p defined by

Original: [Equation]

Original: B Convolution of Gaussian Measures

Original: Theorem 1 (Convolution of Gaussian Measures) .

Original: For any A ∈ R n × k and b ∈ R n it holds that

Original: [Equation]

Original: p

Original: [Equation]

Original: µ

Original: [Equation]

Original: Σ ) , implies that

Original: [Equation] [Equation] [Equation] Proof.
Revised: [No] Nocrowdsourcing/human subjects.

Original: Let’s start by proving (27).

Original: First, we note that

Original: [Equation]

Original: where only the numerator depends on θ .

Original: Thus, the numerator is given by

Original: [Equation]

Original: where c = | 2 π V | − 12 ·| 2 π Σ | − 12 is independent of x and θ .

Original: Using Theorem A.86 in [13], this quadraticform in θ can be rewritten as

Original: [Equation]

Original: , where

Original: [Equation]

Original: Since d ( x ) does not depend on θ , it get incorporated into the normalization constant which proves(27).

Original: In order to prove (28), note that

Original: [Equation]

Original: where the penultimate line follows again from Theorem A.86 in [13] with d ( x ) defined in (29).

Original: Theorem 2 (Woodbury formula) .

Original: Let C be an invertible n × n matrix.

Original: Then, for any A ∈ R n × k and B ∈ R k × n ,

Original: [Equation]

Original: C Invariance gap and posterior-predictive equivalence

Original: In this section, we show that the posterior predictive distribution of the mean-field and invarianceabiding distribution are identical (for identical parametrisation of the likelihood approximation) asstated in (9), and that the difference in the respective KL regularisation terms can be quantified by the

Original: KL defined in (10).

Original: We first discuss the conditions from (7) and (8).

Original: The goal is to use the invariance property of thelikelihood function (4).

Original: Unfortunately, the prior does not generally have the same invariance.

Original: Yet, themean-field approximate posterior defined as the product between prior and likelihood approximation can still has the invariance property as we show for permutation and translation invariance withmean-field Gaussian likelihood approximation and prior.

Original: This condition is defined in (7) for eachmean-field posterior in the integral over r :

Original: [Equation]

Original: where Z mix and Z 0 ( r ) are normalisation constants.

Original: The second line introduced the assumption thatthere exists a surjective mapping r ′ =

Original: φ ( r ) such that the product of the untransformed prior p ( w )and the transformed likelihood approximation g 0 ( t ( w , φ ( r )); θ ) are identical to the product of thetransformed prior and likelihood approximation with r ′ , i.e. as stated in the main text,

Original: [Equation]

Original: = p

Original: [Equation]

Original: This condition may seem quite limiting, however we show that this condition holds for permutationinvariance with the isotropic Gaussian prior and for translation invariance with Gaussian priors and

Original: Gaussian likelihood approximation.

Original: Condition (7) and (8) for permutation invariance.

Original: Consider first permutation invariance(cf.

Original: App. E).

Original: It is easy to see that for the isotropic Gaussian prior: p ( w ) =

Original: p ( P r w ) .

Original: Also, note that

Original: G ( Pw ; η , Λ ) =

Original: G ( w ; P T η , P T ΛP ) .

Original: Thus, the product between the untransformed Gaussian priorand the permuted Gaussian likelihood approximation is

Original: [Equation]

Original: Hence, for permutation invariance with the isotropic Gaussian prior and Gaussian likelihoodapproximation we have

Original: [Equation]

Original: [Equation]

Original: The invariance transformation is thus simply t ( w , φ ( r ))

Original: =

Original: t ( w , r ) =

Original: P r w (i.e. we do not need themapping φ ).

Original: This is because the prior has no preference over the permutation-induced modes of thelikelihood.

Original: Thus, we have

Original: [Equation]

Original: Condition (7) and (8) for translation invariance.

Original: Next, consider translation invariance, and notethat N ( w − v ; µ , Σ ) =

Original: N ( w ; µ + v , Σ ) , and, consequently, G ( w − v ; η , Λ ) = G ( w ; η + Λv , Λ ) .

Original: It is easier to show directly that the product between the untransformed Gaussian prior and translated

Original: Gaussian likelihood approximation can be written as a Gaussian posterior translated , because:

Original: [Equation]

Original: Since the precision matrices are diagonal, (( Λ p + Λ g ) − 1 Λ g )

Original: Br = B (( Λ p + Λ g ) − 1 Λ g ) r .

Original: Consequently, for the translation invariance g 0 ( w ) =

Original: g 0 ( t ( w , r ))

Original: =

Original: g 0 ( w − Br ) , we have

Original: [Equation]

Original: ) where r ′ =

Original: φ ( r ) = (cid:16) ( Λ p + Λ g ) − 1 Λ g (cid:17) r .

Original: Thus, we have

Original: [Equation]

Original: Lemma 1.

Original: For any distribution p ( w ) , likelihood approximation g 0 ( w ; θ ) , q 0 ( w ; θ ) as defined in (6) ,q mix ( w ; θ ) as defined in (6b) and p ( r ) , assume there exists a mapping φ : r (cid:55)→ r ′ such that

Original: [Equation]

Original: as well as

Original: [Equation]

Original: [Equation]

Original: Then,

Original: [Equation]

Original: Proof.

Original: The lemma can be proven by applying the change-of-variables formula and using theinvariance property (42) of g ( · ; θ ) :

Original: [Equation]

Original: 51) where the second line uses (42), the third line changes the order of integration (Fubini’s theorem), thefourth line uses the invariance property ln p ( D | w ) =

Original: ln p

Original: ( D | t ( w , φ ( r ))) , the fifth line then appliesthe change of variables theorem together with (43), the sixth line then uses again the invarianceproperty of the log likelihood, and the seventh line notes that the integral over the normalisationconstants Z 0 ( r ) cancels with the normalisation constant Z mix of the invariance-abiding posteriorsince

Original: [Equation]

Original: 54) where we changed the integration order, resulting in the normalisation constant of the invarianceabiding likelihood approximation.

Original: Lemma 2.

Original: For any distribution p ( w ) , likelihood approximation g 0 ( w ; θ ) , q 0 ( w ; θ ) as defined in(6) , q mix ( w ; θ ) as defined in (6b) and p ( r ) , assume there exist two mappings t : w × r (cid:55)→ w ′ andφ : r (cid:55)→ r ′ such that (42) and (43) hold.

Original: Then,

Original: KL

Original: [Equation]

Original: |

Original: [Equation]

Original: [Equation]

Original: Proof.

Original: First, note that

Original: [Equation]

Original: [Equation]

Original: [Equation]

Original: Now, in the latter KL

Original: [Equation]

Original: we expand the distribution q mix ( w ) as in (42),

Original: [Equation]

Original: where the third line uses the invariance property of the invariance-abiding likelihood approximation,∀ r : g mix ( t ( w , φ ( r )))

Original: =

Original: g mix ( w ) , the change of variables formula is applied to the fourth line forthe volume-preserving transformation, and, the fifth line re-arranges the integration, noting that theintegral over normalisation constants equals one.

Original: The proposition follows by taking the difference between the regularisation terms corresponding tothe respective ELBO objectives:

Original: [Equation]

Original: [Equation]

Original: D Translation invariance in linear models

Original: In this section, we derive the results from Sec.

Original: 4 for a Bayesian linear regression with a single inputvector x and corresponding target observation y .

Original: The result stated in Sec. 4 for x = 1 follows as aspecial case.

Original: We assume we have K latent variables w =

Original: [ w 1 , . .

Original: . , w K ] T and one observation x =

Original: [ x 1 , . .

Original: , x K ] T .

Original: We further assume that the likelihood p ( y, | w , x ) only depend on their inner product, that is, x T w .

Original: Then we know that any change in w which leaves the sum of the elements weighted by x unaffecteddoes not change the likelihood; for example w ′ =

Original: [ w 1 + ∆ , w 2 − ∆ · x 1 x 2 , w 3 , . .

Original: . , w K ] T has theexact same likelihood than w =

Original: [ w 1 , w 2 , w 3 , . .

Original: . , w K ] T .

Original: In general, we can model this translation invariance using an K − 1 dimensional vector ∆ ∈ R K − 1 and noting that

Original: [Equation]

Original: where we used x K − 1 :

Original: =

Original: [ x 1 , . .

Original: , x K − 1 ] T because

Original: [Equation]

Original: D.1 Likelihood Model

Original: Now let us assume that we approximate the likelihood by a function that has Gaussian shape, that isp ( y, | w , x ) ≈ q ( w ) ∝ N ( w ; m , V ) .

Original: In order to model that the likelihood is translation invariant,we compute the marginal q ( w − B∆ ) over p ( ∆ ) =

Original: N (cid:0) ∆ ; 0 , β 2 I (cid:1) and considering the case ofβ → ∞ , that is

Original: [Equation]

Original: According to Theorem 1, for any β ∈ R + this is another Gaussian given by

Original: [Equation]

Original: Note that lim β →∞ q β ( w ) =

Original: g mix ( w ) as defined in Subsection 4.1.

Original: Using the Woodbury formula in

Original: Theorem 2, the inverse of the covariance can be re-written as

Original: [Equation]

Original: D.2 Posterior Model

Original: If we assume a prior p ( w ) =

Original: N ( w ; µ , Σ ) , then the posterior for the Gaussian approximation isgiven by

Original: [Equation]

Original: where we used the identity (cid:0) Σ − 1 + V

Original: [Equation]

Original: [-1 (cid:1) − 
-]
Revised: [-1 (cid:1) − 
-]

Original: = V ( V + Σ )

Original: [Equation]

Original: 1 Σ = Σ ( V + Σ ) − 1 V in the last line.

Original: Similarly, if we use the likelihood approximation which has incorporated the translation invariance, we get

Original: [Equation]

Original: [-Observing that lim β →∞ β − 2 I = 
-]
Revised: [-Observing that lim β →∞ β − 2 I = 
-]

Original: we thus see that

Original: [Equation]

Original: Note again that p ∞ ( w ) =

Original: q mix ( w ; θ ) as defined in (15).

Original: D.2.1 Special Case of Diagonal Likelihood Covariance

Original: Now let us consider the special case where the covariance matrix of the Gaussian likelihoodapproximation is a diagonal matrix, that is V

Original: =

Original: Diag ( λ ) .

Original: Then, observe that

Original: [Equation]

Original: where we used the notation V − 1 K − 1 to denote the diagonal ( K − 1)

Original: × ( K − 1) matrix with all

Original: [Equation]

Original: on the diagonal.

Original: Thus, using Theorem 2 with C =

Original: V − 1 K − 1 , A = x − 2 K

Original: λ − 1 K x K − 1 and B = x T K − 1 , wesee that

Original: [Equation]

Original: Covariance (cid:0) V − 1 B + Σ − 1 (cid:1) .

Original: Thus, for the covariance (67) of the posterior we have

Original: [Equation]

Original: Let’s focus on the expression S first.

Original: Using (68) we have

Original: [Equation]

Original: Similarly, for the expression T using (69) we have

Original: [Equation]

Original: Putting (70) and (71) together, we get

Original: [Equation]

Original: where we repeatedly used that x T Vx − x T K − 1 V K − 1 x K − 1 = λ K x 2 K .

Original: Thus, the covariance in (66)can be written as

Original: [Equation]

Original: where we used Theorem 2 in the third step.

Original: [-Note that (14) is a special case of (72) when using x = 1and observing that V
-]
Revised: [-Note that (14) is a special case of (72) when using x = 1and observing that V
-]

Original: = λ .

Original: Similarly, the covariance in (15) is a special case of (73) when furthernoticing that Σ1 =

Original: σ 2 .

Original: Mean (cid:0)

Original: [Equation]

Original: Σ

Original: [Equation]

Original: µ

Original: [Equation]

Original: In order to derive an efficient update for the mean ofthe posterior, please note that by virtue of (72), V − 1 B can be written as dd T with d = (cid:0) x T Vx (cid:1) − 12 · x .

Original: Thus, using (73) we have

Original: [Equation]

Original: Thus, the location parameter in (15) is a special case of this more general result when using x = 1and noticing again that Σ1 =

Original: σ 2 and V1 = λ , respectively.

Original: It can be seen that the Gaussian productupdates the prior location in the direction Σx .

Original: This is because the likelihood is translation invariantwrt.

Original: all directions perpendicular to x (i.e. the hyper plane determined by the normal vector x ).

Original: The two posterior approximations q 0 and q mix as well as the corresponding likelihood approximationis visualised in Fig.

Original: 3 for two different parametrisations (cf.

Original: Sec. 4.3).

Original: [Figure]

Original: D.3 True posterior and optimal invariance-abiding parameters

Original: For the linear model with a single observation y and inputs x , the true posterior p ( w | x , y ) followsfrom the standard Bayesian update equation for Gaussian linear models (cf.

Original: (27) in App.

Original: B with

Original: A = 1 K · x T ):

Original: [Equation]

Original: [Equation]

Original: For N observations Y :=

Original: { y ( n ) } Nn =1 with identical input x , the posterior is

Original: [Equation]

Original: We want to relate the true posterior (75) to the parameters of the invariance-abiding posteriorq mix ( w ; θ ) .

Original: It suffices to consider the special case x = 1 from the main text.

Original: We will use the form

Original: [Equation]

Original: For the precision matrix Σ − 1 + V − 1mix , using the form in (72) with x = 1 gives

Original: [Equation]

Original: since V1 = λ .

Original: Comparing this to (75), we see that Σ − 1 + V − 1mix has the same structure as theprecision matrix of the true posterior.

Original: By setting the optimal variances as 1 T λ ∗ =

Original: K 2 σ 2 y N , and using

Original: K 2 σ 2 y N =

Original: Kσ 2 y N · 1 T 1 , we see that one possible choice for the optimal variance parameters is

Original: [Equation]

Original: We note that other vectors that are not proportional to 1 and also sum to K 2 σ 2 y N are also valid optima.

Original: Similarly, with µ = 0 , and by noting that Σ − 1 + V − 1mix =

Original: V − 1 p , we see that

Original: [Equation]

Original: [Equation]

Original: E Permutation invariance in Bayesian neural networks

Original: Here we analyse the permutation invariance in BNNs.

Original: This invariance is independent of the data,persisting for any dataset size.

Original: We first describe the set of permutation matrices corresponding to the transformations that leavethe likelihood invariant, i.e. t ( w , r ) =

Original: P r w .

Original: We ignore the biases for simplicity and describe thesepermutations first on a node/neuron level, then layer-wise for the weight matrices, and finally on theweight vector that is obtained by stacking all weight matrices.

Original: We will denote layers with indicesl and the number of layers by L .

Original: Layer-wise permutation matrices are then denoted as ˜ P l and thecorresponding weight matrices are denoted as W l ; the permutation matrix that results when stackingall weights matrices into a vector w is denoted as P .

Original: When necessary, one particular permutationmatrix from the set of all possible permutation matrices for a given architecture is indexed withsuperscript ( i ) , i.e. P = { P ( i ) } i and |P| =

Original: (cid:81)

Original: L − 1 l =1 k l ! , where k l is the number of nodes in layer l .

Original: Node permutations.

Original: Each hidden layer z l ∈ z 1 , . . .

Original: , z L − 1 is a set of nodes that can be permutedin k l ! possible ways, relabelling the corresponding parameters attached to these nodes.

Original: Each of thesek l ! permutations per layer can be combined with any of the permutations of another layer.

Original: Eachpermutation matrix ˜

Original: P l for a layer l corresponds to one of the unique orderings of nodes z l .

Original: Forinstance, the permutation matrix that reorders the first 3 nodes as ( z l, 3 , z l, 1 , z l, 2 )

Original: is˜ P l =     . . .

Original:   .

Original: [Equation]

Original: [Equation]

Original: The remaining entries in the matrix are ones on the diagonal and zeros elsewhere.

Original: Layer-wise permutation of weight matrices.

Original: Next, we describe how node permutations correspondto permutations of weight matrices in terms of permutations to the in- and outgoing weights.

Original: Forone particular instance of permutations (omitting superscript ( i ) ) to each of the hidden layers, thecorresponding weight matrices can be permuted as follows:

Original: [Equation]

Original: The identities corresponding to the first and last layers is because only hidden layer nodes can bepermuted but not the data itself.

Original: Note also that each permutation matrix is applied to two weightmatrices, since permutations to nodes of a particular layer correspond to the simultaneous permutationof the weight matrices from the preceding and subsequent layer.

Original: Permutation of stacked weight vectors.

Original: The layer-wise formulation can be written in terms of thestacked weight vector w = [vec

Original: [Equation]

Original: , vec

Original: [Equation]

Original: L

Original: [Equation]

Original: 1 )] T , using vec( ABC ) = (cid:0) C T ⊗ A (cid:1) vec( B ) :

Original: [Equation]

Original: (82) where we denote the new permutation matrix that is applied to the vectorised weights with a bar andthe subscripts correspond to the two successive layers.

Original: The overall permutation matrix correspondingto the entire weight vector w is given by forming the block-diagonal matrix

Original: P =  

Original: P

Original: [Equation]

Original: P

Original: [Equation]

Original: P 3 ,. . .

Original:  

Original: [Equation]

Original: (83) where each denotes block-matrices of zeros with the respective dimensions and the remaining partsof the matrix are the identity on the diagonal and zeros elsewhere.

Original: Invariance gap.

Original: Note again that the permutation invariance for BNNs is independent of the data.

Original: The invariance gap KL [ q 0 || q mix ] takes values in the range [ 0 , ln |P| ] , where |P| is the factorialnumber of modes.

Original: The gap takes the maximal value when each mode is completely separated fromthe other modes, and the gap is zero when both q 0 ( w ) and q mix ( w ) are identical; this is the case e.g.if q 0 ( w ) reverts to the prior p ( w ) since all modes are then identical, i.e. q 0 ( Pw )

Original: =

Original: q ( w ) .

Original: F Data-related bound on the mean-field KL divergence

Original: Assume the regression setting described in Sec. 2.2, where we assume a finite dataset D

Original: ={ ( x ( n ) , y ( n ) ) }

Original: Nn =1 and a regression setting with fixed homogeneous noise variance σ 2 y .

Original: We canfurther assume that the prior is chosen such that it induces a reasonably bounded output variance ofthe neural network, σ 2 L (cid:0) x ( n ) (cid:1) :

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

Original: ( x ( n ) ; w ) (cid:3) , for some fixed input x ( n ) .

Original: We will then havea finite expected log-likelihood and the ELBO for a mean-field variational approximation q θ is

Original: [Equation]

Original: Let us now consider a hypothetical the worst-case fit in terms of the ELBO objective.

Original: This is whenthe data is completely ignored with q θ ( w ) = p ( w ) as any worse fit could be trivially improved bysetting the approximation to the prior.

Original: This gives then the inequality (using (84))∀ q θ :

Original: [Equation]

Original: Although we can not compute the expected log-likelihood for an arbitrary fit q θ , we can quantify theupper bound given by consider the best-case fit q ∗ , i.e. a hypothetical optimum where the data isperfectly predicted up to the known observation noise variance σ 2 y .

Original: The resulting bound is then

Original: [Equation]

Original: Next, we compute the two expected log-likelihood terms in (86).

Original: We first consider the worst-case,where we have

Original: [Equation]

Original: where z L is the noisy output of the neural network given by

Original: [Equation]

Original: Splitting the expectation of the quadratic term into variance and square of the expectation, we have

Original: [Equation]

Original: Since the last layer computes f ( x ( n ) ; w ) =

Original: w T L, 1 z L − 1 in the regression setting and the prior weightsare independent of z L − 1 and centred at zero, i.e. E w ∼ p [ w L, 1 ] = 0 , it follows that

Original: [Equation]

Original: 92) where we denote the variance of the neural network outputs by σ 2 L (cid:0) x

Original: [Equation]

Original: n

Original: [Equation]

Original: (cid:1)

Original: [Equation]

Original: Hence, the expected log likelihood under the prior is

Original: [Equation]

Original: For the best-case fit in terms of the ELBO, we assume that we completely overfit and perfectly predictthe data by putting the mean of the Gaussian exactly on the data, i.e.

Original: [Equation]

Original: Then, we have

Original: [Equation]

Original: Hence, the expected log likelihood of the best case fit is

Original: [Equation]

Original: Taking the difference between (93) and (97), we obtain the result in (3),

Original: [Equation]