Works · Framework Demo
LLMs often assign high confidence to wrong answers, and most calibration methods only adjust output probabilities. Class-oriented Causal Head Intervention (CoCHI) looks inside the model instead: it links attention heads to patterns of correctness and confidence, then gates the heads behind overconfident errors while preserving correct predictions.
Correcting Overconfident Errors via Causal Head Intervention
Jing Ren†, Bowen Li†, Ziqi Xu*, Xuechao Yang, Feng Xia · † Equal contribution · * Corresponding author
How it works
Predictions on a small calibration set are partitioned by correctness and confidence to isolate overconfident errors.
Head gates are optimized on each class to identify facilitating and non-facilitating heads.
Error and noise heads form a mask, correct-high facilitators are protected, and the top-k heads are selected.
Selected head outputs are gated at inference, while the remaining computation and model weights stay unchanged.
Worked example
Illustrative questions, scores, and head layout — not a live model call, and not results from the paper. CoCHI is evaluated on multiple-choice QA; open-ended generation and multi-turn interaction remain future work. Classes are defined on labeled calibration data, not by access to test-time correctness. The paper's reported results are in the next section.
Reported results
Results from Tables 1 and 6 of the paper ↗, as published on the project page ↗. CoCHI selects 10 heads for LLaMA and 8 for Gemma and Mistral.
Model and benchmark pairs where CoCHI has the best mean of the 5 methods, counted from the table below.
| Method | Accuracy ↑ | ECE ↓ | NLL ↓ | Brier Score ↓ |
|---|
Mean ± standard deviation. ↑ higher is better, ↓ lower is better. Gold bold marks the best mean in each column. TS: Temperature Scaling. ITI: Inference-Time Intervention. ATS: Adaptive Temperature Scaling.