🤖 machine learning

Reading Calibrated Uncertainty from Language Model Trajectories

This paper proposes a method for improving language model uncertainty quantification by extracting scale-invariant geometric features from per-layer MLP update trajectories and feeding them into a sparse linear probe, which outperforms the standard maximum softmax probability approach while providing interpretable insights into how and where errors develop across the model's depth.

Original authors: Aliai Eusebi, Alexander Herzog, Xiaoyu Liang, Marie Vasek, Enrico Mariconti, Lorenzo Cavallaro

Published 2026-05-25

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Aliai Eusebi, Alexander Herzog, Xiaoyu Liang, Marie Vasek, Enrico Mariconti, Lorenzo Cavallaro

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a large language model (LLM) as a very smart, but sometimes overconfident, student taking a multiple-choice test. When the student finishes a question, they usually give you a "confidence score" (like saying, "I'm 95% sure this is the answer"). This is the standard way we check if the model is right or wrong.

However, the paper argues that this confidence score is often a lie. The model might say "95% sure" even when it's completely wrong. It's like a student who guesses wildly but feels very confident about their guess.

The Problem: Looking at the Destination, Not the Journey

Most methods for checking if an AI is confident only look at the final answer (the destination). They treat the model's thinking process as a single snapshot taken at the very end.

The authors of this paper say: "Wait a minute! To understand if someone is truly sure, you shouldn't just look at their final answer. You should watch how they got there."

They compare the model's internal thinking to a hiker walking up a mountain:

The Standard Method (MSP): Only looks at the hiker when they reach the summit. If they look happy at the top, we assume they had an easy, straight hike.
The New Method (Trajectory): Watches the hiker's entire path. Did they walk in a straight line? Did they wander off a cliff and have to double back? Did they zigzag wildly before finally stumbling onto the peak?

The paper claims that even if two hikers reach the same summit with the same happy expression, their paths might tell us very different stories. One might have had a smooth, confident climb (likely correct), while the other might have been lost, confused, and fighting the wind the whole time (likely wrong), even if they both end up smiling at the top.

The Solution: The "Motion Detector"

The researchers built a tool that acts like a motion detector for the model's internal thinking.

Tracking the Steps: As the model processes a question, it goes through many layers of "neurons" (like floors in a building). On each floor, the model makes a small adjustment to its answer.
Drawing the Map: Instead of just looking at the final result, the tool draws a map of the path the model took through these floors.
Measuring the Shape: The tool measures 11 specific things about this path, such as:
- Smoothness: Did the model change its mind suddenly?
- Direction: Did the model keep pushing in the same direction, or did it fight against its own previous thoughts?
- Efficiency: Did the model take a direct route, or did it wander in circles?

The Result: A Smarter "Stop" Button

Using this "motion map," the researchers trained a simple, transparent system (a "sparse linear probe") to predict errors.

Better than the Standard: This new system is much better at spotting when the model is lying about its confidence. In tests, it reduced the number of "high-confidence mistakes" significantly compared to the standard method.
The "Selective Abstention" Trick: Because this system is so good at spotting trouble, it can tell the model: "Hey, you think you're 95% sure, but your internal path was messy. I'm going to tell you to stop and not give an answer." This prevents the model from confidently giving wrong answers.
The "Why" Factor: The coolest part is that because the tool uses simple geometric measurements (like "how much did the path curve?"), we can actually see why it flagged an error. It can tell us, "The model started confident, then suddenly changed its mind in the middle layers, and then tried to force the answer back to the original idea." It's like a doctor being able to say, "The patient isn't just sick; their heart rate spiked at 2 PM and their temperature dropped at 3 PM," rather than just saying, "The patient is sick."

Summary

In short, this paper teaches us to stop trusting the model's final "I'm sure!" statement. Instead, we should watch the movie of how the model thought through the problem. By looking at the "motion" and "shape" of its internal steps, we can catch confident errors that would otherwise slip through the cracks, making AI safer and more reliable without needing expensive or complex new systems.

Technical Summary: Reading Calibrated Uncertainty from Language Model Trajectories

1. Problem Statement

Current approaches to Uncertainty Quantification (UQ) for Large Language Models (LLMs) with structured output face a trade-off between computational efficiency and calibration. The default method, Maximum Softmax Probability (MSP), is computationally free but often miscalibrated, assigning high confidence to incorrect predictions. Conversely, existing methods that probe internal activations typically treat hidden states as static snapshots (either from a single layer or averaged), discarding the layer-wise trajectory by which a representation is formed. This static view ignores how evidence accumulates, reinforces, or reverses across network depth, potentially obscuring uncertainty signals present in the computational path but not in the final output distribution.

2. Methodology

The authors propose a method that treats model activations not as static states but as representational trajectories. The approach involves three main components:

2.1. The Substrate: MLP Write-Vectors

Instead of analyzing raw hidden states, the method focuses on the MLP write-vectors ( $m_\ell$ ) contributed by each transformer block's Multi-Layer Perceptron (MLP) to the residual stream. Prior mechanistic interpretability work suggests MLPs are critical sites for factual knowledge storage.

The method captures the sequence of these vectors at the final prompt position (the readout position) during a single forward pass.
It constructs a discrete trajectory in representation space by computing the cumulative sum of these vectors: $s_\ell = \sum_{k \le \ell} m_k$ .
The trajectory is defined from the initial state ( $s_0 = 0$ ) to the final endpoint ( $s_L$ ).

2.2. Geometric Feature Extraction

The trajectory is summarized using eleven scale-invariant geometric features per layer, grouped into four categories:

Depth Distribution: Measures how computational effort (update magnitude) is distributed across layers (e.g., front-loading vs. back-loading).
Local Trajectory Shape: Captures directional consistency (consecutive cosine), abrupt changes (curvature), and whether updates reinforce or contradict the prior state (update-state alignment).
Relation to Endpoint: Measures how layers relate to the final direction, including convergence (direction to final), alignment of individual updates (update to final), and identification of layers that actively oppose the endpoint (contradictory support).
Efficiency: Measures the ratio of net displacement to total path length (cumulative coherence), indicating how directly the trajectory reaches its endpoint.

2.3. Sparse Linear Probe

These features are fed into a sparse linear probe (sparse logistic regression) to predict the probability of error.

Input: The vector of 11 features per layer ( $\phi(x)$ ) and the model's MSP score ( $p_{msp}$ ).
Interaction: The model includes interaction terms between the features and MSP, allowing the probe to weigh geometric signals differently based on the model's baseline confidence.
Training: The probe is trained to minimize class-balanced binary cross-entropy against an error indicator, regularized with an elastic-net penalty to enforce sparsity. This allows the model to identify predictive features directly without a separate selection stage.

3. Key Contributions

Trajectory-Based UQ: The authors propose a compact set of geometric features describing the evolution of representations across network depth. When fed to a sparse linear probe, this method outperforms MSP under selective abstention, with performance gains scaling with the baseline's miscalibration (up to 21 AURC points improvement).
End-to-End Interpretability: Unlike opaque dense probes on raw activations, every feature has a closed-form geometric meaning. The probe's coefficients reveal how errors occur: identifying which layers commit prematurely, which contradict the running state, and where trajectories drift from their endpoint.
Comprehensive Empirical Evaluation: The method was evaluated across 9 instruction-tuned LLMs (Qwen, Llama, DeepSeek families, ranging from 3B to 72B parameters) on five diverse NLP tasks (MMLU, CosmosQA, HellaSwag, HaluDial, HaluSum).

4. Experimental Results

Performance: The probe improved over MSP on 41 of 45 model-dataset pairs. The gains were most significant where MSP performed worst (high miscalibration), reducing AURC (Area Under the Risk-Coverage curve) by 11.45 to 21.83 points in the most difficult configurations.
Comparison to Baselines: The probe matched or outperformed a high-capacity "raw-activation probe" (trained on full hidden states) on 24 configurations and remained within 2 AURC points on 6 others, despite using only 11 scalar features.
Selectivity: The probe produces smooth, monotone risk-coverage curves, avoiding the unstable low-coverage spikes often seen with MSP. It effectively separates correct from incorrect predictions even when MSP confidence scores are identical.
Interpretability Findings:
- Correct vs. Incorrect: Predictions with identical MSP scores exhibit distinct trajectory signatures. For example, incorrect predictions often show premature alignment with the final direction in early layers, followed by a "double back" or drift.
- Task Specificity: The geometric signatures of errors vary by task. On HaluSum, errors manifested as premature commitment followed by a late, state-breaking correction. On CosmosQA, errors appeared as mid-depth drift orthogonal to the answer direction.
- Coefficient Analysis: The probe's coefficients localize error mechanisms. For instance, on HaluSum, error signals concentrated in the final quarter of the network, while on CosmosQA, they were spread across depth.

5. Significance and Claims

The paper claims that treating activations as motion rather than static snapshots offers a "tractable middle path" between the opacity of dense probes and the limited signal of MSP.

Auditable Uncertainty: Because the features are interpretable by construction, the resulting confidence signal allows practitioners to inspect not just whether a model is uncertain, but where in its computation that uncertainty arises (e.g., specific layers or geometric deviations).
Efficiency: The method adds negligible computational cost at inference, requiring only a single forward pass and simple geometric calculations, avoiding the high cost of sampling-based or ensemble approaches.
Limitations: The authors note that their analysis is restricted to discrete-choice settings (multiple-choice questions) and that the probe is currently fit per (model, task), with transferability across models and tasks remaining an open question.

In summary, the work demonstrates that the geometry of how a language model assembles a prediction contains rich, interpretable signals about uncertainty that are invisible to the final output distribution, enabling more reliable and auditable selective prediction.