# Uncertainty-Aware Transformers: Conformal Prediction for Language Models

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 24 of its 38 references that are in the map.
- First posted on arXiv 10 Apr 2026.

## Builds on (most cited first, 24 of 24)

- [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)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding](https://mlsanity.com/paper/1804.07461/) (2018)
- [Axiomatic Attribution for Deep Networks](https://mlsanity.com/paper/1703.01365/) (2017)
- [SuperGLUE: A Stickier Benchmark for General-Purpose Language Understanding Systems](https://mlsanity.com/paper/1905.00537/) (2019)
- [Attention is not Explanation](https://mlsanity.com/paper/1902.10186/) (2019)
- [A Gentle Introduction to Conformal Prediction and Distribution-Free Uncertainty Quantification](https://mlsanity.com/paper/2107.07511/) (2021)
- [Attention is not not Explanation](https://mlsanity.com/paper/1908.04626/) (2019)
- [Classification with Valid and Adaptive Coverage](https://mlsanity.com/paper/2006.02544/) (2020)
- [A Multiscale Visualization of Attention in the Transformer Model](https://mlsanity.com/paper/1906.05714/) (2019)
- [Least Ambiguous Set-Valued Classifiers with Bounded Error Levels](https://mlsanity.com/paper/1609.00451/) (2016)
- [Deep k-Nearest Neighbors: Towards Confident, Interpretable and Robust Deep Learning](https://mlsanity.com/paper/1803.04765/) (2018)
- [The BigScience ROOTS Corpus: A 1.6TB Composite Multilingual Dataset](https://mlsanity.com/paper/2303.03915/) (2023)
- [On the Societal Impact of Open Foundation Models](https://mlsanity.com/paper/2403.07918/) (2024)
- [Understanding Softmax Confidence and Uncertainty](https://mlsanity.com/paper/2106.04972/) (2021)
- [Criteria of efficiency for conformal prediction](https://mlsanity.com/paper/1603.04416/) (2016)
- [Conformal prediction for text infilling and part-of-speech prediction](https://mlsanity.com/paper/2111.02592/) (2021)
- [On the Expected Size of Conformal Prediction Sets](https://mlsanity.com/paper/2306.07254/) (2023)
- [Transformer Conformal Prediction for Time Series](https://mlsanity.com/paper/2406.05332/) (2024)
- [CONFINE: Conformal Prediction for Interpretable Neural Networks](https://mlsanity.com/paper/2406.00539/) (2024)
- [On the Relationship Between Interpretability and Explainability in Machine Learning](https://mlsanity.com/paper/2311.11491/) (2023)
- [Explorations of the Softmax Space: Knowing When the Neural Network Doesn't Know](https://mlsanity.com/paper/2502.00456/) (2025)

## Abstract (arXiv)

Transformers have had a profound impact on the field of artificial intelligence, especially on large language models and their variants. However, as was the case with neural networks, their black-box nature limits trust and deployment in high-stakes settings. For models to be genuinely useful and trustworthy in critical applications, they must provide more than just predictions: they must supply users with a clear understanding of the reasoning that underpins their decisions. This article presents an uncertainty quantification framework for transformer-based language models. This framework, called CONFIDE (CONformal prediction for FIne-tuned DEep language models), applies conformal prediction to the internal embeddings of encoder-only architectures, like BERT and RoBERTa, while enabling hyperparameter tuning. CONFIDE uses either [CLS] token embeddings or flattened hidden states to construct class-conditional nonconformity scores, enabling statistically valid prediction sets with instance-level explanations. Empirically, CONFIDE improves test accuracy by up to 4.09% on BERT-tiny and achieves greater correct efficiency (i.e., the expected size of the prediction set conditioned on it containing the true label) compared to prior methods, including NM2 and VanillaNN. We show that early and intermediate transformer layers often yield better-calibrated and more semantically meaningful representations for conformal prediction. In resource-constrained models and high-stakes tasks with ambiguous labels, CONFIDE offers robustness and interpretability where softmax-based uncertainty fails. We position CONFIDE as a framework for practical diagnostic and efficiency/robustness improvement over prior conformal baselines.