Original: Pathologies in Priors and Inference for Bayesian Transformers
Revised: Pathologies in Priors and Inference for Bayesian Transformers

Original: Anonymous Author(s)

Original: Abstract
Revised: Abstract

Original: In recent years, the transformer has established itself as a workhorse in manyapplications ranging from natural language processing to reinforcement learning.
Revised: In recent years, the transformer has established itself as a workhorse in many applications ranging from natural language processing to reinforcement learning.

Original: Similarly, Bayesian deep learning has become the gold-standard for uncertaintyestimation in safety-critical applications, where robustness and calibration are crucial.
Revised: Similarly, Bayesian deep learning has become the gold-standard for uncertainty estimation in safety-critical applications, where robustness and calibration are crucial.

Original: Surprisingly, no successful attempts to improve transformer models in termsof predictive uncertainty using Bayesian inference exist.
Revised: Surprisingly, no successful attempts to improve transformer models in terms of predictive uncertainty using Bayesian inference exist.

Original: In this work, we studythis curiously underpopulated area of Bayesian transformers.
Revised: In this work, we study this curiously underpopulated area of Bayesian transformers.

Original: We ﬁnd that weightspace inference in transformers does not work well, regardless of the approximateposterior.
Revised: We ﬁnd that weightspace inference in transformers does not work well, regardless of the approximate posterior.

Original: We also ﬁnd that the prior is at least partially at fault, but that it is veryhard to ﬁnd well-speciﬁed weight priors for these models.
Revised: We also ﬁnd that the prior is at least partially at fault, but that it is very hard to ﬁnd well-speciﬁed weight priors for these models.

Original: We hypothesize thatthese problems stem from the complexity of obtaining a meaningful mapping fromweight-space to function-space distributions in the transformer.
Revised: We hypothesize that these problems stem from the complexity of obtaining a meaningful mapping from weight-space to function-space distributions in the transformer.

Original: Therefore, movingcloser to function-space, we propose a novel method based on the implicit reparameterization of the Dirichlet distribution to apply variational inference directlyto the attention weights.
Revised: Therefore, moving closer to function-space, we propose a novel method based on the implicit reparameterization of the Dirichlet distribution to apply variational inference directly to the attention weights.

Original: We ﬁnd that this proposed method performs competitivelywith our baselines.
Revised: We ﬁnd that this proposed method performs competitively with our baselines.

Original:  Introduction
Revised:  Introduction

Original: The transformer [57] is a deep learning architecture commonly used to process sequences of data,such as text.
Revised: The transformer [58] is a deep learning architecture commonly used to process sequences of data, such as text.

Original: Thanks to multi-head self-attention, the transformer builds contextual embeddingsby capturing the relationships between the sequence elements.
Revised: Thanks to multi-head self-attention, the transformer builds contextual embeddings by capturing the relationships between the sequence elements.

Original: While being most famous for theirstate-of-the-art performance in natural language processing [5, 14], transformers are also used incomputer vision [7, 15, 32, 52, 61], reinforcement learning [6, 37], as well as audio [25, 29, 47] andvideo [63] processing, yielding impressive results.
Revised: While being most famous for their state-of-the-art performance in natural language processing [5, 14], transformers are also used in computer vision [7, 15, 33, 53, 63], reinforcement learning [6, 38], as well as audio [25, 29, 48] and video [65] processing, yielding impressive results

Original: The Bayesian learning paradigm provides a theoretical framework to obtain predictive uncertainty,select the optimal model and improve its calibration.
Revised: The Bayesian learning paradigm provides a theoretical framework to obtain predictive uncertainty, select the optimal model, and improve its calibration.

Original: Furthermore, by designing an informativeprior for the parameters, Bayesian models offer a principled way to incorporate assumptions aboutthe inferred distribution, thus providing regularization.
Revised: Furthermore, by designing an informative prior for the parameters, Bayesian models offer a principled way to incorporate assumptions about the inferred distribution, thus providing regularization.

Original: Finally, recent work [36, 44] has shown that
Revised: Finally, recent work [37, 45] has shown that

Original: Bayesian neural networks (BNN) are often better calibrated than standard neural networks.
Revised: Bayesian neural networks (BNN) are often better calibrated than standard neural networks.

Original: If transformers and Bayesian deep learning are both so popular, why have we not seen any successful
Revised: If transformers and Bayesian deep learning are both so popular, why have we not seen any successful

Original: Bayesian transformer models?
Revised: Bayesian transformer models?

Original: By attempting to implement such models, we make the following  contributions: (i)
Revised: By attempting to implement such models, we make the following It’s contributions: (i)

Original: We ﬁnd that weight space inference in transformers does not provide any improve- ments over a model trained by maximum likelihood.
Revised: We ﬁnd that weight space inference in transformers does not provide any improvements over a model trained by maximum likelihood.

Original: (ii) We show that the prior is at least partiallyat fault for this.
Revised: (ii) We show that the prior is at least partially at fault for this.

Original: (iii)
Revised: (iii)

Original: We propose to perform inference on the attention weights rather than on theparameters, and present a novel variational method for this using the Dirichlet distribution.
Revised: We propose to perform inference on the attention weights rather than on the parameters, and present a novel variational method for this using the Dirichlet distribution.

Original:  Background 
Revised:  Background 2.1 Bayesian deep learning

Original: 2.1 Bayesian deep learning

Original: Bayesian inference computes the posterior distribution as
Revised: Bayesian inference computes the posterior distribution as

Original: [Equation]
Revised: [Equation]

Original: with neural network parameters θ , training data { ( x i , y i ) } Ni =1 , likelihood function P ( y 1: N | θ, x 1: N ) ,prior P ( θ ) , and evidence P ( y 1: N | x 1: N ) .
Revised: with neural network parameters θ , training data { ( x i , y i ) } Ni =1 , likelihood function P ( y 1: N | θ, x 1: N ) , prior P ( θ ) , and evidence P ( y 1: N | x 1: N ) .

Original: The predictive distribution of a new target y ∗ given x ∗ isthen obtained by
Revised: The predictive distribution of a new target y ∗ given x ∗ is then obtained by

Original: [Equation]
Revised: [Equation]

Original: Applied to neural networks, both Eq.
Revised: Applied to neural networks, both Eq.

Original: (1) and Eq.
Revised: (1) and Eq.

Original: (2) are intractable and need to be estimated usingapproximate inference methods, such as variational inference or Monte Carlo sampling.
Revised: (2) are intractable and need to be estimated using approximate inference methods, such as variational inference or Monte Carlo sampling.

Original: 2.2 Bayesian neural network weight space inference
Revised: 2.2 Bayesian neural network weight space inference

Original: [Figure]
Revised: [Figure]

Original: The most commonly used weight space inference methods in BNNs are variational inference(VI) [3, 16, 23,40, 41, 43], the Laplace method[12, 31, 36], and Markov Chain Monte Carlo(MCMC) [8, 45, 58].
Revised: The most commonly used weight space inference methods in BNNs are variational inference (VI) [3, 16, 23, 41, 42, 44], the Laplace method [12, 31, 37], and Markov Chain Monte Carlo (MCMC) [8, 46, 59].

Original: While MCMC methodsdirectly sample from the (unnormalized) posterior, VI and Laplace approximate the posteriorby another distribution.
Revised: While MCMC methods directly sample from the (unnormalized) posterior, VI and Laplace approximate the posterior by another distribution.

Original: As MCMC methods areexpensive computationally and memory-wise,we restrict our focus to VI and Laplace.
Revised: As MCMC methods are expensive computationally and memory-wise, we restrict our focus to VI and Laplace.

Original: Whenapplying these methods to transformers, we ﬁndthat weight space inference fails to improve dataﬁt, calibration and predictive uncertainty compared to a model trained by likelihood maximization (see Fig. 1).
Revised: When applying these methods to transformers, we ﬁnd that weight space inference fails to improve data ﬁt, calibration and predictive uncertainty com- pared to a model trained by likelihood maximization (see Fig. 1).

Original: 2.3 Empirical weight study
Revised: 2.3 Empirical weight study

Original: To understand why weight space inference fails, we study the empirical weight distribution oftransformers trained with stochastic gradient descent (SGD), hoping to obtain better priors.
Revised: To understand why weight space inference fails, we study the empirical weight distribution of transformers trained with stochastic gradient descent (SGD), hoping to obtain better priors.

Original: We followthe framework proposed by Fortuin et al.
Revised: We follow the framework proposed by Fortuin et al.

Original: We ﬁrst examine the marginal weight distributionwhere we especially study the tailedness and modality.
Revised: We ﬁrst examine the marginal weight distribution where we especially study the tailedness and modality.

Original: We also identify the best-ﬁtting distributionand its parameters within the Gaussian, Student, Logistic, Cauchy, and Laplace families.
Revised: We also identify the best-ﬁtting distribution and its parameters within the Gaussian, Student, Logistic, Cauchy, and Laplace families.

Original: Furthermore,we investigate the correlation among layer weights by comparing the empirical covariance matrixand the distribution of off-diagonal covariance elements against samples from an isotropic Gaussian.
Revised: Furthermore, we investigate the correlation among layer weights by comparing the empirical covariance matrix and the distribution of off-diagonal covariance elements against samples from an isotropic Gaussian.

Original:  Methods 
Revised:  Methods 3.1 Variational attention

Original: 3.1 Variational attention

Original: As as alternative to weight-space inference in transformers, we propose to treat self-attention weightsas random variables and approximate their posterior distribution using VI.
Revised: As as alternative to weight-space inference in transformers, we propose to treat self-attention weights as random variables and approximate their posterior distribution using VI.

Original: Previous attention weightinference methods focus on sampling [1], while others explicitly parameterize the attention weightswith a particular distribution [2, 13, 17].
Revised: Previous attention weight inference methods focus on sampling [1], while others explicitly parameterize the attention weights with a particular distribution [2, 13, 17].

Original: Parameters of explicitly reparameterizable distributions suchas the Gaussian [2], Weibull, and Lognormal distributions [17] are learned via VI, while others suchas the Dirichlet [13] require using REINFORCE gradient estimators [53].
Revised: Parameters of explicitly reparameterizable distributions such as the Gaussian [2], Weibull, and Lognormal distributions [17] are learned via VI, while others such as the Dirichlet [13] require using REINFORCE gradient estimators [54].

Original: We implement two baselines for our comparison: 1.
Revised: We implement two baselines for our comparison: 1.

Original: Gaussian attention where the attention logitsare parameterized with a Gaussian distribution and parameters are inferred via VI and 2.
Revised: Gaussian attention where the attention logits are parameterized with a Gaussian distribution and parameters are inferred via VI and 2.

Original: DD , a datadependent conﬁguration where the variational variances of the Gaussian distribution are amortized inorder to support input-dependent (i.e., heteroscedastic ) uncertainties.
Revised: DD , a data [Table] dependent conﬁguration where the variational variances of the Gaussian distribution are amortized in order to support input-dependent (i.e., heteroscedastic ) uncertainties.

Original: 3.2 Implicitly reparameterized Dirichlet attention
Revised: 3.2 Implicitly reparameterized Dirichlet attention

Original: Alternatively, we propose to directly parameterize the attention weights of each position i by a
Revised: Alternatively, we propose to directly parameterize the attention weights of each position i by a

Original: Dirichlet distribution with parameter α = a A i , where a is the sharpness parameter and A i the [Table] of the scaled dot-product attention weights.
Revised: Dirichlet distribution with parameter α = a A i , where a is the sharpness parameter and A i the i th row of the scaled dot-product attention weights.

Original: We then infer a using VI.
Revised: We then infer a using VI.

Original: Samples are obtained by drawingfrom independent Gamma distributions
Revised: Samples are obtained by drawing from independent Gamma distributions

Original: X k ∼ Gamma ( α k , 1) and normalizing ( (cid:80)
Revised: X k ∼ Gamma ( α k , 1) and normalizing ( (cid:80)

Original: Kk =1 X k ) − 1 X ∼
Revised: Kk =1 X k ) − 1 X ∼

Original: Dirichlet ( α ) .
Revised: Dirichlet ( α ) .

Original: We further use contextual Gamma priors such that ˆ α ∝ A
Revised: We further use contextual Gamma priors such that ˆ α ∝ A

Original: i , yielding an analytical KLdivergence as done by Joo et al.
Revised: i , yielding an analytical KL divergence as done by Joo et al.

Original: To obtain gradients of a Gamma random variable with respectto α , we use the implicit gradient reparametrization [18]:
Revised: To obtain gradients of a Gamma random variable with respect to α , we use the implicit gradient reparametrization [18]:

Original: [Equation]
Revised: [Equation]

Original: ) where q α ( z ) is the Gamma density function and F ( z | α ) its CDF.
Revised: (3) where q α ( z ) is the Gamma density function and F ( z | α ) its CDF.

Original: Like Gaussian attention, we considera variation where the sharpness parameter depends on the input, referred to as data dependent.
Revised: Like Gaussian attention, we consider a variation where the sharpness parameter depends on the input, referred to as data dependent.

Original:  Experiments
Revised:  Experiments

Original: We run experiments using the transformer [57] and vision transformer [15] on MNIST imageclassiﬁcation [39], Universal Dependencies part-of-speech (POS) tagging [46] and on syntheticdatasets (M1, M2).
Revised: We run experiments using the transformer [58] and vision transformer [15] on MNIST image classiﬁcation [40], Universal Dependencies part-of-speech (POS) tagging [47] and on synthetic datasets (M1, M2).

Original: We evaluate our models using test log-likelihood, predicted variance meansquared error and the expected mean square error on the synthetic dataset.
Revised: We evaluate our models using test log-likelihood, predicted variance mean squared error and the expected mean square error on the synthetic dataset.

Original: The test log-likelihood,accuracy, F1-score, and expected calibration error (ECE) [26] are used for experiments on the POStagging and MNIST datasets.
Revised: The test log-likelihood, accuracy, F1-score, and expected calibration error (ECE) [26] are used for experiments on the POS tagging and MNIST datasets.

Original: We compare the results obtained by our methods with a transformer(MLE) and an ensemble of 30 transformers both trained by maximum likelihood.
Revised: We compare the results obtained by our methods with a transformer (MLE) and an ensemble of 30 transformers both trained by maximum likelihood.

Original: Further details aregiven in Appendix A.2.4.1 Result 1: Weight-space inference does not improve over MLE
Revised: Further details are given in Appendix A.2. 4.1 Result 1: Weight-space inference does not improve over MLE

Original: Different posteriors do not help.
Revised: Different posteriors do not help.

Original: We ﬁnd that all weight-space VI methods are outperformed byboth maximum-likelihood baselines with respect to all metrics and on all datasets (see Table 1).
Revised: We ﬁnd that all weight-space VI methods are outperformed by both maximum-likelihood baselines with respect to all metrics and on all datasets (see Table 1).

Original: Interestingly, changing the posterior distribution does not signiﬁcantly inﬂuence the performance,considering the large gap between the scores of the VI methods and baselines (see also Table 3 in theappendix).
Revised: Interestingly, changing the posterior distribution does not signiﬁcantly inﬂuence the performance, considering the large gap between the scores of the VI methods and baselines (see also Table 3 in the appendix).

Original: Furthermore, no variational posterior systematically outperforms the others.
Revised: Furthermore, no variational posterior systematically outperforms the others.

Original: Linearized Laplace inference (either on all parameters or just the ﬁnal layer) shows much better resultsthan VI.
Revised: Linearized Laplace inference (either on all parameters or just the ﬁnal layer) shows much better results than VI.

Original: However, it still underperforms our baselines.
Revised: However, it still underperforms our baselines.

Original: Finally, even concrete dropout improves over
Revised: Finally, even concrete dropout improves over

Original: VI and Laplace inference and is more competitive with our baselines.
Revised: VI and Laplace inference and is more competitive with our baselines.

Original: The prior is (at least partially) at fault.
Revised: The prior is (at least partially) at fault.

Original: In our attempt to understand the poor performance of weight-space VI in transformers, we conduct an empirical weight distribution study.
Revised: In our attempt to understand the poor performance of weight-space VI in transformers, we conduct an empirical weight distribution study.

Original: We ﬁnd that the marginal weight distributions are essentially uni-modal, except for someembedding and projection layers which tend to have two or three less signiﬁcant modes.
Revised: We ﬁnd that the marginal weight distributions are essentially uni-modal, except for some embedding and projection layers which tend to have two or three less signiﬁcant modes.

Original: [Table] Furthermore, other than a decrease in tailedness(high degree of freedom) of the last layer, no recurrent pattern in the tailedness of the weight distribution appears across the considered datasets(Fig. 4 in the appendix).
Revised: Furthermore, other than a decrease in tailedness (high degree of freedom) of the last layer, no recurrent pattern in the tailedness of the weight distribution appears across the considered datasets (Fig. 4 in the appendix).

Original: Likewise, no single distribution seems to universally ﬁt the empiricaldistributions of the weights across all datasets(see Fig.
Revised: Likewise, no single distribution seems to universally ﬁt the empirical distributions of the weights across all datasets (see Fig.

Original: 5 and Fig.
Revised: 5 and Fig.

Original: 6 in the appendix).
Revised: 6 in the appendix).

Original: Thissuggests that the shape of the weight distribution strongly depends on the considered dataset.
Revised: This suggests that the shape of the weight distribution strongly depends on the considered dataset.

Revised: [Table]

Original: Usingthe observations from this weight distribution study, we choose more appropriate ("improved") priors.
Revised: Using the observations from this weight distribution study, we choose more appropriate ("improved") priors.

Original: Table 2 shows systematic likelihood improvements.
Revised: Table 2 shows systematic likelihood improvements.

Original: Moreover, we ﬁnd that the performance of VIcritically depends on the prior parameters (see Fig.
Revised: Moreover, we ﬁnd that the performance of VI critically depends on the prior parameters (see Fig.

Original: 3 in the appendix).
Revised: 3 in the appendix).

Original: 2 Results 2: Variational attention is better than weight-space inference
Revised: 4.2 Results 2: Variational attention is better than weight-space inference

Original: Dirichlet attention works well.
Revised: Dirichlet attention works well.

Original: Unlike weight-space inference, we ﬁnd that inference on theattention weights works competitively.
Revised: Unlike weight-space inference, we ﬁnd that inference on the attention weights works competitively.

Original: Indeed, Dirichlet attention strongly outperforms our baselinesin terms of likelihood on the synthetic data and lies between both baselines on the POS tagging and
Revised: Indeed, Dirichlet attention strongly outperforms our baselines in terms of likelihood on the synthetic data and lies between both baselines on the POS tagging and

Original: MNIST (Table 1).
Revised: MNIST (Table 1).

Original: However, the data dependent conﬁguration does not systematically outperformits standard counterpart (see Table 5 in the appendix).
Revised: However, the data dependent conﬁguration does not systematically outperform its standard counterpart (see Table 5 in the appendix).

Original: Moreover, Dirichlet attention outperforms
Revised: Moreover, Dirichlet attention outperforms

Original: Gaussian attention in terms of log-likelihood on the toy data and POS tagging, but not on MNIST.
Revised: Gaussian attention in terms of log-likelihood on the toy data and POS tagging, but not on MNIST.

Original: Variational attention leads to more consistent prior entropies.
Revised: Variational attention leads to more consistent prior entropies.

Original: While investigating the entropyof the predictive distribution when sampling weights from the priors, we ﬁnd that non-improvedpriors yield highly variable entropy distributions, ranging from low values around 1 to higher valuesaround 2 .
Revised: While investigating the entropy of the predictive distribution when sampling weights from the priors, we ﬁnd that non-improved priors yield highly variable entropy distributions, ranging from low values around 1 to higher values around 2 .

Original: 3 bits.
Revised: 3 bits.

Original: However, when sampling from improved priors selected by our weight distributionanalysis, the entropy distribution concentrates very strongly around a high value of 2 . 3 bits.
Revised: However, when sampling from improved priors selected by our weight distribution analysis, the entropy distribution concentrates very strongly around a high value of 2 . 3 bits.

Original: Thissame behavior is observed when sampling from the Gaussian and our proposed Dirichlet attention.
Revised: This same behavior is observed when sampling from the Gaussian and our proposed Dirichlet attention.

Original: This is desirable as the prior predictive should show high uncertainty in function space.
Revised: This is desirable as the prior predictive should show high uncertainty in function space.

Original: [Figure]
Revised: [Figure]

Original: and using variational attention both lead to more consistently high entropies.

Original:  Related work
Revised:  Related work

Original: Bayesian Transformers.
Revised: Bayesian Transformers.

Original: Previous attempts to make the transformer Bayesian have used VI toperform inference on a subset of layers [56, 62].
Revised: Previous attempts to make the transformer Bayesian have used VI to perform inference on a subset of layers [57, 64].

Original: While both methods claim state-of-the-art performance on their respective benchmarks, Tran et al.
Revised: While both methods claim state-of-the-art performance on their respective benchmarks, Tran et al.

Original: [56] do not provide any quantitative resultsand Xue et al.
Revised: [57] do not provide any quantitative results and Xue et al.

Original: [62] initialize their priors to a maximum estimate of the weights which is not strictly
Revised: [64] initialize their priors to a maximum estimate of the weights which is not strictly

Original: Bayesian.
Revised: Bayesian.

Original: Alternatively, Fan et al.
Revised: Alternatively, Fan et al.

Original: [17] parameterize the attention logits of a transformer by a Gaussian distribution and ﬁnetune the deterministic self-attention of language models pretrained on largecorpora.
Revised: [17] parameterize the attention logits of a transformer by a Gaussian distribution and ﬁnetune the deterministic self-attention of language models pretrained on large corpora.

Original: They however only consider ﬁnetuning and not full training using variational attention.
Revised: They however only consider ﬁnetuning and not full training using variational attention.

Original: Orthogonally, Martin et al.
Revised: Orthogonally, Martin et al.

Original: [42] consider attention keys, queries, values, and weights as unobservedrandom variables and use sequential Monte Carlo methods to sample them.
Revised: [43] consider attention keys, queries, values, and weights as unobserved random variables and use sequential Monte Carlo methods to sample them.

Original: Bayesian neural network inference.
Revised: Bayesian neural network inference.

Original: BNN inference has recently advanced in terms of VI methodswith more expressive posteriors [16,40, 41, 43, 55], more efﬁcient inference [22, 23, 54], and greaterstability [35, 59].
Revised: BNN inference has recently advanced in terms of VI methods with more expressive posteriors [16, 41, 42, 44, 56], more efﬁcient inference [22, 23, 55], and greater stability [36, 60].

Original: Likewise, the Laplace inference for BBNs has improved in scalability using further
Revised: Likewise, the Laplace inference for BBNs has improved in scalability using further

Original: GGN approximations [30, 31, 36, 50, 51] and sub-network inference [12, 36].
Revised: GGN approximations [30, 31, 37, 51, 52] and sub-network inference [12, 37].

Original: Orthogonally, MCMCmethods for BNNs have been improved [20, 24, 60, 64], better BNN priors have been studied [19, 21],and even deep ensembles [38] have been cast as approximate inference [9–11, 48, 49].
Revised: Orthogonally, MCMC methods for BNNs have been improved [20, 24, 61, 66], better BNN priors have been studied [19, 21], and even deep ensembles [39] have been cast as approximate inference [9–11, 32, 49, 50, 62].

Original:  Conclusion
Revised:  Conclusion

Original: We have shown that weight space inference in Bayesian transformers does not work well, regardlessof the choice of posterior.
Revised: We have shown that weight space inference in Bayesian transformers does not work well, regardless of the choice of posterior.

Original: We also found that choosing priors according to an empirical weightdistribution analysis improved the performance, suggesting that priors are at least partially at fault.
Revised: We also found that choosing priors according to an empirical weight distribution analysis improved the performance, suggesting that priors are at least partially at fault.

Original: However, we have not found the right priors to make the method competitive.
Revised: However, we have not found the right priors to make the method competitive.

Original: Moreover, we foundevidence that naïve weight-space priors lead to low prior predictive entropy, and therefore do notreﬂect our true beliefs about the output distribution.
Revised: Moreover, we found evidence that naïve weight-space priors lead to low prior predictive entropy, and therefore do not reﬂect our true beliefs about the output distribution.

Original: In order to move closer to the function-spacedistribution, we suggested to perform inference on the attention weights rather than on parameters.
Revised: In order to move closer to the function-space distribution, we suggested to perform inference on the attention weights rather than on parameters.

Original: We proposed a novel method based on the implicit reparameterization of the Dirichlet distribution toapply variational inference on the attention weights, which performed competitively with respect toour baselines.
Revised: We proposed a novel method based on the implicit reparameterization of the Dirichlet distribution to apply variational inference on the attention weights, which performed competitively with respect to our baselines.

Original: A AppendixA.1 Additional experimental resultsA.1.1 Weight space inference
Revised: A Appendix A.1 Additional experimental results A.1.1 Weight space inference

Original: We report in Table 3 the full weight space inference results table discussed in Section 4.1.
Revised: We report in Table 3 the full weight space inference results table discussed in Section 4.1.

Original: [Figure]
Revised: [Figure]

Original: A.1.2 Sub-network variational inference
Revised: A.1.2 Sub-network variational inference

Original: In addition to performing inference on the entire set of model parameters presented in Section 4.
Revised: In addition to performing inference on the entire set of model parameters presented in Section 4.

Original: [Figure]
Revised: [Figure]

Original: we experiment sub-network variational inference.
Revised: we experiment sub-network variational inference.

Original: We ﬁnd that this method performs better than fullinference while still under-performing our baselines.
Revised: We ﬁnd that this method performs better than full inference while still under-performing our baselines.

Original: Interestingly, we observe the same behavior asin full network VI, where the posterior distribution does not signiﬁcantly change the result
Revised: Interestingly [Figure] we observe the same behavior as in full network VI, where the posterior distribution does not signiﬁcantly change the result

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Table]
Revised: [Table]

Original: A.1.3 Variational attention
Revised: A.1.3 Variational attention

Original: We report in Table 5 the full variational attention results table discussed in Section 4.2 including thedata dependent conﬁgurations.
Revised: We report in Table 5 the full variational attention results table discussed in Section 4.2 including the data dependent conﬁgurations [Table]

Original: [Table]
Revised: [Table]

Original: A.1.
Revised: A.1.4 Likelihood sensitivity to the prior

Original: 4 Likelihood sensitivity to the prior

Original: As discussed in Section 4.1, we ﬁnd that the model test likelihood is very sensitive to the choice ofprior (see Figure 3).
Revised: As discussed in Section 4. [Table] we ﬁnd that the model test likelihood is very sensitive to the choice of prior (see Figure 3).

Original: [Figure]
Revised: [Figure]

Original: A.1.5 Weight distribution tailedness
Revised: A.1.5 Weight distribution tailedness

Original: We here present more detailed results regarding the tailedness of the weight distribution.
Revised: We here present more detailed results regarding the tailedness of the weight distribution.

Original: We ﬁnd thatno obvious patterns in the thickness of tails exist across the considered dataset (Figure 4).
Revised: We ﬁnd that no obvious patterns in the thickness of tails exist across the considered dataset (Figure 4).

Original: Q-Q plots ofthe empirical weight distribution against light-tailed Gaussian and heavier tailed Laplace distributionprovide more evidence of this phenomenon.
Revised: Q-Q plots of the empirical weight distribution against light-tailed Gaussian and heavier tailed Laplace distribution provide more evidence of this phenomenon.

Original: Indeed, the weight distributions of transformers trainedon MNIST and POS tagging are well ﬁtted by a Gaussian
Revised: Indeed, the weight distributions of transformers trained on MNIST and POS tagging are well ﬁtted by a Gaussian

Original: [Figure]
Revised: [Figure]

Original: while transformers trained on M1 and
Revised: while transformers trained on M1 and

Original: M2 toy datasets are well ﬁtted by a Laplace (Figure 5 and 6).
Revised: M2 toy datasets are well ﬁtted by a Laplace (Figure 5 and 6).

Original: This suggest that not one universal[Figure] distribution ﬁts the empirical distribution of the weights across all datasets[Figure]
Revised: This suggest that not one universal distribution ﬁts the empirical distribution of the weights across all datasets

Original: u

Original: [Figure]

Original: A tt

Original: [Figure]

Original: li n e a

Original: [Figure]

Original: A tt

Original: [Figure]

Original: li n e a

Original: [Figure]

Original: A tt

Original: [Figure]

Original: ee

Original: [Figure]

Original: w a

Original: [Figure]

Original: A tt

Original: [Figure]
Revised: [Figure]

Original: e

Original: [Figure]
Revised: [Figure]

Original: A tt

Original: [Figure]

Original: u

Original: [Figure]

Original: e

Original: [Figure]
Revised: [Figure]

Original: A tt

Original: [Figure]

Original: 2 v a l u

Original: [Figure]

Original: A tt

Original: [Figure]

Original: li n e

Original: [Figure]

Original: A tt

Original: [Figure]

Original: n e a

Original: [Figure]

Original: [Figure]

Original: [Figure]

Original: A.1.6 Off-diagonal covariance value distributions
Revised: A.1.6 Off-diagonal covariance value distributions

Original: We here present histograms of the off-diagonal empirical covariance elements.
Revised: We here present histograms of the off-diagonal empirical covariance elements.

Original: Covariance valueshave small magnitude and concentrate strongly in distribution around 0.
Revised: Covariance values have small magnitude and concentrate strongly in distribution around 0.

Original: Depending on the dataset,off-diagonal covariance elements are slightly larger than samples from an isotropic Gaussian asshown in Figure 7 and10.
Revised: Depending on the dataset, off-diagonal covariance elements are slightly larger than samples from an isotropic Gaussian as shown in Figure 7 and

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: 1.7 Concrete dropout high
Revised: A.1.7 Concrete dropout high/low entropy samples

Original: low entropy samples

Original: Here we present samples from the MNIST test set which have high (respectively low) predictivedistribution entropy.
Revised: Here we present samples from the MNIST test set which have high (respectively low) predictive distribution entropy.

Original: By visual inspection, we ﬁnd that high predictive entropy samples have anambiguous labeling while low entropy predictive sample labels are much easier to identify
Revised: By visual inspection, we ﬁnd that high predictive entropy samples have an ambiguous labeling while low entropy predictive sample labels are much easier to identify.

Original: [Figure]
Revised: [Equation]

Original: 1 st 2 nd 3 rd 4 th (e) Highest entropy 1 st 2 nd 3 rd 4 th (j) [Figure]
Revised: j) Lowest entropy Dropout : highest/lowest entropy MNIST test samples

Original: A.2 Implementation details
Revised: A.2 Implementation details

Original: In the following section, we present more detailed implementation speciﬁcations.
Revised: In the following section, we present more detailed implementation speciﬁcations.

Original: A.2.1 Architecture & Datasets
Revised: A.2.1 Architecture & Datasets

Original: Toy Data We generate synthetic data from the two following models:
Revised: Toy Data We generate synthetic data from the two following models:

Revised: • Model M1:

Revised: X t +1 = (cid:80) 4 i =0 0 .

Original: [Equation]
Revised: 2 cos(0 . 4 πiX t + 1 / ( i + 1)) + √ 0 . 5 (cid:15) t +

Original: • Model M2: X t +1 =
Revised: • Model M2: X t +1 =

Original: (cid:80) 4 i =

Revised: (cid:80) 4 i =

Original: (cid:80) 

Revised: (cid:80)

Original: j =0 0 .
Revised: j =0 0 .

Original: 5 cos(0 .
Revised: 5 cos(0 .

Original: 8 πjX t − i ) + √ 0 .
Revised: 8 πjX t − i ) + √ 0 .

Original: 1 (cid:15) t [-+
-]
Revised: 1 (cid:15) t [-+
-]+

Original: where ( X i ) 4 i =0 and ( (cid:15) t )
Revised: where ( X i ) 4 i =0 and ( (cid:15) t )

Original: Tt =1 are i.i.d standard Gaussian random variables.
Revised: Tt =1 are i.i.d standard Gaussian random variables.

Original: We use 800 training,80 validation and 80 testing sequences of length 24.
Revised: We use 800 training, 80 validation and 80 testing sequences of length 24.

Original: For experiments with Toy Data, we use thetransformer from [57] with GeLU activation [27], a unique attention block, one attention head anda hidden size of 64.
Revised: For experiments with Toy Data, we use the transformer from [58] with GeLU activation [27], a unique attention block, one attention head and a hidden size of 64.

Original: We train until convergence for 100 epochs.
Revised: We train until convergence for 100 epochs.

Original: We draw 30 posterior samples tocompute the predictive distribution and KL divergence.
Revised: We draw 30 posterior samples to compute the predictive distribution and KL divergence.

Original: We evaluate our model with the test datalog-likelihood, mean squared error of the predicted variance and the expected mean square errorgiven the previous value X t .
Revised: We evaluate our model with the test data log-likelihood, mean squared error of the predicted variance and the expected mean square error given the previous value X t .

Original: Part of Speech Tagging We use the English split of the Part of Speech tagging dataset from the
Revised: Part of Speech Tagging We use the English split of the Part of Speech tagging dataset from the

Original: Universal Dependencies v1.2 corpus [46].
Revised: Universal Dependencies v1.2 corpus [47].

Original: This dataset contains 204’586 train, 25’148 validationand 25’096 test tokens.
Revised: This dataset contains 204’586 train, 25’148 validation and 25’096 test tokens.

Original: We use a maximum length of 40 tokens, pad the shorter sequences and splitthe sentences which exceed the text length.
Revised: We use a maximum length of 40 tokens, pad the shorter sequences and split the sentences which exceed the text length.

Original: For experiments with this dataset, we use a transformerfrom [57] with GeLU activation, a unique attention block, one attention head and a hidden size of 32.
Revised: For experiments with this dataset, we use a transformer from [58] with GeLU activation, a unique attention block, one attention head and a hidden size of 32.

Original: We train until convergence for 100 epochs.
Revised: We train until convergence for 100 epochs.

Original: We draw 10 posterior samples to compute the predictivedistribution and KL divergence.
Revised: We draw 10 posterior samples to compute the predictive distribution and KL divergence.

Original: We evaluate our model with the token level test data log-likelihood,token level accuracy, token level F1-score and expected calibration error (ECE) [26].
Revised: We evaluate our model with the token level test data log-likelihood, token level accuracy, token level F1-score and expected calibration error (ECE) [26].

Original: MNIST We experiment with the MNIST image classiﬁcation dataset [39].
Revised: MNIST We experiment with the MNIST image classiﬁcation dataset [40].

Original: We split the originaldataset into 48’000 training, 12’000 validation and 9’984 testing samples.
Revised: We split the original dataset into 48’000 training, 12’000 validation and 9’984 testing samples.

Original: For experiments withthis dataset, we use a Vision transformer from [15] with GeLU activation, two attention block, oneattention head, a hidden size of 32 and a patch size of 4.
Revised: For experiments with this dataset, we use a Vision transformer from [15] with GeLU activation, two attention block, one attention head, a hidden size of 32 and a patch size of 4.

Original: We train until convergence for 150 epochs.
Revised: We train until convergence for 150 epochs.

Original: We draw 10 posterior samples to compute the predictive distribution and KL divergence.
Revised: We draw 10 posterior samples to compute the predictive distribution and KL divergence.

Original: We evaluateour model with the test data log-likelihood, accuracy, F1-score and ECE.
Revised: We evaluate our model with the test data log-likelihood, accuracy, F1-score and ECE.

Original: A.2.2 Training setup
Revised: A.2.2 Training setup

Original: Our models are trained using the Adam optimizer [34] with the triangular learning rate schedule from[57].
Revised: Our models are trained using the Adam optimizer [35] with the triangular learning rate schedule from [58].

Original: Our weight distribution experiments are conducted using 900 transformers trained by likelihoodmaximization with SGD as done in [21].
Revised: Our weight distribution experiments are conducted using 900 transformers trained by likelihood maximization with SGD as done in [21].

Original: A.2.3 Software packages
Revised: A.2.3 Software packages

Original: We implement and train our models using the JAX [4] and Haiku [28] Python libraries.
Revised: We implement and train our models using the JAX [4] and Haiku [28] Python libraries.