Calibrating Overconfidence Without Sacrificing Confidence: Probe-Conditioned Head Intervention for LLMs

Technical Summary: Probe-Conditioned Head Intervention for LLMs

Problem Statement
Large language models (LLMs) frequently exhibit "verbalized overconfidence," expressing high confidence in incorrect answers. Standard calibration techniques, such as temperature scaling or post-hoc instruction tuning, typically apply global adjustments to the output distribution. While these methods can improve aggregate calibration metrics, they suffer from a critical limitation: they cannot selectively suppress unwarranted confidence on wrong answers without simultaneously eroding warranted confidence on correct answers. This creates a tension where reducing false positives (wrong-but-confident) inevitably increases false negatives (correct-but-unconfident). The paper posits that a more selective approach is required—one that intervenes only when a response is likely to be wrong-but-confident, thereby decoupling the suppression of unwarranted confidence from the loss of warranted confidence.

Methodology: Probe-Conditioned Head Intervention (PCHI)
The authors propose PCHI, an inference-time method that operates internally within the model to selectively rescale attention-head outputs based on evidence of overconfidence. The method consists of three core components:

1. Structured Self-Evaluation: The model is instructed to generate a structured JSON response containing reasoning, answer, and a binary is_confident field (yes/no). The intervention targets the "wrong-yes" (WY) group (incorrect answers with high confidence) while aiming to preserve the "correct-yes" (CY) group.
2. Confidence-Template Probes: A frozen linear probe is trained to detect latent evidence of WY responses. This probe reads hidden states at specific "confidence-template" tokens (the token span between the answer and the final confidence value). The probe predicts the probability that a confident response is actually wrong. The probe layer ($\ell_p$) is selected in the middle of the network where WY/CY separability is strong, ensuring that downstream layers remain available for intervention.
3. Conditional Head Rescaling: During inference, the probe score gates a learned rescaling of downstream attention-head outputs.
   • For each attention head $h$ in layers downstream of $\ell_p$, a learned scalar coefficient $g_h$ (initialized to 1) is applied.
   • The intervention strength $s_{i,t}$ is determined by the probe score $p_{i,t}$. If the score exceeds a threshold $\tau$ (e.g., 0.5), the head output $z_{i,t,h}$ is rescaled: $z'_{i,t,h} = (1 + (g_h - 1)s_{i,t}) z_{i,t,h}$.
   • If the probe does not flag the response as likely wrong-confident ($s_{i,t} = 0$), the model behaves identically to the baseline.
   • An optional attention mask can be applied post-probe to restrict which context regions (prompt, reasoning, answer) the template queries attend to, potentially strengthening the intervention signal.

The learning objective uses a hinge-style loss to push WY examples below a confidence margin while protecting CY examples from being downgraded, regularized by an L1 penalty to keep unnecessary head coefficients near identity.

Key Contributions

• Selective Internal Intervention: Unlike global post-hoc calibration, PCHI intervenes conditionally based on internal evidence, allowing for the correction of specific error types (WY) without globally suppressing confidence.
• Decoupling of Calibration and Accuracy: The method demonstrates that unwarranted confidence can be reduced while largely preserving warranted confidence, addressing the trade-off inherent in uniform rescaling methods.
• Mechanistic Insight: The approach treats verbalized confidence as an internal computation process that can be controlled at specific points (template tokens) rather than just an output distribution to be mapped.

Experimental Results
The method was evaluated on OpenMathInstruct problems using Qwen3-4B-Instruct and Gemma3-4B-it models with a structured binary confidence field.

• Readout-Token Intervention: Applying PCHI at the final confidence token yielded significant improvements. On Qwen3-4B, the Expected Calibration Error (ECE) dropped from 21.9% to 11.7%, and the Area Under the Receiver Operating Characteristic (AUROC) increased from 66.5 to 90.3. Crucially, 82.2% of originally wrong-yes readouts were converted to "no," while only 10.8% of originally correct-yes readouts were damaged.
• Upstream Intervention: Intervening at upstream confidence-template tokens showed variable results. On Qwen3-4B, a joint intervention across tokens 1–5 reduced ECE to 9.2% with 63.3% WY correction and only 5.1% CY damage. On Gemma3-4B, upstream effects were weaker and more dependent on attention masking; restricting the visible context to the answer field improved WY correction to 40.4% but with a slight increase in CY damage.
• Ablation: Controls confirmed that the gains are not merely due to attention masking or generic activation steering. The combination of the probe's evidence and the learned head-specific rescaling is necessary for the observed selectivity.

Significance and Claims
The paper claims that PCHI offers a viable path toward models that can express uncertainty in a trustworthy manner without sacrificing their ability to express confidence when correct. By treating calibration as a selective control problem rather than a global mapping, the authors demonstrate that internal interventions can partially decouple the suppression of unwarranted confidence from the loss of warranted confidence.

The authors maintain a modest scope, noting that their results are specific to mathematical question answering with structured JSON outputs and binary confidence fields. They acknowledge that the method relies on model- and data-specific probe training and that extending this to free-form confidence expressions or open-ended tasks remains an open challenge. The work does not claim to provide a complete mechanistic explanation of verbalized confidence but rather offers evidence that such evidence is decodable and actionable for selective correction.