Improving Confidence Estimation in Language Models with Linear Representations of Ambiguity

As language models (LMs) are increasingly deployed in downstream applications, accurately estimating their likelihood of success on a given example is critical. However, LMs are poor estimators of their own likelihood of success. We hypothesize that calibration is linked to the detection of aleatoric uncertainty, and that LMs possess representations of this uncertainty in the form of linear representations of input ambiguity. Using difference-in-means on minimal pairs of ambiguous and disambiguated sentences, we isolate ambiguity directions in LM activations, and demonstrate that their causal influence selectively modulates the model's judgment of ambiguity. We then show that inference-time steering along ambiguity directions reliably improves confidence estimation across diverse LM families on multiple-choice QA benchmarks, including MMLU and TruthfulQA. Finally, to eliminate the need for inference-time interventions, we propose a LoRA procedure that encourages updates along ambiguity directions; this approach encodes calibration improvements into the parameters of the model, yields further gains in confidence estimation over steering alone, and maintains model accuracy on QA tasks.