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.

References in corpus (23)

• An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale
• Language Models are Few-Shot Learners
• On Calibration of Modern Neural Networks
• Visual Transformers: Token-based Image Representation and Processing for Computer Vision
• TransGAN: Two Pure Transformers Can Make One Strong GAN, and That Can Scale Up
• Multiplicative Normalizing Flows for Variational Bayesian Neural Networks
• Bayesian Deep Learning and a Probabilistic Perspective of Generalization
• Flipout: Efficient Pseudo-Independent Weight Perturbations on Mini-Batches
• VideoGPT: Video Generation using VQ-VAE and Transformers
• Implicit Reparameterization Gradients
• Cyclical Stochastic Gradient MCMC for Bayesian Deep Learning
• Structured and Efficient Variational Deep Learning with Matrix Gaussian Posteriors
• Latent Alignment and Variational Attention
• Uncertainty in Neural Networks: Approximately Bayesian Ensembling
• AST: Audio Spectrogram Transformer
• SLANG: Fast Structured Covariance Approximations for Bayesian Deep Learning with Natural Gradient
• Bayesian Neural Network Priors Revisited
• The k-tied Normal Distribution: A Compact Parameterization of Gaussian Mean Field Posteriors in Bayesian Neural Networks
• BNNpriors: A library for Bayesian neural network inference with different prior distributions
• Bayesian Inference with Anchored Ensembles of Neural Networks, and Application to Exploration in Reinforcement Learning
• Repulsive Deep Ensembles are Bayesian
• Repulsive Attention: Rethinking Multi-head Attention as Bayesian Inference
• On Stein Variational Neural Network Ensembles