🤖 machine learning

From token probabilities to calibrated confidence: An empirical study of mathematical question answering

This empirical study investigates methods for calibrating confidence in large language models for mathematical question answering, finding that while single-pass token probabilities offer limited signal, aggregating them and employing multi-pass strategies like self-verification or Monte Carlo Dropout, followed by post-hoc calibration techniques, can significantly improve the alignment between estimated confidence and actual accuracy.

Original authors: Avery Ma, Lorne Schell, Vin Bhaskara, Leila Pishdad

Published 2026-08-11

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Avery Ma, Lorne Schell, Vin Bhaskara, Leila Pishdad

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 you are playing a high-stakes game of trivia with a super-smart robot friend. You ask a tricky question, and the robot buzzes with an answer. But here's the catch: the robot is so eager to please that it often sounds 100% certain even when it's completely wrong. In the world of Artificial Intelligence, this is called "overconfidence." Scientists are trying to teach these robots how to say, "I'm pretty sure," or "I'm not so sure," so we know when to trust them. This is the challenge of confidence estimation: figuring out how likely a robot's answer is to be correct. To do this, researchers look at the robot's internal "whispers"—tiny numbers called token probabilities that show how likely the robot is to pick the next word in a sentence. The big question is: Can we turn these noisy whispers into a reliable "truth meter" that tells us exactly how much to trust the robot?

A team of researchers decided to put this idea to the test, specifically with math problems. They treated the AI like a student taking a test and asked: "If we listen closely to the student's thought process, can we tell if they got the answer right before they even finish writing it?" They compared two main ways of listening. The first was a "single-pass" listen, where they just analyzed the math problem and the final answer the robot gave in one go. The second was a "multi-pass" listen, where they asked the robot to double-check its work or run the same problem through its brain multiple times with slight variations to see if it got the same result.

Here is what they discovered. First, they found that if you only listen to the very last words of the answer (the final number), the robot sounds incredibly confident, even when it's wrong. It's like a student who writes "42" at the end of a test but got there by guessing. However, if you listen to the entire chain of reasoning—the step-by-step math the robot did to get there—you can spot tiny, consistent differences between right and wrong answers. By averaging these tiny signals across the whole thought process, they found a way to get a much better "truth meter" without needing to ask the robot to do extra work.

The researchers also tried some fancy tricks to get better signals. One trick was to ask the robot, "Are you sure this is right?" (a method called self-verification). They found that asking the robot to re-read its own answer from the very beginning was accurate but very slow and expensive, like making a student re-write their entire essay just to check one line. But they discovered a clever shortcut: instead of re-writing the whole thing, they could just append the question "Is this right?" to the very end of the original answer. This "in-situ" method was just as accurate but saved 88% of the computer effort, making it a huge win for efficiency.

Another trick involved "stochastic passes," where they asked the robot to solve the same problem multiple times with a tiny bit of randomness (like rolling a die) to see if it changed its mind. This method, called MC Dropout, was good at catching uncertainty, but it required the robot to do the work many times over, which is computationally heavy. The researchers found that while this method worked, it often wasn't worth the extra cost compared to the simpler "listen to the whole story" method.

Finally, they tested a "calibration" step, which is like teaching the robot to adjust its confidence dial. They used two classic math tools (Platt scaling and isotonic regression) to tweak the robot's confidence scores so they matched reality better. They found that these tools worked wonders, especially when they had a small set of practice problems (as few as 50 examples) to learn from. However, they also noticed that a calibration learned on easy math problems didn't always work well on hard ones, and a calibration learned on one type of robot didn't always transfer to a different type.

In short, the paper suggests that we don't need to make the robot work twice as hard to know if it's right. By simply paying closer attention to the entire path of its reasoning rather than just the final answer, and by using a clever, low-cost way to ask it to double-check itself, we can build much more trustworthy confidence meters. These findings suggest that with the right tuning, we can make AI systems much more honest about what they know and what they don't.

Title: From Token Probabilities to Calibrated Confidence: An Empirical Study of Mathematical Question Answering

Problem Statement
Large Language Models (LLMs) have achieved success in complex, multi-step tasks but are increasingly deployed in high-stakes domains where knowing when to trust an output is as critical as the output itself. While LLMs generate token probabilities during inference, these raw signals are often overconfident and poorly calibrated for determining the correctness of a completed response. This is because models are trained to optimize next-token prediction rather than to estimate the probability that a finished answer is correct. The paper investigates whether these readily available token probabilities can be transformed into well-calibrated confidence estimates for mathematical question answering (MQA), a domain with verifiable final answers.

Methodology
The authors conduct an empirical study categorizing confidence estimators by the number of inference passes required:

Single-Pass Estimators: These reuse token probabilities from the original generation without additional forward passes. The study analyzes two design choices:
- Token Subset: Comparing estimators based on the full generated sequence (reasoning trajectory + answer) versus those based only on the final answer tokens.
- Aggregation Method: Comparing the arithmetic mean of token probabilities against the length-normalized joint probability (geometric mean).
- Four specific estimators are instantiated: Sequence Joint (SEQ JOINT), Sequence Average (SEQ AVG), Answer Joint (ANS JOINT), and Answer Average (ANS AVG).
Multi-Pass Estimators: These obtain additional signals through repeated forward passes:
- Self-Verification ( $p(\text{True})$ ): The model is re-prompted to judge its own generated answer as "True" or "False." The authors also evaluate an in-situ variant that appends the verification query directly to the original generation trajectory, avoiding the re-encoding of the full context.
- Monte Carlo (MC) Dropout: The model performs multiple stochastic forward passes with dropout enabled while keeping the reasoning trajectory fixed. Uncertainty is measured via the Bayesian Active Learning by Disagreement (BALD) metric, capturing variation in the predictive distributions of answer tokens.
Post-Hoc Calibration: The study evaluates two classic methods to align confidence scores with empirical accuracy:
- Isotonic Regression: A flexible, non-parametric monotonic mapping.
- Platt Scaling: A parametric global adjustment (temperature scaling) applied to logits.

Key Contributions and Findings

Trajectory-Level Aggregation: The paper finds that aggregating probabilities over the full reasoning trajectory (SEQ AVG/JOINT) yields significantly more discriminative and better-calibrated confidence estimates than using only the final answer tokens. While individual token probabilities are often saturated, small but consistent probability differences accumulate across the reasoning path, distinguishing correct from incorrect generations.
Efficiency of In-Situ Verification: The proposed in-situ self-verification method achieves calibration comparable to standard re-prompting ( $p(\text{True})$ ) but reduces token-processing overhead by approximately 88%. This is achieved by reusing the original generation state rather than re-encoding the full question–response pair.
MC Dropout and Distributional Uncertainty: MC Dropout captures uncertainty reflected in the variation of predictive distributions, even when the decoded answer remains unchanged. However, this method requires substantial tuning (dropout rate selection) and multiple stochastic passes (roughly 5x overhead) to stabilize, making it computationally expensive compared to single-pass methods.
Calibration Efficacy and Data Efficiency: Both Isotonic Regression and Platt Scaling substantially reduce in-domain calibration error. Isotonic regression is particularly effective, often requiring as few as 50 calibration examples to yield significant improvements.
Asymmetric Transfer: Calibration mappings transfer asymmetrically across datasets and models. Mappings learned on easier datasets (e.g., GSM8K) generally improve calibration on harder datasets (e.g., GSMHard), but the reverse often degrades performance. Furthermore, transferability depends heavily on the specific model-dataset pair, suggesting that calibration data should ideally match the target model and difficulty level.

Significance and Limitations
The paper demonstrates that token probabilities, when properly aggregated and calibrated, can provide reliable confidence estimates for mathematical reasoning without requiring complex auxiliary models or fine-tuning. The findings highlight that the "overconfidence" of LLMs is not an inherent barrier but a calibration issue that can be addressed through sequence-level aggregation and simple post-hoc methods.

The authors note several limitations: the study is restricted to open-weight models where token probabilities and dropout controls are accessible; it focuses on MQA with verifiable answers rather than open-ended dialogue; and the MC Dropout results rely on specific design choices (e.g., fixed reasoning trajectories) that may not capture uncertainty over alternative reasoning paths. Additionally, the reliance on regex-based answer extraction may introduce errors for non-conforming outputs.

More like this