- The paper introduces a RL-based dynamic abstention mechanism that halts low-value reasoning paths to save compute and improve accuracy.
- It employs a KL-regularized reinforcement learning framework with a lightweight MLP probe to estimate model confidence at each token position.
- Empirical evaluations reveal significant selective accuracy gains on math benchmarks and robust toxicity avoidance, ensuring practical deployment.
Knowing When to Quit: A Principled Framework for Dynamic Abstention in LLM Reasoning
Introduction and Problem Setting
LLMs utilizing Chain-of-Thought (CoT) prompting have established competence in complex reasoning domains such as mathematical problem solving, symbolic logic, and program synthesis. Despite their efficacy, these models frequently generate verbose yet incorrect responses, multiplying computational expenditure with no accuracy guarantee. Empirical findings further suggest that incorrect output traces tend to be longer on average, compounding inefficiency. This situation indicates a dual deficiency: propagation of inaccurate outputs and unnecessary resource consumption.
Abstention, deferring or withholding uncertain outputs, is a natural mitigation. Existing approaches operate either before generation ("input processing," deciding based on the prompt) or after completion ("output processing," using classifier or verifier signals post hoc), but both have significant limitations. Input-only methods lack path-specific knowledge necessary for selective early stopping; output-based abstention comes too late to yield computational savings. Dynamic, mid-generation abstention—terminating unpromising reasoning traces at arbitrary positions—promises improved efficiency and reliability. However, previous proposals lack theoretically sound criteria for executing this decision.
Framework and Methodology
This paper formalizes dynamic abstention in LLMs by embedding the abstention decision within a KL-regularized reinforcement learning (RL) framework, modeling text generation as an MDP with sparse, binary end-of-sequence rewards signifying correctness or failure. The action space is explicitly augmented with an abstention token, L, triggering an immediate halt and a fallback or deferral action with associated scalar reward rabstain​∈(0,1]. This reward quantifies the operational utility of abstention (e.g., handing off to a human) and explicitly tunes the abstention-accuracy-compute efficiency trade-off.
Crucially, the proposed policy, dynamic value-thresholding, abstains whenever the value function (expected future cumulative reward conditioned on current state) falls below rabstain​. This interprets the value function as model confidence in success. The framework fundamentally generalizes sequence generation with reject options, reducing to standard CoT when rabstain​=0. Selective abstention can thus occur at any token position, adapting the stopping point for each trace.
Value function estimation is nontrivial in the context of sparse rewards. The paper demonstrates that, in the setting of binary terminal rewards and zero KL regularization, the value function at any non-terminal prefix equals P(success∣current state). Practically, this conditional probability is efficiently recovered by training a two-layer MLP probe with binary cross-entropy loss on the model's hidden state at each token position, a choice incurring negligible computational overhead.
Theoretical Results
The essential analytical results are as follows:
- Dominance over Baselines: The dynamic value-thresholding policy strictly dominates any non-abstaining policy and all fixed-position abstention baselines (input-based or mid-generation), assuming perfect value estimates. This dominance persists except in degenerate situations where the abstention action is never chosen.
- Optimality: When the base LLM is itself optimal (i.e., achieves the maximum possible expected reward) and KL regularization is zero, the dynamic value-thresholding rule is globally optimal among all policies augmented with an abstention action.
- Linear Improvement Bound: The improvement in expected reward over the base policy is linearly upper bounded in the abstention rate, controlled by the Lipschitz constant of the value function.
- Estimation Guarantees: For binary terminal rewards, minimizing binary cross-entropy at each non-terminal hidden state recovers the value function under realizability assumptions, guaranteeing the correctness of the dynamic abstention criterion.
Empirical Evaluation
Empirical verification uses two competitive LLMs (Qwen2.5-7B-Instruct and Phi-3-mini-4k-Instruct) on both easy (GSM8K) and hard (OlympiadBench) math reasoning benchmarks. Baselines include input-only probes, fixed-position mid-generation probes, self-assessment logits, and LoRA-based abstention heads. The dynamic value-thresholding policy is operationalized using an MLP probe at every token position.
Strong empirical results are reported:
- On OlympiadBench (Phi-3, 16% baseline accuracy), dynamic abstention attains 64% selective accuracy at 90% abstention, compared to only 34% for the strongest baseline. For Qwen (43% baseline), selective accuracy reaches 91% at 90% abstention, compared to 75% for the best baseline.
- On GSM8K, where baselines are already strong, the proposed method remains consistently dominant, achieving selective accuracy up to 99% at 90% abstention.
- Efficiency Trade-offs: Token savings of dynamic abstention approach those of input-only baselines at high abstention rates (92–95% parity at 90% abstention). At lower rates, some savings are traded for increased accuracy; on average, abstention occurs within the first half of generation.
- Robustness: Selective accuracy is invariant to monotonic recalibration of the value estimator; performance degrades gracefully under Gaussian noise corrupting the value function's ranking.
- Cross-Dataset Generalization: The trained value estimator generalizes well between benchmarks, indicating the learned signals correspond to general properties of reasoning quality, not dataset idiosyncrasies.
- Toxicity Avoidance: On RealToxicityPrompts, dynamic abstention achieves 100% non-toxic response rate among non-abstained samples at moderate abstention levels, outperforming all baselines.
Practical and Theoretical Implications
Practically, dynamic abstention enables both reliability improvements and substantial computational savings in real-world CoT settings. By terminating unlikely-to-succeed traces early, production deployments can reduce energy and hardware costs, mitigate risk of confidently erroneous outputs in high-stakes scenarios (e.g., medical, legal, or financial applications), and foster user trust in the model's capability to defer appropriately. The method's overhead is small—one lightweight probe application per position, and the calibration threshold is highly transferable between datasets.
On the theoretical side, the framework generalizes abstention-enabled selective classification to the sequential, autoregressive domain, delivering formal guarantees absent from earlier empirical approaches. The abstention reward rabstain​ is semantically meaningful and directly interpretable. The framework highlights that partial trajectories (not just prompts) encode crucial information about eventual success, consistent with recent findings on LLM self-knowledge. It additionally motivates further research in state-dependent abstention rewards and the extension to continuous preference-model-based rewards.
Future Directions
Two promising research efforts are indicated:
- State-Dependent Abstention Rewards: Adapting the abstention reward to depend upon the input and intermediate trace could align quitting strategies better with pragmatic utility in deployment (e.g., deferring only after partial progress has produced promising partial results).
- Extension to Continuous Reward Tasks: While this work focuses on binary correctness, extension to continuous reward signals (from learned preference models) is straightforward, with the theoretical results generalizing via MSE-based value estimation.
Conclusion
This paper establishes a rigorous and efficient methodology for mid-generation abstention in LLMs, allowing dynamic, per-trace early termination of unpromising reasoning trajectories. The dynamic value-thresholding criterion is grounded in RL value functions, dominates prior abstention strategies both theoretically and empirically, and is validated on challenging benchmarks as well as safety-critical tasks. It offers a principled path forward for cost-effective, reliable deployment of LLMs in settings where compute efficiency and output correctness are both critical (2604.18419).