PaperlayerView exact revision on arXiv

Source-linked AI summary

Beyond Confidence: Stability-Aware Test-Time Adaptation for LLM Reasoning

Bincheng Gu, Min Gao, Zongwei Wang, Yibing Bai, Yulan He, Junliang Yu

arXiv:2609.11393v1cs.AI

TL;DR

Confidence-based test-time adaptation can favor incorrect reasoning when high confidence is locally fragile. TASCO addresses this by optimizing a shared prefix for confidence stability under random or worst-case perturbations while keeping the LLM frozen, improving accuracy and token efficiency across models and benchmarks.

  • Problem

    Entropy-based adaptation can steer LLMs toward highly confident but incorrect reasoning trajectories because confidence alone does not distinguish reliable from fragile reasoning.

  • Method

    TASCO optimizes a shared task-level prefix for confidence and local stability using Random Perturbation and Sharpness-Aware Perturbation, without updating the frozen LLM or using external supervision.

  • Results

    TASCO improves reasoning accuracy and generation efficiency across diverse LLMs and benchmarks, with gains of up to 17.2 percentage points and 28.1% fewer tokens.

  • Takeaways & Limitations

    Local confidence stability complements confidence magnitude as a signal of reasoning reliability, enabling more accurate and concise frozen-model adaptation.

  • Takeaways & Limitations

    The theoretical objectives characterize robustness under sampled and worst-case perturbations rather than correctness, whose relationship with stability is evaluated empirically.

Abstract

from arXiv · show

Test-time adaptation has emerged as a lightweight alternative to costly post-training for improving the reasoning capabilities of Large Language Models (LLMs) on downstream tasks. Predictive entropy provides a model-derived signal for such adaptation, guiding models toward higher-confidence reasoning states without external verifiers or reward models. However, higher confidence does not necessarily imply correctness, as LLMs may remain highly confident along incorrect reasoning trajectories. We observe that high-confidence reasoning is more likely to be correct when confidence remains stable under local perturbations. Based on this observation, we propose Test-Time Adaptation via Stability-Aware Confidence Optimization (TASCO), a framework that incorporates local stability into confidence-based test-time adaptation while keeping the LLM frozen. TASCO operationalizes local stability by optimizing a lightweight task-level prefix under two alternative perturbation strategies: Random Perturbation promotes distributional stability across trajectories induced by nearby perturbed prefixes, whereas Sharpness-Aware Perturbation targets worst-case local sensitivity. Experiments demonstrate that TASCO improves reasoning accuracy and token efficiency across diverse LLMs and reasoning benchmarks, while behavioral analyses show that it maintains stable confidence under local perturbations without prematurely concentrating the model's predictive distribution.

Introduction

Confidence-only adaptation can steer LLMs toward fragile, incorrect high-confidence reasoning, whereas local confidence stability better indicates reliable reasoning. TASCO optimizes a shared prefix under perturbations and improves accuracy and efficiency while keeping the model frozen.

  • Continuous-variable adaptation guides frozen LLMs toward task-specific reasoning states without updating model parameters.
  • Entropy minimization alone cannot distinguish reliable high confidence from spurious confidence along incorrect reasoning trajectories.
  • High confidence is more likely to support correct reasoning when it remains stable under local prefix perturbations.Low-variance queries achieve higher accuracy than high-variance queries, including among queries with comparable mean confidence.
  • TASCO optimizes a shared task-level prefix for frozen LLMs using Random Perturbation or Sharpness-Aware Perturbation.Random Perturbation measures distributional variation across independently decoded perturbed trajectories, while Sharpness-Aware Perturbation targets worst-case local sensitivity.
  • TASCO improves reasoning accuracy and efficiency across diverse LLMs and benchmarks while promoting stable confidence under local perturbations.

Related Work

Prior work improves LLM reasoning through prompting, aggregation, search, verification, and longer deliberation, while test-time adaptation uses inference-time signals. Recent approaches increasingly optimize continuous variables while freezing model parameters.

  • LLM reasoning has been enhanced through CoT prompting, trajectory aggregation, structured search, step-level verification, and longer inference-time deliberation.
  • Test-time adaptation methods use inference-time signals such as majority-vote pseudo-rewards or entropy, but parameter updates can add computation and risk catastrophic forgetting.

Method

TASCO optimizes a shared continuous prefix for confidence and local stability while keeping the LLM frozen. It combines entropy-based confidence optimization with Random Perturbation or Sharpness-Aware Perturbation to control sensitivity under nearby or worst-case prefix changes.

  • Problem Formulation: TASCO optimizes a shared continuous prefix using label-free signals while keeping the LLM parameters frozen.The prefix is prepended to each query embedding and optimized without external supervision or auxiliary models.
  • Confidence Objective: Predictive entropy supplies the confidence objective, but entropy minimization alone can produce locally fragile high-confidence states.The method therefore augments confidence optimization with a local-stability objective.
  • Random Perturbation: Random Perturbation estimates distributional stability by re-decoding trajectories under Gaussian-perturbed prefixes and penalizing their confidence variance.The regularizer averages behavioral sharpness across inputs, while the confidence term encourages low predictive entropy.
  • Sharpness-Aware Perturbation: Sharpness-Aware Perturbation uses one gradient-guided perturbation to target worst-case confidence sensitivity along a fixed trajectory.Its perturbation is obtained from a first-order approximation of the inner maximization, avoiding multiple trajectory generations.
  • Stability Objectives: The two strategies provide complementary stability controls: sampled rollout variation for Random Perturbation and worst-case directional sensitivity for Sharpness-Aware Perturbation.The resulting objectives characterize robustness to prefix perturbations rather than correctness itself.
  • Theoretical Analysis: Under smoothness assumptions, the sharpness-aware perturbation approximates the worst-case local loss increase with O(ρ^2) error, while ideal local increase is governed by largest positive curvature at stationary points.These theoretical results connect gradient-guided perturbations with local sensitivity and curvature.

Experiments

Across general and reasoning-enhanced LLMs, TASCO improves accuracy and token efficiency while its ablations and behavioral analyses support stability-aware prefix optimization as the source of gains.

  • Performance Improvement: TASCO raises average accuracy across general and reasoning-enhanced LLMs, outperforming Zero-Shot CoT and strong comparison baselines.On Qwen2.5-Math-1.5B/7B, gains over Zero-Shot CoT are 14.8/17.2 points; on DeepSeek-R1-Distill-Qwen-1.5B/7B, gains are 6.8/6.0 points.
  • Token Efficiency: 28.1% and 24.0% reductions in average generation length show that TASCO improves accuracy without requiring longer reasoning traces.These reductions are reported for TASCO-R and TASCO-S across general LLMs compared with CoT.
  • Ablation of Prefix Optimization: TASCO-RP outperforms controlled variants, indicating that its gains require combining confidence optimization with local-stability guidance.The controls remove extra prefix positions, stability, or confidence objectives to isolate the source of improvement.
  • Behavioral Effects of Stability-Aware Optimization: TASCO reduces confidence variance and improves answer consistency under prefix perturbations while avoiding premature predictive concentration.On MATH500, TASCO retains higher early decoding entropy and lower top-1 probability before the confidence gap narrows later.
  • Analysis of Cross-Distribution Generalization: Transferred TASCO prefixes consistently outperform Zero-Shot CoT across evaluated source-target pairs, supporting cross-distribution generalization.The transfer experiment applies prefixes to target benchmarks without re-optimization.
  • Performance Across Query Difficulty: TASCO improves hard and medium queries across both tested models while maintaining comparable performance on easy queries.Sensitivity analyses also show both variants outperform confidence-only optimization across broad nonzero perturbation scales, with degradation only at excessively large scales.

Conclusion

The paper concludes that high-confidence reasoning is more reliable when stable under local perturbations, motivating TASCO’s frozen-model optimization of shared task-level prefixes.

  • Conclusion: TASCO combines confidence and local stability in a shared task-level prefix while keeping the LLM frozen.Random Perturbation captures distributional variation, whereas Sharpness-Aware Perturbation targets worst-case local sensitivity.
  • Conclusion: Across diverse LLMs and reasoning benchmarks, stability-aware confidence optimization improves reasoning accuracy and generation efficiency without parameter updates or external supervision.

A. Experimental Details

The experiments optimize only a continuous prefix with frozen model parameters, compare TASCO’s perturbation variants with multiple baselines, and evaluate accuracy under benchmark-specific protocols.

  • Optimization and Generation Settings: TASCO optimizes a continuous prefix using AdamW while keeping all LLM parameters frozen, then evaluates with greedy decoding.The confidence objective uses at most a configured number of initial generated tokens.
  • Perturbation Settings: Random Perturbation uses Gaussian prefix perturbations, while Sharpness-Aware Perturbation uses a fixed perturbation radius; the variants are optimized and evaluated separately.The reported settings are K = 8 and σ = 0.01 for Random Perturbation, and ρ = 0.8 for Sharpness-Aware Perturbation.
  • Baselines: Baselines include standard CoT, representation- and latent-variable adaptation methods, confidence-based methods, and reasoning-control approaches.The compared methods include SLOT, LatentSeek, LTPO, TTSV, s1, Chain of Draft, and α1.

B. Exploring Local Confidence Stability

The analysis measures confidence variance across perturbed-prefix trajectories and finds that lower local variance is associated with higher independently evaluated answer correctness, including after confidence controls.

  • Analysis Design: The analysis tests whether stability remains associated with correctness after accounting for confidence magnitude and query difficulty.Queries are split by model-specific median variance, and additional analysis uses five-way mean-confidence stratification and partial correlation.
  • Stability Measure: Local confidence stability is quantified as the sample variance of length-normalized trajectory confidence across K = 8 Gaussian perturbations.Lower S_i indicates greater local stability, and mean confidence is computed across perturbed trajectories.
  • Experimental Setup: The five-model setup uses confidence-only optimized prefixes, perturbs each query prefix with σ = 0.01, and evaluates correctness from an independently decoded reference trajectory.
  • Association with Correctness: Low-variance queries achieve 17.2–29.0 percentage points higher unperturbed accuracy across all five models, with AUROCs of 0.668–0.745.Stability is estimated from perturbed trajectories while correctness is evaluated independently on the unperturbed trajectory.

Controlling for Confidence Magnitude.

After controlling for mean confidence, local confidence variance generally remains associated with answer reliability, although this relationship is model-dependent.

  • Controlling for Confidence Magnitude.: 11.6 to 30.5 percentage points: Low-Variance queries achieve higher perturbation accuracy after confidence stratification on four of five models.The differences are statistically significant on four models, with partial correlations showing the same pattern after controlling for mean confidence.
  • Controlling for Confidence Magnitude.: Qwen2.5-Math-1.5B is an exception, with a nonsignificant 5.4-point gap and near-zero partial correlation.For this model, local variance is strongly correlated with mean confidence, suggesting little additional information beyond confidence magnitude.
  • Controlling for Confidence Magnitude.: The low-variance group remains more accurate across model families even when queries have comparable mean confidence and difficulty.These analyses indicate that the association is not explained solely by confidence magnitude or easier queries.

C. Proofs and Additional Theoretical Analysis

The theoretical analysis characterizes how Random Perturbation and Sharpness-Aware Perturbation control complementary forms of local confidence sensitivity. These results establish confidence robustness under prefix perturbations, not answer correctness.

  • C. Proofs and Additional Theoretical Analysis.: Sσ(V ; X) decomposes into within-trajectory confidence variation and between-trajectory variation across decoded behaviors.The decomposition separates smooth local changes within a fixed trajectory from variation caused by trajectory switching.
  • C. Proofs and Additional Theoretical Analysis.: Reducing population behavioral sharpness controls the probability of large confidence deviations under the perturbation distribution.This follows from the behavioral-sharpness analysis and is a stability statement rather than a correctness guarantee.
  • C. Proofs and Additional Theoretical Analysis.: The behavioral-change bound is informative when changed and unchanged trajectories have different mean confidence, but not when their means coincide.Thus, the bound does not quantify behavior-change frequency independently of confidence separation.
  • C. Proofs and Additional Theoretical Analysis.: The smooth local expansion is gradient-dominated generally and becomes quadratic-Hessian-dominated at stationary points.The derivation uses local smoothness, Gaussian perturbations, and a fixed decoding neighborhood.
  • C. Proofs and Additional Theoretical Analysis.: Random Perturbation measures distributional variation across perturbed behaviors, while Sharpness-Aware Perturbation targets worst-case sensitivity along one adverse local direction.The two objectives therefore cover sampled neighborhood variation and gradient-guided local sensitivity, respectively.
  • C. Proofs and Additional Theoretical Analysis.: Random Perturbation requires K independently decoded rollouts, whereas SAP uses one gradient-guided perturbation and substantially reduces autoregressive decoding cost.The learned prefix also transfers to unseen datasets without re-optimization, avoiding additional cross-distribution adaptation cost.

E. Further Analysis of Reasoning Effectiveness

TASCO improves reasoning-enhanced models across five mathematical benchmarks while balancing accuracy and generation length. The evaluation also measures variation across independent random seeds.

  • E. Further Analysis of Reasoning Effectiveness: 49.0% from 42.2% and 61.1% from 56.1%: TASCO raises average accuracy on the 1.5B and 7B reasoning-enhanced models.It simultaneously reduces generation length by 14.1% and 8.8%, respectively.
  • E. Further Analysis of Reasoning Effectiveness: TASCO achieves the highest average accuracy at both model scales and consistently outperforms the base model across all five benchmarks.Some reasoning-control baselines generate shorter outputs, but their accuracy gains are smaller or inconsistent.
  • E. Further Analysis of Reasoning Effectiveness: The study evaluates TASCO and representative baselines over three independent random seeds, reporting mean accuracy and standard deviation.The protocol covers no adaptation, latent-variable optimization, and confidence-only prefix optimization.

Semantic Evaluation of Reasoning Processes

The semantic evaluation examines whether TASCO changes the reasoning process, using an external judge to annotate paired trajectories. TASCO improves correctness and informativeness while reducing redundancy and premature commitment, and its performance is broadly stable across perturbation scales.

  • Semantic Evaluation of Reasoning Processes: An external judge labels semantic steps for correctness, informativeness, functional role, and premature commitment using paired generated trajectories.The judge receives the original problem and reasoning process without the ground-truth answer.
  • Semantic Evaluation of Reasoning Processes: Step Correctness measures the percentage of assessable steps labeled correct, excluding uncertain steps.Higher SC indicates a larger proportion of mathematically or logically valid steps.
  • Semantic Evaluation of Reasoning Processes: Higher Step Informativeness indicates more productive reasoning, whereas lower Redundancy Rate indicates less repetitive reasoning.Premature Commitment is also reported, with lower values indicating less commitment before sufficient support.
  • Semantic Evaluation of Reasoning Processes: Across four additional benchmarks, both TASCO variants outperform confidence-only optimization across a broad range of perturbation scales.The default settings σ = 0.01 and ρ = 0.8 achieve the best or near-best accuracy across datasets.
  • Semantic Evaluation of Reasoning Processes: TASCO improves all four reasoning-process metrics across both evaluated datasets.It produces more correct and informative steps while reducing redundant reasoning and premature commitment.

F. Case Studies

The case studies compare Confidence Only with TASCO across worked reasoning examples, showing that TASCO preserves conditions that confidence-only reasoning drops and reaches correct answers.

  • Case Studies: The examples are presented as case studies comparing the reasoning behaviors induced by Confidence Only and TASCO.They illustrate the experimental prompting protocol alongside the behavioral comparison.
  • Prompt and Example: The prompting protocol uses step-by-step reasoning with boxed final answers, while TASCO adds a learned continuous prefix only at the embedding level.The prefix is not introduced as textual instructions.
  • Reasoning Patterns: TASCO’s domain-constraint reasoning correctly requires both nonnegative quadratic terms to vanish simultaneously, rather than combining their roots by union.The two conditions therefore use an intersection of root sets.
  • Reasoning Patterns: Across three case studies, Confidence Only drops a necessary condition, whereas TASCO preserves it and obtains the correct answer.The comparison covers polynomial roots, domain constraints, and coprime factorization.
Loading 2609.11393v1…