Interpreting and Mitigating Unwanted Uncertainty in LLMs

TL;DR

This work investigates unwanted uncertainty in LLMs, the phenomenon where a model revises a correct answer upon re-evaluation. By integrating a Flip-style re-evaluation prompt with a Needle-in-a-Haystack framework, the authors dissect internal mechanisms and show that retrieval heads are not the primary source of stability; instead, a small set of non-retrieval attention heads gate uncertainty. Targeted masking of these heads reduces flip behavior by up to 15% without harming core accuracy, establishing a practical mitigation grounded in mechanistic interpretability. However, downstream evaluations reveal trade-offs: improved confidence in easy tasks but potential overconfidence on harder ones, underscoring context-dependent effects. The findings advance understanding of internal uncertainty mechanisms and offer a comparatively simple, effective technique for increasing LLM reliability in high-stakes contexts.

Abstract

Despite their impressive capabilities, Large Language Models (LLMs) exhibit unwanted uncertainty, a phenomenon where a model changes a previously correct answer into an incorrect one when re-prompted. This behavior undermines trust and poses serious risks in high-stakes domains. In this work, we investigate the mechanisms that drive this phenomenon. We adapt the Needle-in-a-Haystack retrieval framework and integrate a Flip-style re-evaluation prompt to simulate realistic answer-flipping scenarios. We find that retrieval heads are not primarily responsible for avoiding uncertainty. Instead, we identify a small set of non-retrieval attention heads that disproportionately attend to misleading tokens in uncertain contexts. Masking these heads yields significant improvements, reducing flip behavior by up to 15% without introducing incoherence or overcorrection. However, when tested for downstream tasks, we observe trade-offs with flip behavior. Our findings contribute to the growing field of mechanistic interpretability and present a simple yet effective technique for mitigating uncertainty-driven failure modes in LLMs.

Figures (8)

• Figure 1: A visualization of the idea that is described in the paper. (Left) A real example that motivated our idea. The model answers a correct response on a multiple-choice question, but when it is asked if it is sure of its response, it changes its opinion to a false choice. (Right) An illustration of how we adapted this scenario to the Needle-In-A-Haystack task. We add the correct needle inside a large context, and then we ask the model to respond to a question that is answered by this needle. The model, again, responds correctly, but when it is asked to express its certainty, it changes its mind.
• Figure 2: Flow diagram indicating the extension of the needle-in-a-haystack experiment to interpret unwanted uncertainty in LLMs.
• Figure 3: The response-head attention configurations. We consider our setup under the assumption that the response is always correct. The diagram for incorrect cases is symmetrical and explored further in our experiments.
• Figure 4: Effect on confident responses with masking for 400 samples. The unmasked case gives 315 Yes responses and 85 No responses, indicating an inherent model uncertainty.
• Figure 5: Accuracy (% of "Yes" responses) for the different test needles.