Online Reasoning Calibration: Test-Time Training Enables Generalizable Conformal LLM Reasoning

Cai Zhou; Zekai Wang; Menghua Wu; Qianyu Julie Zhu; Flora C. Shi; Chenyu Wang; Ashia Wilson; Tommi Jaakkola

cs.LGApr 1, 2026

Online Reasoning Calibration: Test-Time Training Enables Generalizable Conformal LLM Reasoning

Authors: Cai Zhou, Zekai Wang, Menghua Wu, Qianyu Julie Zhu, Flora C. Shi, Chenyu Wang, Ashia Wilson, Tommi Jaakkola, +1 more

Organizations: Department of Electrical Engineering and Computer Science (MIT EECS) · Computer Science and Artificial Intelligence Laboratory (MIT CSAIL) · Laboratory for Information and Decision Systems (MIT LIDS) · Computational Science and Engineering (MIT CSE)

Abstract

While test-time scaling has enabled large language models to solve highly difficult tasks, state-of-the-art results come at exorbitant compute costs. These inefficiencies can be attributed to the miscalibration of post-trained language models, and the lack of calibration in popular sampling techniques. Here, we present Online Reasoning Calibration (ORCA), a framework for calibrating the sampling process that draws upon conformal prediction and test-time training. Specifically, we introduce a meta-learning procedure that updates the calibration module for each input. This allows us to provide valid confidence estimates under distributional shift, e.g. in thought patterns that occur across different stages of reasoning, or in prompt distributions between model development and deployment. ORCA not only provides theoretical guarantees on conformal risks, but also empirically shows higher efficiency and generalization across different reasoning tasks. At risk level $δ=0.1$ , ORCA improves Qwen2.5-32B efficiency on in-distribution tasks with savings up to 47.5% with supervised labels and 40.7% with self-consistency labels. Under zero-shot out-of-domain settings, it improves MATH-500 savings from 24.8% of the static calibration baseline to 67.0% while maintaining a low empirical error rate, and the same trend holds across model families and downstream benchmarks. Our code is publicly available at https://github.com/wzekai99/ORCA.

Figures & tables

Symbol	Learned in	Description
$\phi_{t}$	–	Step embedding: mean-pooled LLM hidden state at step $t$
$C_{t}$	–	Step label: correctness or consistency; $C_{t}{=}0$ at inference
$W_{0},b_{0}$	outer loop	Initial probe weights, learned via meta-training
$W_{t},b_{t}$	inner loop	Probe weights at step $t$ , updated online during inference
$\theta_{Q,K}$	outer loop	Unified outer view parameters (optional learned projections or identity)
$\eta$	outer loop	Inner learning rate; fixed or learned

Table 1: Notation summary.

Figure 1: Framework of Online Reasoning Calibration (ORCA).

Supervised labels
	$\delta=0.05$		$\delta=0.1$		$\delta=0.15$		$\delta=0.2$
Method	Sav.	Err.	Sav.	Err.	Sav.	Err.	Sav.	Err.
Static Probe	.220	.055	.380	.105	.512	.159	.625	.208
TTT no-QK	.282	.053	.475	.110	.575	.152	.673	.192
TTT QK ( $d_{h}{=}128$ )	.233	.046	.414	.103	.560	.150	.674	.204
Consistent labels (no ground truth)

Table 2: In-distribution early-stopping performance on the 5K test set (Qwen2.5-32B,

\epsilon{=}0.05

). TTT-Probe (no-QK) improves savings by 24.9% relative over the static baseline.

Supervised labels
	MATH-500		GPQA		AIME’24		AIME’25		AIME’26
Method	Sav.	Err.	Sav.	Err.	Sav.	Err.	Sav.	Err.	Sav.	Err.
Static Probe	.248	.008	.643	.270	.158	.050	.139	.000	.147	.050
TTT no-QK	.637	.023	.715	.300	.293	.150	.265	.056	.198	.050
TTT QK ( $d_{h}{=}128$ )	.670	.021	.665	.210	.295	.100	.258	.000	.134	.050
Consistent labels (no ground truth)

Table 3: OOD generalization at

\delta{=}0.1

. The TTT-Probe achieves 2.6–2.7

\times

the baseline savings on MATH-500 under supervised labels.

	Qwen2.5-32B		QwQ-32B		Llama-3.3-70B
Method	Sav.	Err.	Sav.	Err.	Sav.	Err.
Static Probe	.380	.105	.295	.094	.354	.104
TTT no-QK	.475	.110	.394	.081	.424	.090
TTT QK ( $d_{h}{=}128$ )	.414	.103	.376	.076	.378	.081

Table 4: Cross-model results (

\delta{=}0.1

, supervised). TTT-Probe consistently outperforms the static baseline across all three model families.

	Supervised				Consistent
	In-dist		OOD avg		In-dist		OOD avg
Configuration	Sav.	Err.	Sav.	Err.	Sav.	Err.	Sav.	Err.
Static Probe	.380	.105	.267	.076	.345	.098	.241	.093
TTT no-QK	.475	.110	.422	.116	.407	.096	.323	.099
TTT QK ( $d_{h}{=}128$ )	.414	.103	.404	.076	.397	.113	.341	.076

Table 5: Supervised vs. consistent comparison (

\delta{=}0.1

), including average OOD savings and error across five benchmarks.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 2: Compute savings vs. risk tolerance

\delta

for supervised (left) and consistent (right) labels (Qwen2.5-32B). TTT no-QK consistently outperforms the baseline across all risk levels, with the largest gap at low

\delta

Configuration	Architecture	Training	Online update	Savings	Error
Full TTT (no-QK)	Linear	Meta-learn	✓	.475	.110
Standard (no-QK)	Linear	Supervised		.239	.095
Full TTT (QK, $d_{h}{=}128$ )	QK proj.	Meta-learn	✓	.414	.103
No meta-training ∗	QK proj.	None	✓	.254	.099
No meta + no update ∗	QK proj.	None		.173	.091
Static Probe	PCA+LogReg	Supervised		.380	.105

Appendix

Table 6: Core mechanism ablation (supervised,

\delta{=}0.1

). Full TTT achieves the best savings: standard supervised training on the same architecture underperforms the static baseline, and removing either meta-training or online updates substantially degrades performance. Rows marked with

*

use random initialization.

			OOD Savings
Variant	Sav.	Err.	MATH	GPQA	A’24	A’25	A’26
QK ( $d_{h}{=}128$ )	.414	.103	.670	.665	.295	.258	.134
+ LayerNorm	.451	.095	.697	.726	.265	.256	.144
+ LN + Residual	.450	.094	.697	.726	.265	.256	.144
+ Shared QK	.449	.094	.698	.724	.264	.256	.144
+ Learnable $\eta$	.421	.109	.679	.656	.364	.260	.104

Appendix

Table 7: Architecture ablation (supervised,

d_{h}{=}128

\delta{=}0.1

$d_{h}$	Parameters	Savings	Error
32	328K	.440	.105
64	656K	.439	.106
128	1.3M	.414	.103
256	2.6M	.408	.104
512	5.2M	.398	.099
no-QK	5.1K	.475	.110

Appendix

Table 8: Effect of QK projection dimension (supervised,

\delta{=}0.1

Qwen2.5-32B
Configuration	Step Sav.	Token Sav.	$\Delta$
Static Probe	.377	.377	.000
TTT no-QK	.475	.471	$-$ .004
TTT QK ( $d_{h}{=}128$ )	.414	.414	.000
QwQ-32B
Static Probe	.293	.296	+.003

Appendix

Table 9: Step-level vs. token-level savings (

\delta{=}0.1

, supervised). The two metrics are highly consistent across models.

Epoch	10	20	30	40	50
no-QK	.443	.475	.471	.476	.464
QK ( $d_{h}{=}128$ )	.414	.387	.381	.357	.324

Appendix

Table 10: Savings at selected epochs (supervised,

\delta{=}0.1

). The no-QK probe is stable; QK peaks early and overfits.

Figure 3: Actual error rate vs. target risk

\delta

(supervised, Qwen2.5-32B). All methods track the diagonal, confirming valid risk control. Points below the diagonal satisfy the LTT guarantee.

Figure 4: Distribution of per-problem savings at

\delta{=}0.1

(supervised, Qwen2.5-32B, 902 problems). Solid lines: mean; dashed lines: median. TTT no-QK shifts the distribution toward higher savings across the full range.

Figure 5: Probe score trajectories for a test problem (Qwen2.5-32B,

\delta{=}0.1

). The green line marks the first correct step. The static probe (top) never crosses its threshold and saves 0%. The TTT no-QK probe (bottom) crosses the threshold at step 22 and saves 41%.

$n_{\text{train}}$	500	1000	1500	2000	2500	3000
Savings	.385	.424	.438	.450	.458	.475
Error	.092	.093	.101	.102	.111	.110

Appendix

Table 11: Effect of training set size (top,

n_{\text{cal}}{=}1

K fixed) and calibration set size (bottom,

n_{\text{train}}{=}3

K fixed) on in-distribution performance (TTT no-QK, supervised,

\delta{=}0.1

). Savings grow monotonically with

n_{\text{train}}

; small calibration sets yield more conservative thresholds, and risk control is maintained at every size.

Dataset	Savings	Error
In-distribution	.474 $\pm$ .004	.112 $\pm$ .003
MATH-500	.631 $\pm$ .005	.022 $\pm$ .001
GPQA-Diamond	.715 $\pm$ .002	.300 $\pm$ .000
AIME’24	.292 $\pm$ .004	.150 $\pm$ .000
AIME’25	.267 $\pm$ .005	.056 $\pm$ .000
AIME’26	.199 $\pm$ .002	.050 $\pm$ .000

Appendix

Table 12: Mean

\pm

std over 10 random seeds with independent LTT calibration per seed (TTT no-QK, supervised,

\delta{=}0.1

). Variance is small on all benchmarks.

Explore similar work

Unsupervised Confidence Calibration for Reasoning LLMs from a Single Generation
Apr 21, 2026Thomas Zollo, Jimmy Wang, Richard ZemelConfidence Estimation in Language Models Language Model Calibration
Columbia University
Test-time Calibration Learning for Large Language Model Reasoning
Oct 2, 2026Zizhuo Zhang, Xiong Peng, Jingwei Sun +4Language Model Calibration Confidence Estimation in Language Models
TMLR Group, Department of Computer Science, Hong Kong Baptist University · Huawei Noah’s Ark Lab.
Process Supervision of Confidence Margin for Calibrated LLM Reasoning
Apr 25, 2026Liaoyaqi Wang, Chunsheng Zuo, William Jurayj +2Confidence Estimation in Language Models RL for Language Model Reasoning
Johns Hopkins University

Apr 21, 2026J/K move · Enter open · S save

Unsupervised Confidence Calibration for Reasoning LLMs from a Single Generation

Thomas Zollo, Jimmy Wang, Richard Zemel

Columbia University

Confidence Estimation in Language Models Language Model Calibration LLM Reliability LLM Uncertainty Estimation Self-Consistency in Language Models Large Reasoning Models

Open paper AskPDF arXiv

Reasoning language models can solve increasingly complex tasks, but struggle to produce the calibrated confidence estimates necessary for reliable deployment. Existing calibration methods usually depend on labels or repeated sampling at inference time, making them impractical in many settings. We introduce a method for unsupervised confidence calibration of reasoning LLMs when only a single generation is available at inference time. Our approach uses offline sampling on unlabeled data to derive a self-consistency-based proxy target, then distills this signal into a lightweight deployment-time confidence predictor. In a broad evaluation across 5 math and question-answering tasks using 9 reasoning models, our method substantially outperforms baselines, including under distribution shift, and improves downstream performance in selective prediction and simulated downstream decision-making.

Abstract

Explore similar work

Unsupervised Confidence Calibration for Reasoning LLMs from a Single Generation

Test-time Calibration Learning for Large Language Model Reasoning

Process Supervision of Confidence Margin for Calibrated LLM Reasoning