Source-linked AI summary
STEP: Soft Trajectory Estimation of Prediction Correctness via Per-Step Probability Scoring
Chenyang Yu, Zifan Zeng, Qunli Zhang, Feng Liu
TL;DR
LLMs can give incorrect answers with high apparent confidence, while output-based correctness estimates do not capture intermediate reasoning. STEP uses lightweight forward passes to score per-step answer distributions, and its soft trajectory features consistently outperform existing baselines across benchmarks, particularly in AUROC.
Problem
Output-based uncertainty measures use final token statistics and do not capture whether intermediate reasoning is semantically coherent, complicating correctness estimation for high-stakes use.
Source: §1 · p. 1Method
STEP re-scores each prefix of an already-generated chain of thought with a lightweight forward pass, deriving soft trajectory features from answer probabilities in closed-set tasks.
Source: §1 · p. 2Results
Takeaways & Limitations
Preserving continuous probabilities across reasoning steps provides richer correctness information than hard binary indicators.
Source: §1 · p. 2, §5 · p. 9STEP is designed for closed-set tasks with disambiguable first-token candidate answers and relies on structured chain-of-thought, so it may not benefit consistently on unstructured reasoning or exact-computation tasks.
Source: Limitations · p. 9
Abstract
from arXiv · showhide
Estimating the correctness of large language model (LLM) predictions without access to ground truth is critical for reliable deployment in high-stakes settings. Existing approaches, such as entropy and perplexity, operate solely on final output statistics, discarding the rich information encoded in intermediate reasoning steps. To address this limitation, we propose STEP (Soft Trajectory Estimation of Prediction correctness), a training-free framework that scores each reasoning step via lightweight forward passes to obtain a soft probability distribution qi over candidate answers, from which we derive a set of soft trajectory features capturing per-step confidence and its progression toward the final answer. Extensive experiments on various benchmarks demonstrate consistent and substantial improvements, particularly in AUROC, over output-based baselines and hard-indicator trajectory methods. Notably, Temporally-Weighted Consistency, which assigns higher importance to later reasoning steps where answer distributions have stabilized, contributes the largest gains, highlighting the importance of temporal structure in reasoning trajectories. These results further validate the advantage of preserving continuous probability signals over binary discretization.
1 Introduction
LLMs can produce incorrect answers with high apparent confidence, while output statistics do not capture semantic reasoning. STEP estimates correctness in closed-set tasks from soft, per-step answer distributions and shows consistent improvements across benchmarks and models.
1 Introduction
- Incorrect high-confidence answers make inference-time correctness estimation important for reliable deployment, especially in high-stakes settings.Examples include medical question answering, legal reasoning, and automated decision support.Source: §1 · p. 1
- Output entropy, perplexity, and answer probability capture token-level fluency or sharpness, not whether reasoning is semantically coherent or grounded.A plausible but wrong answer can receive high token probability, while a correct answer reached amid uncertainty can receive modest probability.Source: §1 · p. 1
- STEP rescales each prefix of an already-generated reasoning trajectory with a lightweight forward pass, producing soft probabilities over candidate answers without additional generation or resampling.The framework is training-free and targets closed-set tasks with predefined answer candidates.Source: §1 · p. 2
- Its three soft trajectory metrics preserve continuous confidence signals that prior binary-indicator methods discard.The per-step distributions capture which answer is favored and how confidently and consistently the model favors it across the trajectory.Source: §1 · p. 2
- Across four benchmarks and two open-source models, soft trajectory features perform strongly and consistently against existing baselines, with later reasoning steps carrying disproportionately more predictive signal.The experiments report particularly strong AUROC performance.Source: §1 · p. 2
2 Related Work
Prior correctness-estimation work draws on output statistics, sampled responses, hidden-state probes, and reasoning trajectories. STEP is positioned against the costs, fitting requirements, and information losses of these approaches.
2 Related Work
- Output-statistics methods use entropy, length-normalized log-probability, or maximum softmax probability, but measure token-level fluency and sharpness rather than semantic correctness.A superficially fluent wrong answer can receive high token probability, while a correct but uncertain answer can receive low probability.Source: §2 · p. 2
- Semantic entropy clusters sampled responses by meaning, but semantic-entropy and hallucination-detection approaches require multiple samples per prompt, adding inference overhead.The cited work includes Kuhn et al. (2023) and Farquhar et al. (2024).Source: §2 · p. 2
- Hidden-state probes and latent-trajectory methods use internal representations to estimate correctness, but require distribution-specific fitting or white-box access to layer activations.The fitting requirement can limit applicability when the task domain or model changes.Source: §2 · p. 3
- CoVo uses output-probability trajectories to assess consistency and volatility, but samples independent trajectories and reduces confidence to hard binary indicators.STEP instead operates on one trajectory and retains soft probability distributions.Source: §2 · pp. 3–4
- Trajectory-based estimation assumes intermediate steps meaningfully reflect reasoning; STEP probes evolving answer preferences and differs from process reward models that require annotated training data.STEP derives correctness signals from output distributions at each step without training.Source: §2 · p. 4
3 Method
STEP estimates correctness in closed-set tasks by probing a model’s answer preferences at each reasoning step, then summarizing their confidence and progression. Its soft features preserve probability magnitudes and temporal structure that hard indicators discard.
3.1 Problem Formulation
- STEP scores a final answer against ground truth without labels, using a scalar correctness score for a question, candidate set, and generated chain-of-thought.The task is defined for a predefined set of candidate answers.Source: §3.1 · p. 4
3.2 Per-Step Answer Distribution
- The method segments an already-generated chain-of-thought into steps and probes each truncated prefix with one forward pass to obtain candidate-answer probabilities.This requires T lightweight forward passes and no additional text generation; candidates must be distinguishable by their first token.Source: §3.2 · p. 4
3.3 STEP: Soft Trajectory Estimation of Prediction Correctness
- Soft consistency averages confidence in the final predicted answer across steps, while margin-based consistency captures the gap between the two most probable answers.The features aggregate per-step distributions and are normalized by the number of steps; α controls the strictness of the margin feature.Source: §3.3 · pp. 4–5
- Temporally-Weighted Consistency gives later steps greater weight because their answer distributions are described as more discriminative between correct and incorrect trajectories.As β approaches zero, this feature recovers soft consistency.Source: §3.3 · p. 5
3.4 Relationship to Hard Indicators
- Unlike hard consistency, STEP’s soft features retain confidence magnitudes; margin and temporal features also capture decisiveness and trajectory structure that binary indicators cannot encode.The predefined candidate set permits single-trajectory probing without resampling, simplifying the prior cross-trajectory approach.Source: §3.4 · p. 5
4 Experiments
STEP is evaluated across four benchmarks and two models, with comparisons against output-based, latent-trajectory, and hard token-trajectory methods. Its soft trajectory features perform consistently well, and the experiments examine how confidence evolves across reasoning steps and how temporal weighting affects performance.
4.1 Experimental Setup
- Experiments use four reasoning benchmarks and two 7–8B instruction-tuned models, reusing one fixed chain-of-thought rollout per query across methods.Source: §4.1 · pp. 5–6
4.2 Main Results
- STEP’s soft trajectory features achieve the best overall performance across datasets and models, surpassing output-statistics, latent-trajectory, and hard token-trajectory baselines in most settings.Source: Table 1 · p. 6, §4.2 · pp. 7–8
4.1 Experimental Setup
- Task accuracy spans 61.13–89.89%, indicating meaningfully imbalanced rather than near-ceiling or near-chance evaluation.Source: Table 2 · p. 6, §4.1 · p. 6
- STEP adds T lightweight forward passes and no generated tokens; unlike N-sample methods, its overhead is O(T) forward passes.The number of intermediate steps, T̄, typically remains in the single digits.Source: §4.1 · pp. 6–7
4.3 Trajectory Analysis
- Correct trajectories maintain higher q* throughout and converge toward near-certainty, while incorrect trajectories plateau lower across datasets and models.Soft Consistency averages q* across the chain of thought.Source: Figure 2 · p. 7, §4.3 · p. 8
4.4 Ablation Study
- AUROC increases monotonically as β moves from negative to positive across datasets; β > 0 weights later steps, while β = 0 recovers Consoft.Source: Figure 3 · p. 8, §4.4 · p. 8
- At large β, AUROC gains plateau around β ≈5 to 10; negative β hurts across datasets, as early answer distributions remain noisy and less committed.Source: §4.4 · pp. 8–9
5 Conclusion
STEP estimates correctness in closed-set reasoning tasks by deriving soft trajectory features from per-step probability distributions using lightweight forward passes on an existing chain of thought. Across benchmarks and models, its variants outperform existing baselines in AUROC, with Temporally-Weighted Consistency strongest; extending the approach to open-ended settings is left for future work.
5 Conclusion
- Across multiple benchmarks and two models, all three STEP instantiations consistently outperform existing baselines in AUROC while achieving competitive or improved AUPR performance.Source: §5 · p. 9
- STEP derives soft trajectory features from per-step candidate-answer probabilities using lightweight forward passes on an already-generated chain of thought, at no additional generation cost.Source: §5 · p. 9
- Temporally-Weighted Consistency performs strongest, supporting the conclusion that later reasoning steps carry disproportionately more predictive signal than hard binary indicators preserve.Source: §5 · p. 9
- Future work may extend soft trajectory features to open-ended settings by dynamically constructing candidate sets from the model’s top-k predictions at each step.Source: §5 · p. 9
Limitations
STEP’s applicability is limited by its closed-set answer format, the structure of the reasoning trajectory, and access to next-token logits.
Scope of applicability
- STEP is designed for closed-set tasks and is not directly applicable to open-ended generation.Candidate answers must also be distinguishable by their first token; shared prefixes require additional aggregation strategies.Source: Limitations · p. 9
Dependence on reasoning structure
- STEP relies on structured chain-of-thought reasoning, and may not benefit consistently on unstructured or highly variable reasoning and exact-computation tasks.The concern is formal symbolic tasks where correctness depends on exact computation rather than progressive semantic narrowing.Source: Limitations · p. 9
Access requirements
- STEP requires access to next-token logits, which are available for open-source models but not black-box APIs.Source: Limitations · p. 9
A Answer-Elicitation Prompt
The answer-elicitation prompt is appended to each truncated reasoning context so a single forward pass can estimate the model’s candidate-answer preferences at that step. A separate generation template structures reasoning into newline-separated steps for scoring.
A Answer-Elicitation Prompt
- An answer-elicitation prompt is appended to each truncated context before a forward pass computes the per-step answer distribution qi.Source: A Answer-Elicitation Prompt · p. 10
- The model’s logits over candidate answer tokens are normalized with softmax over the candidate set to obtain qi.Only one forward pass is used, with no further tokens generated.Source: A Answer-Elicitation Prompt · p. 10
- The elicitation prompt targets the model’s implicit answer preference given its reasoning so far, without eliciting additional reasoning.Source: A Answer-Elicitation Prompt · p. 10
B Chain-of-Thought Generation Prompt
- The chain-of-thought template asks for a few statements, each representing a small logical improvement and appearing on its own line.The input-data placeholder is replaced with the dataset-specific question and choices.Source: B Chain-of-Thought Generation Prompt · p. 10
- The response must end with “Answer: X,” where X is one of the dataset’s candidate labels, and include nothing after that line.Source: B Chain-of-Thought Generation Prompt · p. 10
- Newline splitting treats each line except the final answer as one reasoning step sk, forming a trajectory of length T for confidence scoring.The prompt avoids numbered lists and structured headers to support this segmentation.Source: B Chain-of-Thought Generation Prompt · p. 11
C Evaluation Metrics
The evaluation uses AUROC and two precision-recall measures to assess correctness scoring and error detection. The measures differ in which prediction class is treated as positive.
AUROC
- AUROC measures whether correct predictions receive higher scores than incorrect ones, with 1.0 indicating perfect separation and 0.5 random performance.Source: C Evaluation Metrics · p. 11
AUPR-C
- AUPR-C summarizes precision-recall trade-offs when correct predictions are positive and complements AUROC under class imbalance.It reflects the ability to confirm high-confidence correct predictions.Source: C Evaluation Metrics · p. 11
AUPR-E
- AUPR-E measures error detection when incorrect predictions are positive, a minority class on benchmarks where accuracy ranges from 60% to 90%.The passage describes it as particularly informative for practical error-detection evaluation.Source: C Evaluation Metrics · p. 11
D Dataset Statistics
The evaluation covers four multiple-choice datasets spanning academic subjects, commonsense reasoning, reading comprehension, and science. Table 3 summarizes dataset statistics.
D Dataset Statistics
- The four evaluation datasets are MMLU, CommonsenseQA, Belebele, and ARC-Challenge, all used as multiple-choice question-answering benchmarks.Source: D Dataset Statistics · p. 11
- MMLU spans 57 academic subjects, while CommonsenseQA tests everyday commonsense knowledge among plausible answer choices.Source: D Dataset Statistics · p. 11
- Belebele evaluates reading comprehension using its English subset, and ARC-Challenge comprises science questions designed to resist simple retrieval and word co-occurrence methods.Source: D Dataset Statistics · p. 11
- Table 3 summarizes the key statistics for the evaluation datasets.Source: Table 3 · p. 11, D Dataset Statistics · p. 11
E Relationship to CoVo Consistency
In the closed-set setting, STEP adapts CoVo consistency by using the predefined candidate set and per-step answer distributions. This yields a single-trajectory binary measure, while STEP’s soft features preserve continuous confidence information.
E Relationship to CoVo Consistency
- Using the predefined candidate set and normalized per-step distribution q_i replaces CoVo’s sampled-answer distance comparison with a candidate-preference comparison.The original formulation uses N sampled trajectories to populate K distinct answers; the closed-set candidate set is fixed.Source: E Relationship to CoVo Consistency · pp. 11–12
- Dropping CoVo’s group-level summation for a single trajectory yields Conhard, which retains a binary consistency structure and needs one forward pass per reasoning step.Intermediate states are reindexed from i = 1 to T in this paper’s notation.Source: E Relationship to CoVo Consistency · p. 12
- STEP’s soft features replace Conhard’s binary indicator with continuous q*_i values, preserving per-step confidence information.CoVo consistency is defined as the fraction of intermediate states whose own final answer is closer than any alternative.Source: E Relationship to CoVo Consistency · pp. 11–12
F Runtime and Hardware
Experiments ran without model training on a single NVIDIA A100 GPU. Dataset inference, including generation and downstream scoring, took 20 minutes to approximately 3–4 hours, depending on dataset size and model.
F Runtime and Hardware
- 20 minutes to approximately 3–4 hours per dataset covered chain-of-thought generation and downstream scoring, depending on dataset size and model.Experiments used a single NVIDIA A100 GPU and involved no model training.Source: F Runtime and Hardware · p. 12
- STEP’s extra forward passes process fixed truncated prefixes without producing tokens, avoiding autoregressive generation of additional full reasoning trajectories.Source: F Runtime and Hardware · p. 12