Uncertainty-Aware Transformers: Conformal Prediction for Language Models

Abhiram Vellore; Niraj K. Jha

cs.LGApr 10, 2026

Uncertainty-Aware Transformers: Conformal Prediction for Language Models

Authors: Abhiram Vellore, Niraj K. Jha

Organizations: Department of Electrical and Computer Engineering, Princeton University

Abstract

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.

Figures & tables

Figure 1: Failure case on BoolQ (BERT-tiny). CONFIDE fails to distinguish between same-class and different-class neighbors, resulting in a high-uncertainty prediction.

Method	CoLA		MNLI		MRPC		QNLI
	Test Acc	Top Corr Eff	Test Acc	Top Corr Eff	Test Acc	Top Corr Eff	Test Acc	Top Corr Eff
Original NN	0.5858	–	0.6949	–	0.7010	–	0.7895	–
1-Nearest Neighbor	0.5791	0.5762	0.3403	0.3254	0.6103	0.6103	0.5067	0.5050
NM (2)	0.5858	0.5724	0.6949	0.6936	0.7010	0.7010	0.7895	0.7878
CONFIDE-A (ours)	0.6903	0.6903	0.6982	0.6907	0.7304	0.7181	0.8157	0.8104
CONFIDE-C (ours)	0.6903	0.6903	0.6975	0.6932	0.7279	0.7279	0.8131	0.8129

Table 1: Performance on GLUE benchmarks using BERT-tiny: CoLA, MNLI, MRPC, and QNLI. Hyperparameters used can be found in Appendix B.1

Method	QQP		RTE		SST-2		WNLI
	Test Acc	Top Corr Eff	Test Acc	Top Corr Eff	Test Acc	Top Corr Eff	Test Acc	Top Corr Eff
Original NN	0.8116	–	0.5668	–	0.8303	–	0.5634	–
1-Nearest Neighbor	0.6423	0.6391	0.4693	0.4549	0.5011	0.4794	0.2394	0.2394
NM (2)	0.8214	0.8073	0.5668	0.5668	0.8303	0.8280	0.5634	0.5634
CONFIDE-A (ours)	0.8566	0.8530	0.5957	0.5848	0.8372	0.8326	0.8360	0.8245
CONFIDE-C (ours)	0.8566	0.8530	0.5884	0.5884	0.8349	0.8349	0.8326	0.8280

Table 2: Performance on GLUE benchmarks using BERT-tiny: QQP, RTE, SST-2, and WNLI.

Method	BoolQ		CB		MultiRC
	Test Acc	Top Corr Eff	Test Acc	Top Corr Eff	Test Acc	Top Corr Eff
Original NN	0.7275	–	0.6250	–	0.6640	–
1-Nearest Neighbor	0.5602	0.5590	0.4286	0.4286	0.5192	0.4792
NM1/NM2	0.7275	0.6269	0.6250	0.6250	0.6640	0.6621
CONFIDE-A (ours)	0.7297	0.6104	0.7143	0.6786	0.6625	0.6557
CONFIDE-C (ours)	0.7272	0.7245	0.7143	0.7143	0.6599	0.6582

Table 3: Performance on SuperGLUE benchmarks using BERT-small: BoolQ, CB, and MultiRC.

Method	CoLA		MNLI		MRPC		QNLI
	Test Acc	Top Corr Eff	Test Acc	Top Corr Eff	Test Acc	Top Corr Eff	Test Acc	Top Corr Eff
Original NN	0.8408	–	–	–	0.8824	–	0.9231	–
1-Nearest Neighbor	0.6069	0.6021	–	–	0.5564	0.5490	–	–
NM1/NM2	0.8408	0.8370	0.8608	0.8497	0.8824	0.8627	0.9231	0.9098
CONFIDE-A (ours)	0.8495	0.8092	0.8608	0.8497	0.8848	0.8725	0.9235	0.9087
CONFIDE-C (ours)	0.8495	0.8495	0.8588	0.8512	0.8848	0.8848	0.9235	0.9217

Table 4: Performance on GLUE benchmarks using RoBERTa-base: CoLA, MNLI, MRPC, and QNLI.

Method	QQP		RTE		SST-2		WNLI
	Test Acc	Top Corr Eff	Test Acc	Top Corr Eff	Test Acc	Top Corr Eff	Test Acc	Top Corr Eff
Original NN	0.9069	–	0.7220	–	0.9323	–	0.4366	–
1-Nearest Neighbor	–	–	0.4693	0.4657	0.5092	0.5034	0.1972	0.1831
NM1/NM2	0.9069	0.8961	0.7220	0.7112	0.9323	0.9255	0.4366	0.4366
CONFIDE-A (ours)	0.9075	0.8955	0.7365	0.7256	0.9346	0.9289	–	–
CONFIDE-C (ours)	0.9070	0.8961	0.7329	0.7292	0.9323	0.9323	–	–

Table 5: Performance on GLUE benchmarks using RoBERTa-base: QQP, RTE, SST-2, and WNLI.

Method	BoolQ		CB		MultiRC
	Test Acc	Top Corr Eff	Test Acc	Top Corr Eff	Test Acc	Top Corr Eff
Original NN	0.8092	–	0.6786	–	0.7473	–
1-Nearest Neighbor	–	–	0.4464	0.4286	–	–
NM1/NM2	0.8092	0.7654	0.6786	0.6607	0.7473	0.7471
CONFIDE-A (ours)	0.8095	0.5483	0.7500	0.7500	0.4299	0.4292
CONFIDE-C (ours)	0.8089	0.8089	0.7500	0.7500	0.4299	0.4292

Table 6: Performance on SuperGLUE benchmarks using RoBERTa-base: BoolQ, CB, and MultiRC.

Method	CoLA		MRPC		RTE		WNLI
	Test Acc	Top Corr Eff	Test Acc	Top Corr Eff	Test Acc	Top Corr Eff	Test Acc	Top Corr Eff
Original NN	0.8418	–	0.8578	–	0.8087	–	0.5634	–
1-Nearest Neighbor	0.6069	0.6021	0.5564	0.5490	0.4693	0.4657	0.1972	0.1831
NM1/NM2	0.8418	0.8293	0.8578	0.8578	0.8087	0.8087	0.5634	0.5493
CONFIDE-A (ours)	0.8466	0.8437	0.8710	0.8578	0.8123	0.7978	0.5634	0.5493
CONFIDE-C (ours)	0.8466	0.8437	0.8652	0.8652	0.8051	0.8051	0.5493	0.5493

Table 7: Performance on GLUE and SuperGLUE benchmarks using RoBERTa-large: CoLA, MRPC, RTE, and WNLI.

Panel of Figure 2 — Figure 2: Classwise and aggregate coverage curves for BERT-tiny on CoLA. The false class severely undercovers, violating conformal validity.

Panel of Figure 3 — Figure 3: Classwise and aggregate coverage curves for BERT-tiny on BoolQ (non-classwise). The false class shows massive undercoverage due to biased calibration, despite almost diagonal aggregated statistics.

Panel of Figure 4 — Figure 4: Coverage curves on BoolQ using classwise calibration. The classwise method reduces undercoverage but does not fully resolve it.

Panel of Figure 5 — Figure 5: Coverage curves on CoLA using classwise calibration.

Panel of Figure 6 — Figure 6: Coverage curves on CoLA using RoBERTa models. Undercoverage and classwise calibration perform marginally better than smaller models but still poorly.

Panel of Figure 7 — Figure 7: RoBERTa-large on CB (classwise): neutral class coverage fails entirely for neutral labels despite overall decent accuracy.

Panel of Figure 8 — Figure 8: Classwise coverage curves for CoLA using the NM2 prior method. Both BERT-tiny and RoBERTa-base show persistent undercoverage, particularly for the “unacceptable” class.

Panel of Figure 9 — Figure 9: Classwise coverage curves for BoolQ using BERT-tiny with two prior methods. Both NM2 and Vanilla NN fail to maintain valid coverage, especially for the “false” class.

Task / Variant	Attention		Flattened
	Acc.	Corr. Eff.	Acc.	Corr. Eff.
MNLI – BERT-tiny
CONFIDE-A	0.6982	0.6907	0.6993	0.6497
CONFIDE-C	0.6975	0.6932	0.6980	0.6971
SST-2 – BERT-small
CONFIDE-A	0.8933	0.8647	0.8968	0.8853

Table 8: Comparison of CONFIDE-A and CONFIDE-C using Attention and Flattened representations for MNLI (BERT-tiny) and SST-2 (BERT-small). Bolded values highlight the best result for each task and metric.

Appendix figures & tables30 assets

Supplementary material from the paper’s appendix.

Appendix

Dataset	Model	Method	Hyperparameters
glue_cola	BERT-small	CONFIDE-A	$l=81,\ k=10,\ \text{flattened},\ \text{cosine},\ \text{PCA}:No$
		CONFIDE-C	$l=81,\ k=10,\ \text{flattened},\ \text{cosine},\ \text{PCA}:No$
	BERT-tiny	CONFIDE-A	$l=16,\ k=50,\ \text{flattened},\ \text{Mahalanobis},\ \text{PCA}:Yes$
		CONFIDE-C	$l=16,\ k=50,\ \text{flattened},\ \text{Mahalanobis},\ \text{PCA}:Yes$
	RoBERTa-base	CONFIDE-A	$l=190,\ k=20,\ \text{flattened},\ \text{cosine},\ \text{PCA}:No$
		CONFIDE-C	$l=\text{softmax},\ T=1.0,\ k=1,\ \text{flattened},\ \text{cosine},\ \text{PCA}:Yes$

Appendix

Table 9: Best CONFIDE hyperparameters across datasets and models.

BERT-tiny (BoolQ)
$l$	Semantic alias
16	block.0.attn_out
27	block.1
34	block.1.attn_out
38	block.1.mlp_in
41	block.1.output

Appendix

Table 19

RoBERTa-base (BoolQ)
$l$	Semantic alias
45	block.2.attn_out
63	block.3.attn_out
65	block.3.mlp_in
99	block.5.attn_out
191	block.10.attn_in

Appendix

Table 20

Figure 10: BERT-small hyperparameter testing on the CoLA dataset - accuracy

Figure 11: BERT-small hyperparameter testing on the MRPC dataset - correct efficiency

Figure 12: RoBERTa-large hyperparameter testing on the MRPC dataset - accuracy

Figure 13: BERT-tiny hyperparameter testing on the RTE dataset - correct efficiency

Figure 14: BERT-tiny hyperparameter testing on the SST2 dataset - accuracy

Figure 15: RoBERTa-base hyperparameter testing on the CB dataset - accuracy

Dataset	$\Delta$ calibration time / block (s)	$\Delta$ test time / block (s)
SST-2	1.58 (R 2 =0.983)	0.075 (R 2 =0.966)
QNLI	5.40 (R 2 =0.961)	0.980 (R 2 =0.994)
MRPC	0.048 (R 2 =0.735)	0.021 (R 2 =0.709)

Appendix

Table 10: RoBERTa-base timing vs depth. Slope from least-squares fit of time to encoder block index.

Model	Dataset	Test latency (ms / example)
BERT-tiny	MRPC	1.48
RoBERTa-base	MRPC	3.10

Appendix

Table 11: Per-example test time averaged across probed blocks; exact

k

NN; same hardware.

Figure 16: Success case on BoolQ (BERT-tiny). CONFIDE separates neighbors by class, yielding a confident and credible singleton prediction.

Method	BoolQ		CB		MultiRC
	Test Acc	Top Corr Eff	Test Acc	Top Corr Eff	Test Acc	Top Corr Eff
Original NN	0.6596	–	0.5357	–	0.6219	–
1-Nearest Neighbor	0.5602	0.5590	0.4286	0.4286	0.5192	0.4792
NM (2)	0.6596	0.6529	0.5357	0.5000	0.6219	0.6209
CONFIDE-A (ours)	0.6636	0.6630	0.7321	0.7143	0.6275	0.5918
CONFIDE-C (ours)	0.6636	0.6630	0.7321	0.7143	0.6244	0.6229

Appendix

Table 12: Performance on SuperGLUE benchmarks using BERT-tiny: BoolQ, CB, and MultiRC.

Method	CoLA		MNLI		MRPC		QNLI
	Test Acc	Top Corr Eff	Test Acc	Top Corr Eff	Test Acc	Top Corr Eff	Test Acc	Top Corr Eff
Original NN	0.7421	–	0.7914	0.7805	0.7230	–	0.8691	–
1-Nearest Neighbor	0.5791	0.5762	0.3403	0.3254	0.6103	0.6103	0.5067	0.5050
NM1/NM2	0.7421	0.7392	0.7914	0.7805	0.7230	0.7034	0.8691	0.8671
CONFIDE-A (ours)	0.7546	0.7546	0.7904	0.7865	0.7672	0.7525	0.8704	0.8669
CONFIDE-C (ours)	0.7546	0.7546	0.7904	0.7865	0.7598	0.7574	0.8699	0.8697

Appendix

Table 13: Performance on GLUE benchmarks using BERT-small: CoLA, MNLI, MRPC, and QNLI.

Method	QQP		RTE		SST-2		WNLI
	Test Acc	Top Corr Eff	Test Acc	Top Corr Eff	Test Acc	Top Corr Eff	Test Acc	Top Corr Eff
Original NN	0.8893	–	0.5884	–	0.8911	–	0.4648	–
1-Nearest Neighbor	0.6423	0.6391	0.4693	0.4549	0.5011	0.4794	0.2394	0.2394
NM1/NM2	0.8893	0.8739	0.5884	0.5884	0.8911	0.8830	0.4648	0.4648
CONFIDE-A (ours)	0.8933	0.8820	0.6245	0.6173	0.8933	0.8647	0.5915	0.5493
CONFIDE-C (ours)	0.8901	0.8846	0.6245	0.6173	0.8899	0.8911	0.5634	0.5634

Appendix

Table 14: Performance on GLUE benchmarks using BERT-small: QQP, RTE, SST-2, and WNLI.

Panel of Figure 17 — Figure 17: Classwise and aggregate coverage curves for BERT-tiny on MNLI. All classes are well calibrated and closely follow the diagonal, indicating robust conformal validity.

Panel of Figure 18 — Figure 18: Improvement performances for CoLA with classwise calibration.

Panel of Figure 19 — Figure 19: Improvement performances for BoolQ compared to smaller models

Panel of Figure 20 — Figure 20: Coverage curves on CoLA using BERT-small models. Undercoverage and classwise calibration perform as poorly as BERT-tiny models.

Explore similar work

Beyond Surface Statistics: Robust Conformal Prediction for LLMs via Internal Representations
Apr 17, 2026Yanli Wang, Peng Kuang, Xiaoyu Han +2Confidence Estimation in Language Models Conformal Prediction
Imperial College London · Zhejiang University · University of Illinois Urbana-Champaign +1
U-Space: Uncovering When and Why Uncertainty Arises in Language Models
Oct 6, 2026Tobias Braun, Nils Loose, Alexander Herzog +4LLM Uncertainty Estimation Confidence Estimation in Language Models
Technische Universität Darmstadt · University College London · Universität zu Lübeck +2
Improved Confidence Estimates for Black-Box Large Language Models
Aug 19, 2026Sokhna Diarra Mbacke, Mouloud Belbahri, Gabriel Loaiza-GanemConfidence Estimation in Language Models LLM Uncertainty Estimation
Layer 6 AI · TD Insurance

Apr 17, 2026J/K move · Enter open · S save

Beyond Surface Statistics: Robust Conformal Prediction for LLMs via Internal Representations

Yanli Wang, Peng Kuang, Xiaoyu Han, Kaidi Xu+1

Imperial College London · Zhejiang University · University of Illinois Urbana-Champaign+1

Confidence Estimation in Language Models Conformal Prediction Language Model Calibration LLM Reliability Conformal Risk Control LLM Uncertainty Estimation+1

Open paper AskPDF arXiv

Large language models are increasingly deployed in settings where reliability matters, yet output-level uncertainty signals such as token probabilities, entropy, and self-consistency can become brittle under calibration--deployment mismatch. Conformal prediction provides finite-sample validity under exchangeability, but its practical usefulness depends on the quality of the nonconformity score. We propose a conformal framework for LLM question answering that uses internal representations rather than output-facing statistics: specifically, we introduce Layer-Wise Information (LI) scores, which measure how conditioning on the input reshapes predictive entropy across model depth, and use them as nonconformity scores within a standard split conformal pipeline. Across closed-ended and open-domain QA benchmarks, with the clearest gains under cross-domain shift, our method achieves a better validity--efficiency trade-off than strong text-level baselines while maintaining competitive in-domain reliability at the same nominal risk level. These results suggest that internal representations can provide more informative conformal scores when surface-level uncertainty is unstable under distribution shift.

Abstract

Explore similar work

Beyond Surface Statistics: Robust Conformal Prediction for LLMs via Internal Representations

U-Space: Uncovering When and Why Uncertainty Arises in Language Models

Improved Confidence Estimates for Black-Box Large Language Models