Confidence-Based Response Abstinence: Improving LLM Trustworthiness via Activation-Based Uncertainty Estimation

TL;DR

The paper tackles abstention in retrieval-augmented generation within high-stakes domains by introducing an activation-based confidence model that treats confidence as a sequence-classification problem over FFN activations. An LSTM-based probe processes activations from selected layers (e.g., layer 16 or 32) to produce a confidence score, with the objective optimized via $L_{Total} = L_{CE} + \lambda L_{Huber}$ to improve robustness against noisy supervision. Empirical results on a proprietary financial knowledge base show that activation-based confidence can achieve precision > 0.95 with masking around 30% of outputs, while markedly reducing latency when using shallower activations. The approach offers a scalable, architecture-aware pathway to trustworthy RAG deployment, enabling abstention decisions that balance utility and safety in regulated settings.

Abstract

We propose a method for confidence estimation in retrieval-augmented generation (RAG) systems that aligns closely with the correctness of large language model (LLM) outputs. Confidence estimation is especially critical in high-stakes domains such as finance and healthcare, where the cost of an incorrect answer outweighs that of not answering the question. Our approach extends prior uncertainty quantification methods by leveraging raw feed-forward network (FFN) activations as auto-regressive signals, avoiding the information loss inherent in token logits and probabilities after projection and softmax normalization. We model confidence prediction as a sequence classification task, and regularize training with a Huber loss term to improve robustness against noisy supervision. Applied in a real-world financial industry customer-support setting with complex knowledge bases, our method outperforms strong baselines and maintains high accuracy under strict latency constraints. Experiments on Llama 3.1 8B model show that using activations from only the 16th layer preserves accuracy while reducing response latency. Our results demonstrate that activation-based confidence modeling offers a scalable, architecture-aware path toward trustworthy RAG deployment.

Figures (5)

• Figure 1: Diagram of the proposed Retrieval Augmented Generation (RAG) with the confidence model. When a user asked a question, the IR component retrieves related context from a database. The prompt is then constructed and sent into a question and answering LLM. A confidence score would be generated by the confidence model and being used to control whether or not to show the result to the user.
• Figure 2: Landscape of Uncertainty Quantification Methods
• Figure 3: Motivation to use activations instead of token probabilities.
• Figure 4: Confidence model based on the activations of large language models. Our method first feed the <Question, Context, Response> pair in an LLM. We then extract the activations from the 32-th or 16-th layer, and feed the activations into an LSTM and a classification head. The classification logit serves as the confidence score.
• Figure 5: Features of our Knowledge base. (a) Complex structure of our knowledge articles; (b) Process of SME validated training data creation; (c) Example SME validated LLM Response