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

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.

## introduction in high-stakes applications like financial customer support, it is often more desirable and trustworthy for a retrieval augmented generation (rag) system to abstain from answering than to risk providing an incorrect response. although not responding to a query reduces the system's immediate utility, it is a necessary trade-off to ensure accuracy and preserve user trust. the guiding principle is that the reputational and financial cost of providing a wrong answer is significantly higher than the cost of not providing one. this challenge requires a principle of abstention. one way to achieve the abstention is to have a confidence measure that correlates with correctness of the response, and mask the response when the confidence score is below a threshold. uncertainty of the model while generating the response is a viable source of signal for building a confidence measure. to develop a practical solution, it is crucial to identify the primary source of this uncertainty. in highly regulated fields, the error is rarely due to aleatoric uncertainty (randomness inherent in the data), as knowledge bases are typically vetted by legal and subject-matter experts. the more probable source is epistemic uncertainty (the model's own lack of knowledge), which arises when the model's parametric knowledge, acquired during pre-training or fine-tuning, conflicts with or misinterprets the provided context. while existing approaches @xcite @xcite to uncertainty estimation in retrievalaugmented generation (rag) have shown promise, they often fall short when the target response is long and narrative in nature. this challenge becomes especially pronounced in sensitive domains such as finance, where queries can be ambiguous or underspecified. for instance, a question like "what is the deadline to make a payment on card type a?" may retrieve multiple similar documents, each corresponding to different subcategories of the card type. in such cases, both the query and the retrieved context exhibit ambiguity, which can propagate through the rag pipeline. simply measuring uncertainty based on generated response is insufficient to ensure correctness. also, methods relying on sampling @xcite , are less practical at scale. these techniques rely on generating a response multiple times with slight variations to measure the model's consistency, a process that introduces prohibitive computational costs and latency in a production environment. for rag systems that must serve users in real-time, such multi-generational approaches are not a viable solution. uncertainty and correctness, while related, are fundamentally distinct concepts @xcite . a model's low uncertainty in its output does not necessarily imply correctness, just as a model may generate a correct response with a high uncertainty. this distinction becomes particularly salient in retrieval-augmented generation (rag) applications, where correctness often hinges on factual grounding rather than surface-level fluency. our goal is to utilize the model's internal uncertainty signals to generate a confidence score that correlates strongly with the correctness of the response generated by an llm. we build our confidence model using the raw activation signals inside the feedforward layers of llm which include the activations of knowledge neurons @xcite . thus, our model captures the relationship between the autoregressive properties of activations and inherent uncertainty of the model in generating a response. we propose a supervised framework to train a sequence classifier model and generate a confidence score that correlates with response correctness. figure 1 illustrates the practical utility of integrating a confidence model into our rag pipeline. the primary goal of the system is to provide users with accurate answers. however, in cases where there is insufficient epistemic or aleatoric knowledge to reliably answer a question, the system's next best action is to abstain from answering. this behavior is enabled by a controller that filters responses based on their confidence scores, allowing the system to avoid potentially incorrect or misleading outputs. this system is deployed in production for large-scale use that achieves high precision while maintaining an acceptable display rate (defined as the ratio of response pass the confidence filter to total responses generated by the system). experimental results show that our confidence model outperforms multiple baselines, reaching a precision of 0.95 with 70.1% display rate (masking 29.9% of the total responses). furthermore, when compared to ground truth, displayed responses exhibit a significantly higher rouge score than masked responses. azaria and mitchell (2023) show that the llms internal parameters show tell-tale signs when generating text with uncertainty. when a model's generation path falls into a speculative region, evidenced by competition between two or more plausible next tokens, its confidence is adversely affected. they introduce small input perturbations to induce trajectory shifts, and monitor corresponding changes in token-generation activations and outputs. they label the speculative generation as lying, and propose that activation patterns can shed light on this speculative generation. this method requires white-box access to the model to obtain token-level probability traces. tian et al. (2023) have empirically shown that llm's self generated confidence score while giving a response could be calibrated by sampling over perturbed questions. specifically, they show that prompting a model to produce several answer choices before giving its confidence scores helps in calibration of verbalized probabilities. in a related direction, @xcite generate multiple variants of a prompt using diverse prompting strategies such as chain-of-thought (cot), self-probing, and top-k sampling. they utilize a separate llm as a "judge model" to evaluate each variant and assign a confidence score. variations in these scores are then used to predict the confidence of the target model's original response. similarly, @xcite proposed a confidence measurement based on the perturbation of the question. the variation in model's answer generation probabilities for various perturbations of the question for the same context is used as a measure to generate a verbalized confidence score. several recent studies adopt a classificationbased approach to estimate response plausibility, offering a more computationally efficient alternative by avoiding multiple generations. for example, hhem @xcite uses an entailmentbased model to measure the semantic coherence between the input and the generated output. this approach operates under black-box constraints, requiring only the input-output pair from the target llm to assess the correctness of the response. other methods focus on linguistic cues as indicators of ambiguity in llm outputs. @xcite argue that model confidence does not always correlate with correctness and show that linguistic calibration of input prompts can significantly influence a model's confidence. they introduce a calibration score that helps generate more accurate responses by aligning linguistic features with expected confidence levels. their evaluations were performed on factoid qa datasets, where there is a zero-sum approach towards correctness. we argue that when the parametric knowledge of the llm is mainly contributing to the style of the response, and the key facts come from the input, confidence can serve as an effective signal for correctness. our method draws inspiration from prior work on activation-based knowledge tracing @xcite , generation trajectory modeling (azaria and mitchell, 2023), and importance-weighted token probabilities @xcite . @xcite highlight how feedforward network (ffn) activations encode key factual information, showing that the activation of certain neurons is positively correlated with knowledge expression. building on this insight, we treat ffn activations as autoregressive signals and train a recurrent neural network (rnn) to predict the probability that a model-generated response is correct. a score closer to 1 indicates greater model confidence in the response's correctness. ## method for a generated response sequence s of length l for the given input x to a model m with parameters θ, the probability of generating the sequence is given as follows: to compare sequence probability across different lengths of generated output, previous approaches have normalized the score based on the length of the response. the length-normalized score, used in prior uncertainty estimation (ue) methods (malinin and gales, 2020): here all the tokens contribute are given equal importance irrespective of the length of sequence. the risk with this approach is that a single lowprobability unusual word can disproportionately lower the overall sequence score, even if subsequent tokens have high probabilities. several of the methods that perform uncertainty estimation taking token-logits perform similar weighing and they have shown great results in factoid question answering. these methods do not scale for longer answers, where there are multiple sentences and few tokens don't hold the key to correctness. also, multiple generations needed to quantify the confidence score make them prohibitively expensive in a large scale settings. our goal is to estimate the correctness of the generated response in a single shot using uncertainty estimation. we prefer using ff-layer activations rather than token probabilities because token probabilities are computed by applying the decoder head (a linear projection) followed by a softmax transformation. this projection compresses the rich internal representation into a vocabulary space and the softmax operation further distorts the signal by normalizing it into a probability distribution (see figure 3 ), potentially obscuring fine-grained differences in the model's internal state. in contrast, raw activations preserve the high-dimensional representation prior to this compression, providing a more direct view of the model's internal dynamics during response generation. ## experimentation we have conducted experiments to identify the optimal masking ratio in order to maintain utility and precision of the system. ## discussion in this work, we present an approach for constructing a confidence score that aligns with the correctness of responses generated by large language models (llms). such a measure is particularly critical in high-stakes domains such as finance and healthcare, where the cost of an incorrect response far exceeds that of withholding a response. our method extends prior works in uncertainty quantification (uq) @xcite by leveraging model activation patterns to predict correctness more robustly. figure 3 illustrates our motivation for using raw activation signals from the feed-forward network (ffn) layers as auto-regressive features, rather than token logits or probabilities. token probabilities are obtained after a linear projection and softmax transformation. the projection step reduces dimensionality, discarding non-vocabulary-aligned features, while the softmax normalization saturates probability values, erasing scale information and compressing relative differences. using activations directly, we retain the full representational capacity of the internal state of the model. our application setting involves customer support agents consulting a proprietary knowledge base to resolve customer queries using specialized internal tools. the knowledge base contains documents vetted across multiple dimensions, including risk and legal compliance, making factual errors in the content highly unlikely. however, strict permissions govern which documents an agent can access. figure 5(a) shows the complexity of document formats and fine-grained entitlements that impact retrieval and downstream generation. we model confidence estimation as a classification problem over sequences of activations. specifically, we employ a lightweight recurrent neural network (lstm) that consumes ffn activations as auto-regressive signals. the classification logit from the lstm head serves as the confidence score (see figure 4 ). to enhance robustness against noisy supervision, we introduce a huber loss regularizer l huber alongside the cross-entropy loss l ce . the huber loss's ability to behave quadratically for small errors and linearly for large errors makes it well-suited for smoothing gradients and mitigating the influence of outliers @xcite . results in in real-world deployment, retrieval-augmented generation (rag) pipelines must meet strict latency requirements, as the llm prompt length is constrained by model context limits and thousands of queries are processed daily. tables 4 and 5 summarize our performance-latency trade-offs. reducing the number of llama 3.1 8b layers from 32 to 16 while keeping context size fixed preserves accuracy while reducing latency by approximately 42.5%. when the context size is reduced, alethic errors increase due to incomplete retrieval, raising the model's masking rate (i.e., instances where no answer is returned due to low confidence). nevertheless, the 16-layer configuration achieves comparable performance to the 32-layer setup at lower computational cost. we observe a slight improvement in response latency when hosting the model using vllm inference compared to hugging face's inference api, likely due to vllm's optimized memory management and continuous batching capabilities. overall, our approach leveraging ffn activations as auto-regressive signals, modeling them with an lstm, and regularizing with l huber proves effective in long-form rag settings. this method improves the trustworthiness of llm-generated responses and holds strong potential for safe deployment in sensitive, domain-specific applications. ## limitations our work pushes the boundary of confidence estimation in retrieval-augmented generation (rag) for sensitive domains, but several practical considerations remain. ideally, a rag system should generate both the response and its confidence score in a single pass. in our current implementation, the confidence score requires a second run of the system, which introduces additional computational and latency overhead. while this design choice enables deeper access to model internals, it also necessitates operating in a white-box setting, as the confidence model relies on activation signals from the llm to assess correctness. furthermore, the method is customized to the specific architecture of the target model, meaning that adaptation to other llms may require reconfiguration and retraining. these limitations also present opportunities for future research: integrating confidence estimation directly into the generation process, reducing computational cost, and developing architecture-agnostic approaches that preserve the performance benefits of activationbased probing methods. a limitation of this study is that the dataset cannot be made publicly available. the data contains sensitive and proprietary information pertaining to internal financial tools and knowledge resources used by service agents within a financial institution. this restriction is mandated by internal data governance policies to protect confidential and regulated financial information.