Source-linked AI summary
ProbPlug: A Plugin Uncertainty Network for Reliable Confidence in LLM Binary Classification
Jianzong Wang, Chuhang Liu, Botao Zhao, Zuheng Kang, Xulong Zhang, Xiaoyang Qu, Junqing Peng, Zhiewei Ye, Yayun He
TL;DR
LLM classifiers often lack reliable confidence estimates despite strong classification performance, limiting their use in reliability-sensitive settings. ProbPlug estimates whether binary outputs are correct from hidden token representations of a frozen LLM, and experiments report better calibration, classification performance, and cross-task generalization with negligible overhead.
Problem
LLM-based classification lacks reliable confidence estimates, although confidence calibration is important for reliability-sensitive deployment.
Method
ProbPlug is a lightweight plug-in framework that estimates binary prediction correctness from hidden representations of a frozen LLM without modifying the base model.
Results
Experiments across multiple benchmarks report better-calibrated confidence estimates, improved downstream classification performance, negligible inference overhead, and promising cross-task generalization.
Takeaways & Limitations
ProbPlug provides a practical confidence-estimation approach for LLM-based binary classification while preserving the original model.
Abstract
from arXiv · showhide
Large language models (LLMs) have achieved strong performance across a broad range of classification settings, yet the reliability of their predictions remains a major obstacle to deployment in high-stakes scenarios. Although confidence estimation for LLMs has been widely studied, confidence calibration for LLM-based classification remains underexplored. We introduce ProbPlug, a lightweight confidence estimation framework for LLM-based binary classification, which predicts whether an output is correct using internal token features extracted from a frozen LLM. ProbPlug employs a self-attention module to aggregate hidden representations and can be integrated into the original inference pipeline without modifying the base model. Experiments across multiple tasks involving both text-based and multimodal large models show that ProbPlug provides more reliable confidence estimates, improves classification performance with negligible additional overhead, and exhibits strong generalization across tasks. These results indicate that ProbPlug serves as a practical solution for confidence estimation in LLM-based classification. Our code is publicly available at Github.
1 Introduction
ProbPlug addresses unreliable confidence in LLM classification by estimating prediction correctness from internal representations of a frozen model. Experiments report improved calibration, classification performance, and cross-task generalization with negligible overhead.
- Unreliable confidence estimates hinder deploying LLM classifiers in reliability-sensitive applications.
- The framework targets binary classification with newly generated Yes/No tokens to improve generalization.
- ProbPlug predicts confidence from layer-wise hidden representations while leaving the base LLM frozen and unmodified.
- Experiments across multiple benchmarks show better-calibrated confidence estimates, improved downstream classification performance, negligible inference overhead, and promising cross-task generalization.
2 Problem Statement
LLM-generated binary decisions indicate Yes or No but do not naturally provide calibrated confidence. ProbPlug therefore estimates answer correctness from complementary hidden representations across Transformer layers.
- Prompted LLMs generate discrete Yes/No responses, but these outputs do not naturally provide calibrated confidence scores.
- ProbPlug uses hidden representations from different Transformer layers to estimate whether the generated binary answer is correct.
- The framework learns a lightweight confidence estimator over layer-wise representations of a frozen LLM.
3 Method
ProbPlug aggregates hidden signals from a frozen LLM’s Transformer layers to estimate confidence for its generated binary output. The method uses the final generated token as the principal representation.
- ProbPlug combines Token Compression, Layer Information Aggregation, and a classification head while keeping the backbone LLM frozen during training.
- The frozen LLM produces a binary output token, whose hidden representations are extracted from all Transformer layers.
- Traditional LLM inference generates a textual response through next-token prediction from a prompted input.
Y E S
ProbPlug estimates confidence from final-token hidden representations across Transformer layers while keeping the backbone LLM frozen. It compresses token features, adaptively aggregates layer information, and supports binary, multi-class, and cross-task classification.
- Y E S: ProbPlug uses the final generated token’s hidden states across layers, represented as (N, S, L, D), with S = 1 by default.The final token directly reflects the model’s decision and empirically provides stronger cross-task generalization.
- Y E S: Token Compression transforms and aggregates token-level features into a single representation for each Transformer layer.The module uses nonlinear mapping followed by fully connected layers that expand S to 2S and compress it to 1.
- Y E S: Multi-head attention adaptively fuses layer-wise representations because Transformer layers contribute unequally to decision confidence.A learnable query vector produces attention weights for combining information from different layers.
- Y E S: The aggregated representation is passed to a lightweight classifier that outputs the estimated probability that the generated answer is correct.The backbone LLM remains frozen during training, while the external confidence network is trained separately.
- Y E S: Because many classification tasks can be reformulated as Yes/No decisions, the trained confidence estimator can transfer to unseen tasks by modifying the prompt template.This supports cross-task generalization without changing the frozen base model.
- Y E S: ProbPlug supports multi-class classification by decomposing it into one-vs-rest binary decisions and applying softmax normalization to class confidence scores.The formulation also enables extension to multi-class and multimodal scenarios.
4 Experimental Results and Analysis
Experiments evaluate ProbPlug across five public text datasets and the multimodal IEMOCAP dataset, using multiple confidence-estimation and inference baselines. The reported comparisons cover in-task and cross-task settings, with F1-score and related performance measures.
- 4 Experimental Results and Analysis: Experiments cover spam, sentiment, toxicity, and product-review classification across five public text datasets.The multimodal evaluation uses IEMOCAP, containing 5,531 utterances from four emotion categories.
- 4 Experimental Results and Analysis: Table 2 reports ProbPlug comparisons across Civil Comments and Amazon-polarity in both in-task and cross-task scenarios.The table evaluates the proposed method on the two datasets described in its caption.
- 4 Experimental Results and Analysis: The evaluation compares Qwen3-8B with verbalization, logit, self-consistency, CISC, and SAPLMA-based confidence-estimation approaches.Model performance is quantified using F1-score and additional reported evaluation measures.
Baselines
ProbPlug consistently outperforms competing confidence-estimation methods across text and multimodal classification settings, including multi-class speech emotion recognition.
- ProbPlug outperforms competitive confidence-estimation methods on in-task datasets in both F1-score and AUPRC.
- Across cross-task evaluations, ProbPlug generally achieves better performance than current state-of-the-art approaches without additional training.Cross-task transfer requires only modifying prompt templates.
- ProbPlug requires one inference pass, unlike CISC’s repeated inferences, while achieving better AUPRC on SST2.
- In multimodal speech emotion recognition, ProbPlug shows significant advantages across all reported metrics over non-LLM and competing confidence-estimation methods.The evaluation uses Qwen2-audio on the Iemocap task; UA and WA are reported metrics.
4.4 Calibration Experiments
ProbPlug provides better confidence calibration than SAPLMA, with lower average ECE and Brier Score and notable gains on SMS-SPAM and Amazon-polarity.
- ProbPlug achieves lower average ECE (0.0805) and Brier Score (0.0824) than SAPLMA.ECE measures agreement between confidence and accuracy, while Brier Score measures probability-label discrepancy.
- ProbPlug shows notable calibration gains on SMS-SPAM and Amazon-polarity.
4.5 Ablation Studies
Ablations show that token compression and layer-wise aggregation improve performance, while using only the final token improves cross-task robustness by reducing task-specific noise.
- Both the token compression module and layer-wise information aggregation module improve performance.
- Using 10 tokens slightly improves in-task performance over S=1 but noticeably reduces cross-task generalization.
- Using only the final token, S=1, provides better cross-task robustness than including preceding tokens.The paper attributes this to preceding tokens carrying task-specific semantic information that can cause source-domain overfitting.
4.6 Layer-wise Attention Visualization and Interpretability
Layer-wise attention reveals that confident predictions rely primarily on early lexical and final reasoning layers, while intermediate layers receive little weight.
- In high-confidence cases, attention concentrates on early lexical layers L0–L5 and final reasoning layers L30–L35.
- Intermediate layers receive little attention, suggesting they contribute limited information and may introduce task-specific noise.
- Suppressing less informative layers helps ProbPlug maintain focused attention and improve confidence robustness.
5 Conclusion
ProbPlug estimates confidence for LLM-based binary classification using hidden representations from frozen base models. Across text and multimodal tasks, it improves confidence reliability, classification performance, and cross-task generalization with limited overhead.
- ProbPlug estimates prediction confidence from hidden token representations without modifying the frozen base model’s original architecture.It is designed as a plug-in framework for LLM-based binary classification.
- Across multiple classification tasks, ProbPlug provides more reliable confidence estimates for standard and multimodal large models.The experiments include Qwen3 and Qwen2-Audio.
- ProbPlug delivers consistent downstream classification gains with only limited additional inference overhead.
- ProbPlug demonstrates promising cross-task generalization, reducing the need for extensive retraining when adapting to new tasks.