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)
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
ϕt
–
Step embedding: mean-pooled LLM hidden state at step t
Ct
–
Step label: correctness or consistency; Ct=0 at inference
W0,b0
outer loop
Initial probe weights, learned via meta-training
Wt,bt
inner loop
Probe weights at step t , updated online during inference
θQ,K
outer loop
Unified outer view parameters (optional learned projections or identity)
η
outer loop
Inner learning rate; fixed or learned
Table 1: Notation summary.
Figure 1: Framework of Online Reasoning Calibration (ORCA).
δ=0.05
δ=0.1
δ=0.15
δ=0.2
Method
Sav.
Err.
Sav.
Err.
Sav.
Err.
Sav.
Err.
Supervised labels
Static Probe
.220
.055
.380
.105
.512
.159
.625
.208
TTT no-QK
.282
.053
.475
.110
.575
.152
.673
.192
TTT QK ( dh=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, ϵ=0.05 ). TTT-Probe (no-QK) improves savings by 24.9% relative over the static baseline.
MATH-500
GPQA
AIME’24
AIME’25
AIME’26
Method
Sav.
Err.
Sav.
Err.
Sav.
Err.
Sav.
Err.
Sav.
Err.
Supervised labels
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 ( dh=128 )
.670
.021
.665
.210
.295
.100
.258
.000
.134
.050
Consistent labels (no ground truth)
Table 3: OOD generalization at δ=0.1 . The TTT-Probe achieves 2.6–2.7 × 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 ( dh=128 )
.414
.103
.376
.076
.378
.081
Table 4: Cross-model results ( δ=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 ( dh=128 )
.414
.103
.404
.076
.397
.113
.341
.076
Table 5: Supervised vs. consistent comparison ( δ=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 δ 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 δ .
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, dh=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, δ=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.
Table 8: Effect of QK projection dimension (supervised, δ=0.1 ).
Configuration
Step Sav.
Token Sav.
Δ
Qwen2.5-32B
Static Probe
.377
.377
.000
TTT no-QK
.475
.471
− .004
TTT QK ( dh=128 )
.414
.414
.000
QwQ-32B
Static Probe
.293
.296
+.003
Appendix
Table 9: Step-level vs. token-level savings ( δ=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 ( dh=128 )
.414
.387
.381
.357
.324
Appendix
Table 10: Savings at selected epochs (supervised, δ=0.1 ). The no-QK probe is stable; QK peaks early and overfits.
Figure 3: Actual error rate vs. target risk δ (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 δ=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, δ=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%.
ntrain
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, ncal=1 K fixed) and calibration set size (bottom, ntrain=3 K fixed) on in-distribution performance (TTT no-QK, supervised, δ=0.1 ). Savings grow monotonically with ntrain ; small calibration sets yield more conservative thresholds, and risk control is maintained at every size.
Dataset
Savings
Error
In-distribution
.474 ± .004
.112 ± .003
MATH-500
.631 ± .005
.022 ± .001
GPQA-Diamond
.715 ± .002
.300 ± .000
AIME’24
.292 ± .004
.150 ± .000
AIME’25
.267 ± .005
.056 ± .000
AIME’26
.199 ± .002
.050 ± .000
Appendix
Table 12: Mean ± std over 10 random seeds with independent LTT calibration per seed (TTT no-QK, supervised, δ=0.1 ). Variance is small on all benchmarks.
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.
Reliable large language models (LLMs) must not only produce accurate answers but also express confidence that faithfully reflects their probability of being correct. Such calibration is essential for identifying uncertain predictions and supporting reliable decision-making in real-world deployment. Recent studies incorporate calibration learning into reinforcement learning (RL), jointly optimizing answer correctness and verbalized confidence using ground-truth correctness supervision. However, their reliance on labeled data limits their applicability in practical test-time settings, where ground-truth labels are unavailable and calibration may need to adapt to newly encountered target tasks. To address this challenge, we propose Test-Time Calibration Learning (TTCL), a label-free framework that jointly adapts reasoning accuracy and verbalized confidence directly on unlabeled target-task data. Specifically, TTCL derives self-supervision signals for both correctness and calibration from multiple model-generated responses, enabling calibration learning at test time without ground-truth labels. Theoretical analysis further establishes TTCL as a bounded surrogate for the ideal calibration objective. Extensive experiments on mathematical reasoning and factual question answering demonstrate that TTCL consistently improves both accuracy and calibration across diverse models and tasks. On base models, TTCL achieves an average relative accuracy improvement of +40.13% and an ECE reduction of +70.80% across eight benchmarks. Moreover, TTCL can further improve both accuracy and calibration for already calibrated models under domain shift, particularly when source-domain calibration transfers poorly to target tasks. In the math-to-factQA setting, TTCL achieves an average relative accuracy gain of +20.35% and reduces ECE by +53.83%. The source code is released at https://github.com/tmlr-group/TTCL.
Zizhuo Zhang, Xiong Peng, Jingwei Sun +4
TMLR Group, Department of Computer Science, Hong Kong Baptist University · Huawei Noah’s Ark Lab.
Scaling test-time computation with reinforcement learning (RL) has emerged as a reliable path to improve large language models (LLM) reasoning ability. Yet, outcome-based reward often incentivizes models to be overconfident, leading to hallucinations, unreliable confidence-based control, and unnecessary compute allocation. We introduce Reinforcement Learning with Confidence Margin (RLCM), a calibration-aware RL framework that jointly optimizes correctness and confidence reliability via a margin-enhanced process reward over intermediate-budget completions. Rather than aligning confidence to correctness likelihoods, RLCM encourages to widen the confidence margin between correct and incorrect steps within a single reasoning trajectory. Across mathematical, code, logic and science benchmarks, our method substantially improves calibration while maintaining or improving accuracy. We further show that, with calibrated confidence signals, the resulting models enable more efficient conformal risk control and effective confidence-weighted aggregation.