TLDR: This research paper investigates how Large Language Models (LLMs) manage and express their confidence, known as calibration, throughout their internal processing layers. Contrary to previous beliefs that calibration primarily occurs in the final output layer, the study reveals a “confidence correction phase” in the upper layers where models actively recalibrate their predictions even after accuracy has stabilized. The authors also identify a specific “calibration direction” within the model’s internal data flow (residual stream) that can be adjusted to improve calibration without affecting accuracy, suggesting that confidence regulation is a distributed and dynamic process across the network’s depth.
Large Language Models (LLMs) have shown a remarkable ability to be well-calibrated, meaning their predicted probabilities align closely with the correctness of their answers. This is a surprising finding, especially when compared to earlier deep neural networks that often exhibited overconfidence. Previous research has pointed to specific components in the final layer of LLMs, such as ‘entropy neurons’ or the ‘null space’ of the unembedding matrix, as key players in this calibration.
However, a new study titled Calibration Across Layers: Understanding Calibration Evolution in LLMs offers a fresh perspective. Researchers Abhinav Joshi, Areeb Ahmad, and Ashutosh Modi from IIT Kanpur investigated how calibration isn’t just a final-layer phenomenon but rather a process that evolves throughout the entire depth of the network.
The Journey of Confidence: A Layer-by-Layer Look
The team analyzed several popular open-weight models, including Phi-2, LLaMA-3, LLaMA-2, and Mistral-7B, using the MMLU benchmark. They used a technique similar to the ‘Logit Lens’ to observe the internal workings of these models. Essentially, they looked at the ‘residual stream’—the main pathway of information flow—at each layer and projected it back to the vocabulary space to see what the model was ‘thinking’ and how confident it was at different stages.
Their findings revealed a consistent and fascinating pattern: while a model’s accuracy typically stabilizes in its middle layers (for instance, around layers 22-26 in Phi-2), its calibration scores (measured by Expected Calibration Error, ECE, and Maximum Calibration Error, MCE) continue to change significantly in the later layers. Initially, these scores might even worsen, indicating a phase of overconfidence, before sharply improving towards the final layers. The researchers termed this a ‘confidence correction phase’—a period where the model actively recalibrates its confidence, even after it has largely settled on its prediction.
Uncovering a ‘Calibration Direction’
The study also explored the role of the ‘unembedding matrix,’ which translates the model’s internal representations into final token probabilities. While previous work suggested its ‘null space’ (components with small singular values) might be involved in calibration, this research found that removing these components led to fluctuations in calibration, supporting their role but not as the sole mechanism.
Perhaps the most intriguing discovery was a specific ‘calibration direction’ within the residual stream. This low-dimensional direction, identified by analyzing the differences in successive layer outputs in the final layers, appears to govern how confidence is modulated. When the researchers intentionally perturbed the residual stream along this direction during inference, they observed a significant improvement in calibration metrics (lower ECE and MCE) without negatively impacting the model’s accuracy.
Remarkably, a calibration direction identified using the MMLU-Humanities dataset also generalized and improved calibration on other datasets, including TruthfulQA. This suggests the existence of a task-agnostic ‘calibration subspace’—a dedicated part of the model’s internal representation that it uses to regulate confidence, separate from the part responsible for making predictions.
Also Read:
- Unifying LLM Control: How In-Context Learning and Activation Steering Shape Model Beliefs
- Understanding Knowledge Dynamics in LLM Explanations with a New Framework
Implications for Understanding and Controlling LLMs
These findings challenge the notion that calibration is solely an output-layer property. Instead, it appears to be a dynamic and distributed process, shaped throughout the network’s forward pass. This new understanding could pave the way for more interpretable and controllable LLMs, allowing developers to fine-tune their confidence levels without compromising accuracy.
While the identified calibration directions showed promising results within individual models and some datasets, they didn’t directly generalize across all architectures (e.g., Mistral or LLaMA-2). This indicates that the specific mechanisms for confidence regulation might vary between different model designs, opening up new avenues for future research into more universal confidence-modulating features.
In essence, this work provides a deeper, layer-wise understanding of how LLMs manage their uncertainty, moving us closer to building more reliable and trustworthy AI systems.
