Modern large language models (LLMs) often exhibit overconfident errors, assigning high confidence to incorrect predictions and thereby undermining reliability in real-world use. Existing calibration methods typically rely on post-hoc adjustment or retraining, operating at the output level without addressing the underlying causes of miscalibration. We propose CoCHI, (Class-oriented Causal Head Intervention), a framework that improves calibration by identifying and intervening on internal components responsible for distinct confidence behaviours. By analysing predictions through correctness-confidence patterns, CoCHI enables targeted intervention on model internals, reducing overconfident errors while preserving correct predictions. Our approach provides a mechanistic perspective on calibration, linking reliability to identifiable structures within the model rather than treating it as an output-level issue. Experiments across multiple LLMs show consistent improvements in calibration metrics, including Expected Calibration Error, Negative Log-Likelihood, and Brier Score, while maintaining accuracy on five multiple-choice question answering benchmarks. These results suggest that effective calibration requires not only adjusting outputs, but also addressing the underlying causal mechanisms within the model.
Figures & tables
Figure 1: The framework of CoCHI.
MMLU
MedMCQA
Method
Accuracy ↑
ECE ↓
NLL ↓
BS ↓
Accuracy ↑
ECE ↓
NLL ↓
BS ↓
Gemma-2-9B
TS
0.6653 ± 0.0050
0.0660 ± 0.0173
0.5060 ± 0.0335
0.1683 ± 0.0125
0.5607 ± 0.0186
0.0556 ± 0.0261
0.6114 ± 0.0151
0.2122 ± 0.0066
ITI
0.6693 ± 0.0060
0.0605 ± 0.0062
0.5029 ± 0.0233
0.1675 ± 0.0105
0.5673 ± 0.0200
0.0451 ± 0.0140
0.6085 ± 0.0154
0.2122 ± 0.0070
Therm.
0.6653 ± 0.0050
0.1831 ± 0.0567
1.8729 ± 1.1455
0.2302 ± 0.0429
0.5607 ± 0.0186
0.2849 ± 0.0106
2.9242 ± 0.4188
0.3293 ± 0.0165
ATS
0.6653 ± 0.0050
0.3784 ± 0.0052
0.9191 ± 0.0066
0.3553 ± 0.0030
0.5607 ± 0.0186
0.3534 ± 0.0783
4.1559 ± 1.3606
0.3702 ± 0.0396
Table 1: Results on four benchmarks (MMLU, MedMCQA, OpenBookQA, and MedQA-USMLE). Best results are highlighted in bold.
Mask Rule
Masked Acc ↑
Masked ECE ↓
Wrong-high facilitating only
0.2044 ± 0.0038
0.4745 ± 0.1477
Correct-high non-facilitating only
0.2389 ± 0.0398
0.1518 ± 0.0474
Correct-high interfering only
0.4267 ± 0.0462
0.0694 ± 0.0488
Canonical union (broad)
0.2189 ± 0.0271
0.3850 ± 0.1816
Random size-matched
0.2111 ± 0.0077
0.4330 ± 0.2736
CoCHI
0.6120 ± 0.0626
0.0490 ± 0.0089
Table 2: Ablation study of different masking strategies on MMLU.
τc
Accuracy ↑
ECE ↓
NLL ↓
BS ↓
0.2
0.5933 ± 0.0217
0.0555 ± 0.0096
0.5385 ± 0.0075
0.1808 ± 0.0032
0.3
0.6060 ± 0.0059
0.0643 ± 0.0102
0.5319 ± 0.0043
0.1781 ± 0.0017
0.4
0.5860 ± 0.0118
0.0665 ± 0.0108
0.5276 ± 0.0058
0.1761 ± 0.0025
0.5
0.6120 ± 0.0626
0.0490 ± 0.0089
0.5168 ± 0.0236
0.1743 ± 0.0008
0.7
0.5873 ± 0.0377
0.0705 ± 0.0099
0.5548 ± 0.0212
0.1873 ± 0.0083
0.9
0.5933 ± 0.0081
0.0696 ± 0.0277
0.5255 ± 0.0059
0.1748 ± 0.0022
Table 3: Sensitivity analysis of the confidence threshold τc on MMLU.
Method
GPU
Cal. time (s) ↓
Artifact size ↓
Inf.
Prunable
TS
No
0.5
4 B
Negligible
No
ATS
No
25.8
16 B
Negligible
No
ITI
Yes
18.4
16 KB
O(k⋅d)
No
Therm.
No
19.4
1.0 MB
O(H⋅C)
No
CoCHI
Yes
69.9
40 B
None
Yes
Table 4: Efficiency comparison of calibration methods using 40 calibration samples on MMLU.
Figure 2: ECE comparison between the unmasked model and CoCHI on correct_low and wrong_high subsets.
Appendix figures & tables3 assets
Supplementary material from the paper’s appendix.
Appendix
Dataset
Domain
# Questions
MMLU
Multi-domain knowledge
14,042
MedQA-USMLE
Medical exam
11,451
OpenBookQA
Science
∼ 4,200
SciQ
Science
13,679
MedMCQA
Medical exam
187,005
Appendix
Table 5: Statistical summary of the datasets.
Method
Acc ↑
ECE ↓
NLL ↓
BS ↓
Gemma-2-9B
TS
0.9173 ± 0.0012
0.0579 ± 0.0184
0.4699 ± 0.4226
0.0641 ± 0.0094
ITI
0.9253 ± 0.0083
0.1027 ± 0.0091
0.2337 ± 0.0076
0.0617 ± 0.0025
Therm.
0.9223 ± 0.0021
0.0609 ± 0.0066
1.0836 ± 0.1824
0.0599 ± 0.0068
ATS
0.9173 ± 0.0012
0.0607 ± 0.0067
1.1578 ± 0.1183
0.0602 ± 0.0065
CoCHI
0.9353 ± 0.1069
0.0563 ± 0.0199
0.2309 ± 0.1276
0.0548 ± 0.0472
Appendix
Table 6: Results on SciQ. Best results are highlighted in bold.
Model
Cal. time (s) ↓
Min
Max
Peak GPU (GiB)
Mistral-7B-v0.1
39 ± 21
21
80
28.1
Llama-3.1-8B
58 ± 20
41
93
30.4
Gemma-2-9B
88 ± 23
72
159
46.2
Appendix
Table 7: Calibration time (s; mean ± std) of CoCHI across three model families, averaged over 15 runs per model. Costs include class-conditioned forward passes and exclude model loading, and scale with model size while remaining practical.