Fixing Overconfidence in LLMs with Adaptive Temperature Scaling

Large Language Models (LLMs) have revolutionized artificial intelligence, demonstrating remarkable fluency and reasoning capabilities. However, a persistent issue plagues even the most advanced models: calibration.

Ideally, when an LLM says it is 80% confident in an answer, it should be correct 80% of the time. Unfortunately, this is rarely the case. Modern LLMs, particularly those fine-tuned with Reinforcement Learning from Human Feedback (RLHF), tend to be notoriously overconfident. They might hallucinate a completely incorrect fact but assign it a 99% probability score. In high-stakes fields like medicine, law, or autonomous coding, this disconnect between confidence and accuracy is dangerous.

While unsupervised pre-training often yields well-calibrated models, the very process that makes models helpful—RLHF—breaks this calibration.

In this deep dive, we explore a research paper from Stanford University, “Calibrating Language Models with Adaptive Temperature Scaling.” The researchers introduce a novel, post-hoc method called Adaptive Temperature Scaling (ATS). Unlike traditional methods that apply a “one-size-fits-all” fix, ATS dynamically adjusts the model’s confidence for each specific token it generates, restoring reliability without sacrificing performance.

The Problem: RLHF and the Calibration Gap

To understand the solution, we must first understand the damage caused by fine-tuning.

Base models (like a raw Llama-2 pre-trained on the web) are generally well-calibrated. They learn to predict the next token based on statistical likelihoods inherent in the training data. However, these base models aren’t very good at following instructions.

To fix this, we use RLHF. This aligns the model with human preferences, making it helpful and chatty. But recent studies show that this alignment comes at a cost. RLHF pushes the model distributions away from the natural probabilities of language, resulting in miscalibration. The model becomes a “people pleaser”—it learns to sound authoritative and confident, even when it is wrong.

Why Traditional Fixes Fail

The standard solution for miscalibration in machine learning is Temperature Scaling (TS).

In a neural network, the final layer produces “logits” (unnormalized scores). These logits are usually passed through a Softmax function to turn them into probabilities (confidence scores). Temperature Scaling involves dividing these logits by a single scalar value, \(T\) (temperature), before the Softmax.

• If \(T > 1\), the distribution flattens (lower confidence).
• If \(T < 1\), the distribution sharpens (higher confidence).

In traditional classification tasks (like identifying cats vs. dogs), finding a single optimal \(T\) for the whole validation set works well.

However, Language Modeling is not Image Classification. In language generation, the “difficulty” varies wildly from token to token. Predicting that “The capital of France is…” will be followed by “Paris” is easy. Predicting the result of a complex math problem or an obscure historical fact is hard.

A single, global temperature parameter assumes the model’s level of overconfidence is constant across all contexts. This is a false assumption for LLMs. Some topics trigger massive overconfidence, while others do not. Using a single scalar fails to capture the dynamic nature of language.

The Solution: Adaptive Temperature Scaling (ATS)

The core innovation of this paper is moving from a static parameter to a dynamic, predicted one. Instead of learning one temperature \(T\) for the entire model, ATS predicts a unique temperature \(\tau\) for every single token.

The Architecture

How does the model know what temperature to use? It looks at the context.

The authors propose using a “Calibration Head”—a small auxiliary neural network that sits on top of the main LLM.

1. Input: The main LLM processes the input text and generates hidden states (vectors representing the semantic meaning of the text).
2. Features: ATS takes the last hidden state of the model as its input features. The researchers found that hidden states contain much richer information about uncertainty than the raw logits alone.
3. Prediction: The Calibration Head processes these hidden states and outputs a scalar value \(\tau\) for that specific step.
4. Scaling: The original logits \(\hat{z}\) are effectively scaled by this predicted temperature to produce calibrated logits \(\hat{q}\). \[ \hat{q} = \hat{z} \circ e^{\tau} \] Because \(\tau\) is predicted from the hidden states, the calibration adapts to the nuance of the current sentence. If the model is entering a domain where it typically hallucinates (like citing fake case law), the Calibration Head can detect those patterns in the hidden states and lower the temperature to reduce overconfidence.

The Special Sauce: Selective Smoothing Loss

Designing the architecture is only half the battle. How do you train this Calibration Head? You need a loss function that specifically targets calibration.

Standard Cross-Entropy loss (used to train the LLM itself) encourages the model to be 100% confident in the correct token. However, when calibrating, simply maximizing likelihood can be counterproductive if the model is fundamentally wrong.

The authors introduce a loss function called Selective Smoothing.

Here is how to interpret this equation:

1. Case 1: The Model is Correct (\(\arg \max \hat{q} = y\)). If the model’s top prediction is actually the correct token, we use a standard cross-entropy term. We want the model to be confident here. The term \(-(1-\alpha)\) weights this part of the loss.

2. Case 2: The Model is Incorrect. If the model predicts the wrong token as the most likely option, we do not want to maximize the probability of the correct label indiscriminately, nor do we want to sharpen the distribution on the wrong answer. Instead, the target becomes a Uniform Distribution (1/|V|). This effectively tells the model: “If you are wrong, you should not be confident. You should be uncertain.” This flattens the distribution, reducing overconfidence.

This “selective” nature helps the model distinguish between justified confidence and unjustified overconfidence.

Experimental Results

The researchers evaluated ATS using Llama-2-Chat (7B and 13B) and Qwen-Chat-7B. They trained the calibration head on the Alpaca GPT-4 dataset, which consists of diverse instruction-following data.

Crucially, they tested the calibration on entirely different downstream benchmarks: MMLU (Massive Multitask Language Understanding), TriviaQA, and TruthfulQA. This tests whether the calibration generalizes to new tasks.

Quantitative Improvements

The metrics used were Expected Calibration Error (ECE) (lower is better) and Brier Score (BS) (lower is better).

The results, shown below, are striking.

Let’s look at Llama-2-7b-Chat:

• Uncalibrated ECE: 0.298 (Very miscalibrated)
• Standard Temperature Scaling ECE: 0.270 (Minor improvement)
• ATS (Ours) ECE: 0.125

ATS reduced the calibration error by over 50% compared to the uncalibrated model, and significantly outperformed vector scaling and binning methods. This pattern holds true across different models and datasets. It is worth noting that ATS achieves this without changing the accuracy of the model—it simply makes the confidence scores honest.

Visualizing Calibration: Reliability Diagrams

Numbers are good, but Reliability Diagrams show the full story. In these graphs:

• X-axis: The confidence the model claims (e.g., “I am 90% sure”).
• Y-axis: The actual accuracy (e.g., “I was right 90% of the time”).
• Blue Line: Perfect calibration (Confidence = Accuracy).
• Bars: The actual performance of the model in those confidence bins.

Llama-2-7b-Chat (Before Calibration):

In the uncalibrated MMLU diagram above, notice how the red bars are consistently below the blue line. When the model says it is 90% or 100% confident (far right), its actual accuracy is much lower. This is classic overconfidence.

Llama-2-7b-Chat (After ATS Calibration):

After applying ATS, the bars align beautifully with the blue diagonal line. When the model claims high confidence, it is actually correct. When it is unsure, it lowers its confidence score.

The improvement is even more visible on TriviaQA, a dataset requiring specific factual knowledge.

Before (TriviaQA):

After (TriviaQA):

Notice in the “After” image how the distribution of confidence spreads out. The model is no longer blindly confident; it utilizes the lower confidence bins (left side of the graph) more frequently when it encounters difficult questions.

Why Does It Work? Ablation Studies

The researchers performed ablation studies to justify their design choices. They tested different smoothing methods, loss weights, and head architectures.

1. The Importance of Selective Smoothing

They compared their “Selective Smoothing” against standard Cross-Entropy (No smoothing) and Full Label Smoothing (always smoothing, regardless of correctness).

The data shows that Selective Smoothing is superior (ECE 0.125 vs 0.226 for no smoothing). Standard cross-entropy struggles because, in an effort to minimize loss on incorrect predictions, it might paradoxically increase confidence in wrong answers or fail to flatten the distribution sufficiently.

2. Loss Weighting (\(\alpha\))

The hyperparameter \(\alpha\) controls how much weight is given to the uniform distribution target when the model is wrong.

A higher \(\alpha\) (around 0.5) is necessary. This suggests that to correct the extreme overconfidence of RLHF models, the calibration method must aggressively penalize confidence when the model makes mistakes.

3. Head Architecture

Finally, does the complexity of the calibration head matter? Can we just use a simple linear layer?

Using a Transformer layer works best. This confirms that calibration requires context. The head needs to aggregate hidden state values from prior tokens to understand the nuance of the current prediction. A simple linear projection of the current token’s state isn’t enough to capture the complexity of “why” the model is uncertain.

Broader Implications and Conclusion

The “Adaptive Temperature Scaling” method presents a significant step forward in making Large Language Models trustworthy.

By accepting that LLM uncertainty is dynamic and context-dependent, ATS moves beyond the limitations of static temperature scaling. It treats calibration as a supervised learning problem, utilizing the rich internal representations of the model to predict how confident it should be.

Key Takeaways:

1. RLHF breaks calibration: Making models “helpful” makes them overconfident.
2. Context matters: You cannot use a single temperature for every token. Some tokens are harder to predict than others.
3. Hidden States are key: The internal state of the model contains the cues necessary to predict its own reliability.
4. Selective Loss: Training the calibration head requires a special loss function that penalizes confidence specifically when the model is wrong.

For students and practitioners, this paper highlights a crucial layer of the AI stack: Reliability. As we deploy these models into the real world, it is not enough for them to be smart; they must also know what they don’t know. ATS offers a computationally efficient way (adding just one small layer) to achieve that self-awareness.