Calibration Across Layers: Understanding Calibration Evolution in LLMs

AI summary

Overview

Research area: Mechanistic interpretability and uncertainty quantification in large language models, specifically how model calibration (confidence versus correctness) develops across transformer depth.

Technical level: Intermediate. The paper assumes familiarity with transformer internals such as the residual stream, LayerNorm, the unembedding matrix, and the Logit Lens technique, though the calibration metrics (ECE, MCE) are defined from scratch in the background section.

Scope: A layer-by-layer empirical analysis of calibration in four open-weight LLMs evaluated primarily on MMLU, leading to the identification of a late-layer "confidence correction phase" and a low-dimensional "calibration direction" in the residual stream that can be perturbed to improve calibration without hurting accuracy.

What This Paper Is About

Prior work has explained LLM calibration through components of the final layer, such as entropy neurons and the null space of the unembedding matrix. This paper asks a complementary question: how does calibration actually evolve as information passes through the network's depth, not just at the output projection? The authors track accuracy, Expected Calibration Error (ECE), and Maximum Calibration Error (MCE) at every layer and find that confidence is actively adjusted in the later layers even after prediction accuracy has already stabilized.

Key Contributions

1. A layerwise analysis of calibration dynamics. The authors show that confidence is not simply correlated with accuracy but passes through a distinct confidence correction phase, in which models become temporarily overconfident before self-adjusting in later layers.

2. Identification of a "calibration direction" in the residual stream. They isolate a direction that governs confidence modulation and show that small perturbations along it improve ECE and MCE without sacrificing accuracy.

3. A complementary perspective to final-layer accounts. Rather than locating calibration solely in the unembedding matrix or output head, they show calibration behavior is distributed across network depth, including intermediate layers that have received limited attention.

4. Released code. The implementation is publicly available at https://github.com/Exploration-Lab/LLM-Calibration-Mechanism.

Main Findings

• A three-phase calibration pattern. The paper describes an initial decision formation phase, a subsequent phase of overconfidence, and a final confidence correction phase in the upper/later layers.

• Accuracy saturates before calibration settles. In Phi-2 on MMLU Humanities, performance stays near random (25%, for four-option questions) in the initial layers, begins to rise from layer 22, and saturates at layer 26, with only minor changes across layers 26-31. In the conclusion the authors state that accuracy plateaus beyond layer 24 in Phi-2.

• Calibration follows a different trajectory than accuracy. Figure 1 reports that ECE and MCE first rise across layers 25-28 and then decline across layers 28-31; the Figure 3 caption reports the rise as layers 26-28 and the decline as layers 29-31. The pattern holds across MMLU STEM, Humanities, Social Science, and Others, and similar trends are reported for Llama-3-8B, Mistral-7B, and Llama-2-7B in the appendix.

• Truncating the unembedding null space leaves accuracy intact but moves calibration. Reconstructing the unembedding matrix from only the top 85%, 90%, and 95% singular values produced largely unchanged accuracy but fluctuations in calibration metrics, especially MCE. The authors note the effects go in both directions across middle layers, so there is no clear indication that ECE/MCE simply increase when the null space is removed.

• The calibration direction is not the null space. The identified direction is shown not to be aligned with the low-singular-value (null space) directions of the unembedding matrix, indicating a distinct mechanism. A separate analysis of log eigenvalues shows writing not only in the null space but throughout.

• Perturbing the direction improves calibration. Adding the calibration direction to the residual stream shifts ECE and MCE to lower values across MMLU Humanities, MMLU STEM, MMLU Others, and TruthfulQA, while preserving accuracy.

• Cross-dataset transfer, but not cross-model transfer. A direction computed on the MMLU Humanities split also improved calibration on other MMLU subsets and on TruthfulQA. However, the direction does not generalize to other models such as Mistral or LLaMA-2.

• A different trend from vision models. The authors contrast their results with Wang and Zhang (2024), who found a U-shaped calibration trend in vision architectures such as ResNets and VGG, with better calibration in middle layers. Here, calibration is regulated in the later/upper layers before final predictions.

Methodology in Plain English

The models are evaluated on multiple-choice question answering, where the input consists of a query and a set of answer options, and the model must produce the correct answer-choice token next. The number of options depends on the dataset (four in the case of MMLU). Answer option order is randomized to reduce position bias.

To see inside the model, the authors use a Logit Lens-style approach: after each transformer block they take the residual stream representation, pass it through LayerNorm, and project it into the vocabulary space with the unembedding matrix. This produces a per-layer prediction distribution without training any additional classifiers, so no extra supervision is introduced. Accuracy, ECE, and MCE are computed at every layer, and reliability diagrams are used to visualize how confidence tracks empirical accuracy.

For the unembedding analysis, they apply Singular Value Decomposition to the unembedding matrix and reconstruct it while discarding the smallest 5%, 10%, and 15% of singular values (keeping the top 85%, 90%, and 95%), to test the role of the low-rank tail.

To find the calibration direction, they take the normalized differences between successive layer residual outputs and average the differences from the final three layers (layers 29, 30, 31). This direction captures the internal shift the model undergoes as calibration improves. They then test its functional role by adding a scaled version of it to the residual stream during inference, with a positive scaling coefficient, and re-measuring accuracy, ECE, and MCE.

Why This Matters

Impact on research. The work reframes calibration as a distributed, dynamic process rather than a property of the final projection. It extends prior findings on entropy neurons and the unembedding null space by showing that related confidence-regulating behavior appears at multiple depths, and it offers a mechanistic, unsupervised alternative to probing with trained classifiers. It also flags a practical caution: intermediate layers can show high accuracy but poor calibration, so they should not be trusted naively for downstream decision-making.

Real-world applications:

• Reliable confidence estimates for question-answering and knowledge-retrieval systems built on LLMs.
• Selective prediction and abstention, where a model defers to a human when confidence is unreliable.
• Early-exit and efficient inference schemes that need to know whether a layer's answer is trustworthy, not just correct.
• Post-hoc confidence adjustment in deployed models without retraining or degrading accuracy.

Industry relevance. A low-dimensional direction that improves calibration metrics while preserving accuracy is attractive for deployment because it suggests a cheap intervention rather than fine-tuning or expensive ensembling. The finding that intermediate-layer representations can be accurate but miscalibrated is directly relevant to systems that read from or act on intermediate computation, such as speculative decoding and early-exit serving.

Future Directions

• Extending the analysis beyond single-token multiple-choice classification to multi-token generation, where calibration over token sequences and temporal dynamics would need new metrics.
• Finding more universal, architecture-agnostic confidence-modulating features, since the discovered direction does not transfer to Mistral or LLaMA-2 and appears partly model- and domain-specific.
• Disentangling calibration from competence on reasoning-based datasets, where accuracy increases gradually rather than saturating in mid-to-late layers, making the confidence correction phase harder to isolate.
• Developing more principled methods for finding the calibration direction, such as optimization, gradient sensitivity to ECE loss, or attribution techniques, instead of the simple layerwise residual difference used here.
• Investigating how these layerwise calibration mechanisms arise during pretraining and whether similar correction dynamics generalize across model families and sizes.

Target Audience

Researchers working on mechanistic interpretability, uncertainty quantification, and calibration in large language models, as well as practitioners who need trustworthy confidence estimates from deployed models. Readers without background in transformer internals will need the background section (residual stream, LayerNorm, unembedding matrix, ECE/MCE definitions) before the results are accessible.

Authors’ abstract

Large Language Models (LLMs) have demonstrated inherent calibration capabilities, where predicted probabilities align well with correctness, despite prior findings that deep neural networks are often overconfident. Recent studies have linked this behavior to specific components in the final layer, such as entropy neurons and the unembedding matrix null space. In this work, we provide a complementary perspective by investigating how calibration evolves throughout the network depth. Analyzing multiple open-weight models on the MMLU benchmark, we uncover a distinct confidence correction phase in the upper/later layers, where model confidence is actively recalibrated after decision certainty has been reached. Furthermore, we identify a low-dimensional calibration direction in the residual stream whose perturbation significantly improves calibration metrics (ECE and MCE) without harming accuracy. Our findings suggest that calibration is a distributed phenomenon, shaped throughout the network forward pass, not just in its final projection, providing new insights into how confidence-regulating mechanisms operate within LLMs.