Back to AI Research

AI Research

Beyond Confidence: Stability-Aware Test-Time Adapta... | AI Research

Key Takeaways

  • Beyond Confidence: Stability-Aware Test-Time Adaptation for LLM Reasoning Large Language Models (LLMs) are powerful, but they often struggle with complex rea...
  • 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.
Paper AbstractExpand

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.

Beyond Confidence: Stability-Aware Test-Time Adaptation for LLM Reasoning
Large Language Models (LLMs) are powerful, but they often struggle with complex reasoning tasks. While researchers have developed ways to adapt these models at test time—without the high cost of retraining—many of these methods rely solely on "confidence." They assume that if a model is highly confident in its output, it is likely correct. However, this paper demonstrates that LLMs can be dangerously overconfident even when their reasoning is flawed. The authors introduce a new framework called TASCO, which improves reasoning by ensuring that a model’s confidence is not just high, but also stable when faced with minor perturbations.

The Problem with Confidence Alone

Current test-time adaptation methods often use "predictive entropy" to guide a model toward higher confidence. The authors observe that this approach is limited because it cannot distinguish between genuine, reliable reasoning and "spurious" confidence—where a model is simply stuck in a fragile, incorrect state. By analyzing various models, the researchers found that high-confidence reasoning is significantly more likely to be correct if that confidence remains steady when the input is slightly altered. If a model’s confidence fluctuates wildly under local changes, it is often a sign of unreliable reasoning. The same large language models question is explored in Efficient Test-Time Adaptation through Human-AI Interaction, which adds a research perspective.

How TASCO Works

TASCO (Test-Time Adaptation via Stability-Aware Confidence Optimization) improves reasoning by optimizing a small, shared "prefix"—a set of continuous variables prepended to the model's input—while keeping the LLM’s main parameters frozen. To ensure this prefix leads to stable, reliable reasoning, the framework uses two distinct strategies:

  • Random Perturbation: This method tests how the model behaves when the prefix is slightly modified with random noise. By penalizing large variations in confidence across these different versions, the model is encouraged to find a "flatter," more stable reasoning path.

  • Sharpness-Aware Perturbation: Instead of sampling random noise, this approach identifies the "worst-case" direction—the specific change that would most negatively impact the model's confidence. By optimizing the prefix to remain stable even against this worst-case sensitivity, the model becomes more robust. The same large language models question is explored in Rethinking On-Policy Distillation of Large Language..., which adds a research perspective.

Performance and Efficiency

The researchers tested TASCO across a variety of mathematical and scientific benchmarks, including high-school and olympiad-level problems. The results show that TASCO consistently outperforms standard confidence-based methods. For example, on Qwen2.5-Math models, TASCO achieved accuracy gains of up to 17.2 percentage points compared to standard prompting. Beyond just getting more answers right, the framework also improved token efficiency, meaning the models were able to reach correct conclusions while generating fewer tokens.

Key Takeaways

The primary contribution of this work is the shift from purely confidence-driven optimization to stability-aware optimization. By treating local stability as a vital signal for reliability, TASCO provides a way to steer frozen LLMs toward better performance without requiring external verifiers or reward models. The authors note that this approach is effective across different model families and reasoning paradigms, suggesting that stability is a fundamental indicator of how well an LLM is actually "thinking" through a problem. The same ai evaluation question is explored in Beyond Aggregate Scores, which adds a research perspective. as detailed in the full paper on Arxiv

Comments (0)

No comments yet

Be the first to share your thoughts!