Pathologies in Priors and Inference for Bayesian Transformers
Abstract
In recent years, the transformer has established itself as a workhorse in many applications ranging from natural language processing to reinforcement learning.
Similarly, Bayesian deep learning has become the gold-standard for uncertainty estimation in safety-critical applications, where robustness and calibration are crucial.
Surprisingly, no successful attempts to improve transformer models in terms of predictive uncertainty using Bayesian inference exist.
In this work, we study this curiously underpopulated area of Bayesian transformers.
We find that weight-space inference in transformers does not work well, regardless of the approximate posterior.
We also find that the prior is at least partially at fault, but that it is very hard to find well-specified weight priors for these models.
We hypothesize that these problems stem from the complexity of obtaining a meaningful mapping from weight-space to function-space distributions in the transformer.
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.
We find that this proposed method performs competitively with our baselines.
1. Introduction
The transformer (Vaswani et al., 2017) is a deep learning architecture commonly used to process sequences of data, such as text.
Thanks to multi-head self-attention, the transformer builds contextual embeddings by capturing the relationships between the sequence elements.
While being most famous for their state-of-the-art performance in natural language processing (Brown et al., 2020; Devlin et al., 2019), transformers are also used in computer vision (Chen et al., 2020;
Dosovitskiy et al., 2021; Jiang et al., 2021; Strudel et al., 2021; Wu et al., 2020), reinforcement learning (Chen et al., 2021; Kumar et al., 2020), as well as audio (Gong et al., 2021; Huang et al., 2018; Payne, 2019) and video (Yan et al., 2021) processing, yielding impressive results.
The Bayesian learning paradigm provides a theoretical framework to obtain predictive uncertainty, select the optimal model, and improve its calibration.
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.
Finally, recent work (Kristiadi et al., 2020; Mitros and Namee, 2019) has shown that Bayesian neural networks (BNN) are often better calibrated than standard neural networks.
If transformers and Bayesian deep learning are both so popular, why have we not seen any successful Bayesian transformer models?
By attempting to implement such models, we make the following contributions: (i)
We find that weight space inference in transformers does not provide any improvements over a model trained by maximum likelihood.
(ii) We show that the prior is at least partially at fault for this.
(iii)
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.
2. Background 2.1. Bayesian deep learning
Bayesian inference computes the posterior distribution as
[Equation]
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 ) .
The predictive distribution of a new target y ∗ given x ∗ is then obtained by
[Equation]
Applied to neural networks, both Eq.
(1) and Eq.
(2) are intractable and need to be estimated using approximate inference methods, such as variational inference or Monte Carlo sampling.
2.2. Bayesian neural network weight space inference
[Figure]
MLE captures the mean of the generative process well (MSE: 0.66), but VI does not (MSE: 9.58).
The most commonly used weight space inference methods in BNNs are variational inference (VI) (Blundell et al., 2015; Dusenberry et al., 2020; Gal et al., 2017; Louizos and
Welling, 2016, 2017; Mishkin et al., 2019), the Laplace method (Daxberger et al., 2021;
Immer et al., 2021b; Kristiadi et al., 2020), and Markov Chain Monte Carlo (MCMC) (Chen et al., 2014; Neal, 2012; Welling and
Teh, 2011).
While MCMC methods directly sample from the (unnormalized) posterior,
VI and Laplace approximate the posterior by another distribution.
As MCMC methods are expensive computationally and memory-wise, we restrict our focus to VI and Laplace.
When applying these methods to transformers, we find that weight space inference fails to improve data fit, calibration and predictive uncertainty compared to a model trained by likelihood maximization (see Fig. 1).
2.3. Empirical weight study
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.
We follow the framework proposed by Fortuin et al.
(2021b).
We first examine the marginal weight distribution where we especially study the tailedness and modality.
We also identify the best-fitting distribution and its parameters within the Gaussian, Student, Logistic,
Cauchy, and Laplace families.
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.
3. Methods 3.1. Variational attention
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.
Previous attention weight inference methods focus on sampling (An et al., 2020), while others explicitly parameterize the attention weights with a particular distribution (Bahuleyan et al., 2018;
Deng et al., 2018; Fan et al., 2020).
Parameters of explicitly reparameterizable distributions such as the Gaussian Bahuleyan et al.
(2018), Weibull, and Lognormal distributions (Fan et al., 2020) are learned via VI, while others such as the Dirichlet (Deng et al., 2018) require using REINFORCE gradient estimators (Sutton et al., 2000).
We implement two baselines for our comparison: 1.
Gaussian attention where the attention logits are parameterized with a Gaussian distribution and parameters are inferred via VI and 2.
DD , a data dependent configuration where the variational variances of the
Gaussian distribution are amortized in order to support input-dependent (i.e., heteroscedastic ) uncertainties.
3.2. Implicitly reparameterized Dirichlet attention
Alternatively, we propose to directly parameterize the attention weights of each position i by a 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.
We then infer a using VI.
Samples are obtained by drawing from independent Gamma distributions
X k ∼ Gamma ( α k , 1) and normalizing ( (cid:80)
Kk =1 X k ) − 1 X ∼ Dirichlet ( α ) .
We further use contextual Gamma priors such that ˆ α ∝ A
i , yielding an analytical KL divergence as done by Joo et al.
To obtain gradients of a Gamma random variable with respect to α , we use the implicit gradient reparametrization (Figurnov et al., 2019):
[Equation]
(3) where q α ( z ) is the Gamma density function and F ( z | α ) its CDF.
Like Gaussian attention, we consider a variation where the sharpness parameter depends on the input, referred to as data dependent. [Table]
4. Experiments
We run experiments using the transformer Vaswani et al.
(2017) and vision transformer
Dosovitskiy et al.
(2021) on MNIST image classification (LeCun et al., 2010), Universal
Dependencies part-of-speech (POS) tagging (Nivre et al., 2015) and on synthetic datasets (M1, M2).
We evaluate our models using test log-likelihood, predicted variance mean squared error and the expected mean square error on the synthetic dataset.
The test log-likelihood, accuracy, F1-score, and expected calibration error (ECE) (Guo et al., 2017) are used for experiments on the POS tagging and MNIST datasets.
We compare the results obtained by our methods with a transformer (MLE) and an ensemble of 30 transformers both trained by maximum likelihood.
Further details are given in Appendix A.2.
4.1. Result 1: Weight-space inference does not improve over MLE Different posteriors do not help.
We find 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).
Interestingly, changing the posterior distribution does not significantly influence the performance, considering the large gap between the scores of the VI methods and baselines (see also Table 3 in the appendix).
Furthermore, no variational posterior systematically outperforms the others.
Linearized Laplace inference (either on all parameters or just the final layer) shows much better results than VI.
However, it still underperforms our baselines.
Finally, even concrete dropout improves over VI and Laplace inference and is more competitive with our baselines.
The prior is (at least partially) at fault.
In our attempt to understand the poor performance of weight-space VI in transformers, we conduct an empirical weight distribution study.
We find that the marginal weight distributions are essentially uni-modal, except for some embedding and projection layers which tend to have two or three less significant modes.
[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).
Likewise, no single distribution seems to universally fit the empirical distributions of the weights across all datasets (see Fig.
5 and
Fig.
6 in the appendix).
This suggests that the shape of the weight distribution strongly depends on the considered dataset.
Using the observations from this weight distribution study, we choose more appropriate ("improved") priors.
Table 2 shows systematic likelihood improvements.
Moreover, we find that the performance of VI critically depends on the prior parameters (see Fig.
3 in the appendix).
4.2. Results 2: Variational attention is better than weight-space inference Dirichlet
attention works well.
Unlike weight-space inference, we find that inference on the attention weights works competitively.
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 MNIST (Table 1).
However, the data dependent configuration does not systematically outperform its standard counterpart (see Table 5 in the appendix).
Moreover,
Dirichlet attention outperforms Gaussian attention in terms of log-likelihood on the toy data and POS tagging, but not on MNIST.
Variational attention leads to more consistent prior entropies.
While investigating the entropy of the predictive distribution when sampling weights from the priors, we find that non-improved priors yield highly variable entropy distributions, ranging from low values around 1 to higher values around 2 . 3 bits.
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.
This same behavior is observed when sampling from the Gaussian and our proposed Dirichlet attention.
This is desirable as the prior predictive should show high uncertainty in function space.
[Figure]
i
[Figure]
5. Related work
Bayesian Transformers.
Previous attempts to make the transformer Bayesian have used
VI to perform inference on a subset of layers (Tran et al., 2019; Xue et al., 2021).
While both methods claim state-of-the-art performance on their respective benchmarks, Tran et al.
(2019) do not provide any quantitative results and Xue et al.
(2021) initialize their priors to a maximum estimate of the weights which is not strictly Bayesian.
Alternatively, Fan et al.
(2020) parameterize the attention logits of a transformer by a Gaussian distribution and finetune the deterministic self-attention of language models pretrained on large corpora.
They however only consider finetuning and not full training using variational attention.
Orthogonally, Martin et al.
(2020) consider attention keys, queries, values, and weights as unobserved random variables and use sequential Monte Carlo methods to sample them.
Bayesian neural network inference.
BNN inference has recently advanced in terms of
VI methods with more expressive posteriors (Dusenberry et al., 2020; Louizos and Welling, 2016, 2017; Mishkin et al., 2019; Tomczak et al., 2020), more efficient inference (Gal and
Ghahramani, 2016; Gal et al., 2017; Swiatkowski et al., 2020), and greater stability (Kingma et al., 2015; Wen et al., 2018).
Likewise, the Laplace inference for BBNs has improved in scalability using further GGN approximations (Immer et al., 2021a,b; Kristiadi et al., 2020;
Ritter et al., 2018a,b) and sub-network inference (Daxberger et al., 2021; Kristiadi et al., 2020).
Orthogonally, MCMC methods for BNNs have been improved (Fortuin et al., 2021a;
Garriga-Alonso and Fortuin, 2021; Wenzel et al., 2020; Zhang et al., 2019), better BNN priors have been studied (Fortuin, 2021; Fortuin et al., 2021b), and even deep ensembles
Lakshminarayanan et al.
(2017) have been cast as approximate inference (Ciosek et al., 2019;
D’Angelo and Fortuin, 2021; D’Angelo et al., 2021; Izmailov et al., 2021; Pearce et al., 2018, 2020; Wilson and Izmailov, 2020).
6. Conclusion
We have shown that weight space inference in Bayesian transformers does not work well, regardless of the choice of posterior.
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.
However, we have not found the right priors to make the method competitive.
Moreover, we found evidence that naïve weight-space priors lead to low prior predictive entropy, and therefore do not reflect our true beliefs about the output distribution.
In order to move closer to the function-space distribution, we suggested to perform inference on the attention weights rather than on parameters.
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.
Appendix A. Appendix A.1. Additional experimental results A.1.1. Weight space inference
We report in Table 3 the full weight space inference results table discussed in Section 4.1. [Table] [Figure] [Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure] [Table] [Figure] [Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table] [Figure] [Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
A.1.2. Sub-network variational inference
In addition to performing inference on the entire set of model parameters presented in
Section 4.
[Figure]
we experiment sub-network variational inference.
We find that this method performs better than full inference while still under-performing our baselines.
Interestingly, we observe the same behavior as in full network VI, where the posterior distribution does not significantly change the result
[Figure]
priors vs baselines
[Table]
A.1.3. Variational attention
We report in Table 5 the full variational attention results table discussed in Section 4.
including the data dependent configurations.
[Table]
A.1.4. Likelihood sensitivity to the prior
As discussed in Section 4.
[Table]
we find that the model test likelihood is very sensitive to the choice of prior (see Figure 3)
[Table]
[Figure]
A.1.5. Weight distribution tailedness
We here present more detailed results regarding the tailedness of the weight distribution.
We find that no obvious patterns in the thickness of tails exist across the considered dataset (Figure 4).
Q-Q plots of the empirical weight distribution against light-tailed Gaussian and heavier tailed Laplace distribution provide more evidence of this phenomenon.
Indeed, the weight distributions of transformers trained on MNIST and POS tagging are well fitted by a Gaussian
[Figure]
while transformers trained on M1 and M2 toy datasets are well fitted by a
Laplace (Figure 5 and 6).
This suggest that not one universal distribution fits the empirical distribution of the weights across all datasets
[Figure]
E m b e d d i n g
A tt - 1 f e e d f o r w a r d
A tt - 1 k e y s
A tt - 1 q u e r i e s
A tt - 1 v a l u e s
A tt - 1 li n e a rA tt - 1 li n e a rP r o j e c t i o n
[Figure]
D e g
[Figure]
ee
[Figure]
f f r
[Figure]
E m b e d d i n g
A tt - 1 f e e d f o r w a r d
A tt - 1 k e y s
A tt - 1 q u e r i e s
A tt - 1 v a l u e s
A tt - 1 li n e a rA tt - 1 li n e a rA tt - 1 f e e d f o r w a r d
A tt - 2 k e y s
A tt - 2 q u e r i e s
A tt - 2 v a l u e s
A tt - 2 li n e a rA tt - 2 li n e a rP r o j e c t i o nP r
o j e c t
i o n
[Figure]
[Figure] [Figure]
A.1.6. Off-diagonal covariance value distributions
We here present histograms of the off-diagonal empirical covariance elements.
Covariance values have small magnitude and concentrate strongly in distribution around 0.
Depending on the dataset, off-diagonal covariance elements are slightly larger than samples from an isotropic Gaussian as shown in Figure 7 and
[Figure] [Figure]
[Figure] [Figure]
[Figure]
A.1.7. Concrete dropout high/low entropy samples
Here we present samples from the MNIST test set which have high (respectively low) predictive distribution entropy.
By visual inspection, we find that high predictive entropy samples have an ambiguous labeling while low entropy predictive sample labels are much easier to identify.
1 st 2 nd 3 rd 4 th (e) Highest entropy 1 st 2 nd 3 rd 4 th
[Figure]
A.2. Implementation details
In the following section, we present more detailed implementation specifications.
A.2.1. Architecture & Datasets Toy Data
We generate synthetic data from the two following models:
• Model M1:
X t +1 = (cid:80) 4 i =0 0 .
2 cos(0 .
4 πiX t + 1 / ( i + 1))
+ √ 0 .
5 ϵ t + • Model M2: X t +1 =
(cid:80) 4 i =
(cid:80)
j =0 0 .
5 cos(0 .
8 πjX t − i ) + √ 0 .
1 ϵ t +
where ( X i ) 4 i =0 and ( ϵ t )
Tt =1 are i.i.d standard Gaussian random variables.
We use 800 training, 80 validation and 80 testing sequences of length 24.
For experiments with Toy Data, we use the transformer from (Vaswani et al., 2017) with GeLU activation (Hendrycks and Gimpel, 2020), a unique attention block, one attention head and a hidden size of 64.
We train until convergence for 100 epochs.
We draw 30 posterior samples to compute the predictive distribution and KL divergence.
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 .
Part of Speech Tagging
We use the English split of the Part of Speech tagging dataset from the Universal Dependencies v1.2 corpus (Nivre et al., 2015).
This dataset contains 204’586 train, 25’148 validation and 25’096 test tokens.
We use a maximum length oftokens, pad the shorter sequences and split the sentences which exceed the text length.
For experiments with this dataset, we use a transformer from (Vaswani et al., 2017) with GeLU activation, a unique attention block, one attention head and a hidden size of 32.
We train until convergence for 100 epochs.
We draw 10 posterior samples to compute the predictive distribution and KL divergence.
We evaluate our model with the token level test data log-likelihood, token level accuracy, token level F1-score and expected calibration error (ECE) (Guo et al., 2017).
MNIST
We experiment with the MNIST image classification dataset (LeCun et al., 2010).
We split the original dataset into 48’000 training, 12’000 validation and 9’984 testing samples.
For experiments with this dataset, we use a Vision transformer from (Dosovitskiy et al., 2021) with GeLU activation, two attention block, one attention head, a hidden size of 32 and a patch size of 4.
We train until convergence for 150 epochs.
We draw 10 posterior samples to compute the predictive distribution and KL divergence.
We evaluate our model with the test data log-likelihood, accuracy, F1-score and ECE.
A.2.2. Training setup
Our models are trained using the Adam optimizer (Kingma and Ba, 2017) with the triangular learning rate schedule from (Vaswani et al., 2017).
Our weight distribution experiments are conducted using 900 transformers trained by likelihood maximization with SGD as done in (Fortuin et al., 2021b).
A.2.3. Software packages
We implement and train our models using the JAX (Bradbury et al., 2018) and Haiku (Hennigan et al., 2020)
Python libraries.