Stability aware adaptation improves large language model reasoning accuracy
Beyond Confidence: Stability-Aware Test-Time Adaptation for LLM Reasoning
Artificial Intelligence
Summary
Large language models sometimes get very confident about answers that are actually wrong. The authors found that when a model's confidence stays steady even after small changes to the input, the answer is more likely to be right. They created a method called TASCO that adjusts the model at test time to focus on this stability, without changing the main model itself. This improves how well models reason and saves computing effort.
What this means in practice
- •For natural language processing engineers: Enhance reasoning accuracy of language models during inference without costly retraining by using stability-aware confidence adjustments.
- •For software developers integrating llms: Improve output reliability of LLM-based applications by applying lightweight task-specific prefix tuning adapting model confidence and stability at test time.
Authors
Bincheng Gu, Min Gao, Zongwei Wang, Yibing Bai, Yulan He, Junliang Yu
Abstract
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.