When Does Learning to Stop Help? A Cost-Aware Study of Early Exits in Reasoning Models

2026-06-29

AI Summary

This work addresses the challenge of achieving optimal performance in reasoning models under minimal computational cost by proposing LearnStop, a lightweight, learning-based early-stopping mechanism that operates without access to hidden states. LearnStop dynamically predicts whether the current reasoning prefix is correct by online aggregation of multidimensional features—such as confidence, entropy, and answer stability—at fixed budget points to decide whether to terminate inference early. Theoretical analysis and experiments demonstrate that LearnStop significantly outperforms conventional threshold-based methods specifically in tasks lacking reliable single-scalar signals yet containing early correct answers, such as open-ended mathematical reasoning; on GSM8K, it improves accuracy by 2.8% over the strongest scalar baseline and extends the performance frontier under fixed budgets. However, its advantage is limited in multiple-choice or extremely difficult problems. The study systematically delineates the effectiveness boundary of learned stopping strategies.

Abstract

Reasoning models spend different amounts of useful computation across instances, but it remains unclear when a learned stopping rule improves over simple confidence or convergence thresholds. We study this question with LearnStop, a hidden-state-free checkpoint stopper for reasoning language models. At fixed budget checkpoints, LearnStop probes a short answer from the current reasoning prefix and predicts prefix correctness from online features such as answer confidence, entropy, prefix vote share, answer stability, and backtracking-marker density. Across 18 task-model settings spanning GSM8K, MATH-500, MMLU-Pro, AIME-90, GPQA, Qwen3, and DeepSeek-R1 distillations, the answer is task-dependent. On free-form math, learned multi-feature stopping improves the fixed-budget frontier and often beats scalar exits: on GSM8K with Qwen3-32B, the empirical frontier reaches a post-hoc peak adapt gain of +0.157, validation-selected operating points preserve positive gains, and the paired gain over the strongest scalar baseline is +0.028. On multiple-choice and very hard settings, scalar confidence, entropy, or stability rules are competitive or stronger. We therefore frame learned stopping not as a universal replacement for scalar exits, but as a tool whose value depends on trajectory structure. We further provide validation-selected operating points, paired bootstrap tests, finite-grid lost-correct risk calibration, cost accounting under KV-fork, prefix-cache, and black-box regimes, H100 serving profiles, checkpoint-schedule sweeps, transfer analyses, and robustness checks. The main practical finding is that learned stopping is useful when many questions become correct before full budget but do not exhibit a single reliable scalar stopping signal; its benefits largely disappear when confidence or answer convergence already solves the stopping problem.

Problem

• early exiting
• reasoning models
• stopping criteria
• cost-aware learning
• adaptive inference

Innovation

• learned stopping
• early exit
• reasoning models
• cost-aware inference
• multi-feature prediction