# Transformer Uncertainty Estimation with Hierarchical Stochastic Attention

arXiv: https://arxiv.org/abs/2112.13776 · HTML: https://mlsanity.com/paper/2112.13776/

## Key facts

- Cited by 21 papers in the map (data through 14 Oct 2026).
- Builds on 33 of its 56 references that are in the map.
- First posted on arXiv 27 Dec 2021.

## Builds on (most cited first, 30 of 33)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [A Baseline for Detecting Misclassified and Out-of-Distribution Examples in Neural Networks](https://mlsanity.com/paper/1610.02136/) (2016)
- [What Uncertainties Do We Need in Bayesian Deep Learning for Computer Vision?](https://mlsanity.com/paper/1703.04977/) (2017)
- [Informer: Beyond Efficient Transformer for Long Sequence Time-Series Forecasting](https://mlsanity.com/paper/2012.07436/) (2020)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [Neural Network Acceptability Judgments](https://mlsanity.com/paper/1805.12471/) (2018)
- [Efficient Transformers: A Survey](https://mlsanity.com/paper/2009.06732/) (2020)
- [A Survey of Uncertainty in Deep Neural Networks](https://mlsanity.com/paper/2107.03342/) (2021)
- [Neural Turing Machines](https://mlsanity.com/paper/1410.5401/) (2014)
- [Transformer in Transformer](https://mlsanity.com/paper/2103.00112/) (2021)
- [Variational Dropout and the Local Reparameterization Trick](https://mlsanity.com/paper/1506.02557/) (2015)
- [Pitfalls of In-Domain Uncertainty Estimation and Ensembling in Deep Learning](https://mlsanity.com/paper/2002.06470/) (2020)
- [On the Properties of the Softmax Function with Application in Game Theory and Reinforcement Learning](https://mlsanity.com/paper/1704.00805/) (2017)
- [Fast Transformers with Clustered Attention](https://mlsanity.com/paper/2007.04825/) (2020)
- [Pre-trained Language Models in Biomedical Domain: A Systematic Survey](https://mlsanity.com/paper/2110.05006/) (2021)
- ['In-Between' Uncertainty in Bayesian Neural Networks](https://mlsanity.com/paper/1906.11537/) (2019)
- [H-Transformer-1D: Fast One-Dimensional Hierarchical Attention for Sequences](https://mlsanity.com/paper/2107.11906/) (2021)
- [Estimating Uncertainty and Interpretability in Deep Learning for Coronavirus (COVID-19) Detection](https://mlsanity.com/paper/2003.10769/) (2020)
- [Doubly Stochastic Variational Inference for Neural Processes with Hierarchical Latent Variables](https://mlsanity.com/paper/2008.09469/) (2020)
- [Probabilistic Gradient Boosting Machines for Large-Scale Probabilistic Regression](https://mlsanity.com/paper/2106.01682/) (2021)
- [Tactical Decision-Making in Autonomous Driving by Reinforcement Learning with Uncertainty Estimation](https://mlsanity.com/paper/2004.10439/) (2020)
- [State-Regularized Recurrent Neural Networks](https://mlsanity.com/paper/1901.08817/) (2019)
- [Conversations with Search Engines: SERP-based Conversational Response Generation](https://mlsanity.com/paper/2004.14162/) (2020)
- [Wizard of Search Engine: Access to Information Through Conversations with Search Engines](https://mlsanity.com/paper/2105.08301/) (2021)

## Built on by (most cited first, 21 of 21)

- [Calibration in Deep Learning: A Survey of the State-of-the-Art](https://mlsanity.com/paper/2308.01222/) (2023)
- [Uncertainty in Natural Language Processing: Sources, Quantification, and Applications](https://mlsanity.com/paper/2306.04459/) (2023)
- [Energy-Based Transformers are Scalable Learners and Thinkers](https://mlsanity.com/paper/2507.02092/) (2025)
- [Interpreting Predictive Probabilities: Model Confidence or Human Label Variation?](https://mlsanity.com/paper/2402.16102/) (2024)
- [Uncertainty Quantification for Machine Learning in Healthcare: A Survey](https://mlsanity.com/paper/2505.02874/) (2025)
- [Calibrating Verbalized Probabilities for Large Language Models](https://mlsanity.com/paper/2410.06707/) (2024)
- [Probabilistic Transformer: Modelling Ambiguities and Distributions for RNA Folding and Molecule Design](https://mlsanity.com/paper/2205.13927/) (2022)
- [Benchmarking Uncertainty and its Disentanglement in multi-label Chest X-Ray Classification](https://mlsanity.com/paper/2508.04457/) (2025)
- [HyTRec: A Hybrid Temporal-Aware Attention Architecture for Long Behavior Sequential Recommendation](https://mlsanity.com/paper/2602.18283/) (2026)
- [Uncertainty-Instructed Structure Injection for Generalizable HD Map Construction](https://mlsanity.com/paper/2503.23109/) (2025)
- [Multimodal Graph Representation Learning for Robust Surgical Workflow Recognition with Adversarial Feature Disentanglement](https://mlsanity.com/paper/2505.01766/) (2025)
- [State-Regularized Recurrent Neural Networks to Extract Automata and Explain Predictions](https://mlsanity.com/paper/2212.05178/) (2022)
- [S2-Track: A Simple yet Strong Approach for End-to-End 3D Multi-Object Tracking](https://mlsanity.com/paper/2406.02147/) (2024)
- [Dropouts in Confidence: Moral Uncertainty in Human-LLM Alignment](https://mlsanity.com/paper/2511.13290/) (2025)
- [Calibrating Scientific Foundation Models with Inference-Time Stochastic Attention](https://mlsanity.com/paper/2604.19530/) (2026)
- [On the Calibration and Uncertainty with P'olya-Gamma Augmentation for Dialog Retrieval Models](https://mlsanity.com/paper/2303.08606/) (2023)
- [A Saliency-based Clustering Framework for Identifying Aberrant Predictions](https://mlsanity.com/paper/2311.06454/) (2023)
- [Stochastic Vision Transformers with Wasserstein Distance-Aware Attention](https://mlsanity.com/paper/2311.18645/) (2023)
- [BayesJudge: Bayesian Kernel Language Modelling with Confidence Uncertainty in Legal Judgment Prediction](https://mlsanity.com/paper/2404.10481/) (2024)
- [Can Uncertainty Quantification Improve Learned Index Benefit Estimation?](https://mlsanity.com/paper/2410.17748/) (2024)
- [TCUQ: Single-Pass Uncertainty Quantification from Temporal Consistency with Streaming Conformal Calibration for TinyML](https://mlsanity.com/paper/2508.12905/) (2025)

## Abstract (arXiv)

Transformers are state-of-the-art in a wide range of NLP tasks and have also been applied to many real-world products. Understanding the reliability and certainty of transformer model predictions is crucial for building trustable machine learning applications, e.g., medical diagnosis. Although many recent transformer extensions have been proposed, the study of the uncertainty estimation of transformer models is under-explored. In this work, we propose a novel way to enable transformers to have the capability of uncertainty estimation and, meanwhile, retain the original predictive performance. This is achieved by learning a hierarchical stochastic self-attention that attends to values and a set of learnable centroids, respectively. Then new attention heads are formed with a mixture of sampled centroids using the Gumbel-Softmax trick. We theoretically show that the self-attention approximation by sampling from a Gumbel distribution is upper bounded. We empirically evaluate our model on two text classification tasks with both in-domain (ID) and out-of-domain (OOD) datasets. The experimental results demonstrate that our approach: (1) achieves the best predictive performance and uncertainty trade-off among compared methods; (2) exhibits very competitive (in most cases, improved) predictive performance on ID datasets; (3) is on par with Monte Carlo dropout and ensemble methods in uncertainty estimation on OOD datasets.