Knowing When to Quit: A Principled Framework for Dynamic Abstention in LLM Reasoning

1. Problem Statement

Large Language Models (LLMs) utilizing Chain-of-Thought (CoT) reasoning often generate long, verbose responses. While this improves performance, it incurs significant computational costs. Crucially, incorrect reasoning traces tend to be longer than correct ones, meaning models waste compute generating incorrect answers that are eventually delivered to users.

Existing solutions for "abstention" (withholding low-quality outputs) suffer from timing limitations:

• Input Processing: Decides to abstain before generation based solely on the prompt. This is efficient but often inaccurate because the prompt alone provides little signal about the specific reasoning path the model will take.
• Output Processing: Decides to abstain after generation is complete. This wastes all compute resources.
• Fixed-Position Mid-Generation: Decides to abstain at a pre-determined token count (e.g., after 20 tokens). This is suboptimal because the point at which a reasoning trace becomes "unrecoverable" varies significantly across different problems.

The Gap: There is a lack of a principled, dynamic framework that allows the model to decide to stop generation at any token position based on the current state of the reasoning trace, balancing computational savings against the risk of discarding a potentially correct answer.

2. Methodology

The authors propose a formal framework based on KL-regularized Reinforcement Learning (RL) to model dynamic abstention.

A. MDP Formulation with Abstention

The text generation process is modeled as a Markov Decision Process (MDP):

• State: The prompt $x$ and the current token prefix $y_{1:t}$.
• Action Space: Augmented to include a special abstention token ($\perp$) alongside the vocabulary $V$.
• Reward Structure:
  • Terminal Reward: Sparse and binary. If the final answer is correct, reward = 1; otherwise, 0.
  • Abstention Reward ($r_\perp$): A fixed scalar reward ($0 < r_\perp \le 1$) given immediately if the model chooses to abstain. This represents the utility of a fallback mechanism (e.g., deferring to a human or a stronger model).
  • KL Regularization: Penalizes deviation from a reference policy ($\pi_{ref}$), though the theoretical analysis often focuses on the $\beta=0$ case (pure correctness).

B. The Dynamic Value-Thresholding Rule

The core contribution is a decision rule derived from the Value Function $V_\beta(x, y_{1:t})$, which represents the expected future reward from the current state.

• The Rule: At every token position $t$, the model calculates the value function. If $V_\beta(x, y_{1:t}) < r_\perp$, the model abstains immediately. Otherwise, it continues generating.
• Interpretation: The value function acts as the model's "confidence" in eventual success. If the expected reward of continuing drops below the guaranteed reward of the fallback ($r_\perp$), continuing is mathematically wasteful.
• Control: The parameter $r_\perp$ allows practitioners to directly tune the trade-off between selective accuracy (lower $r_\perp$) and compute savings (higher $r_\perp$).

C. Value Function Estimation

Since the true value function is unknown, the authors propose an efficient estimation method:

• Target: For binary rewards and $\beta=0$, the value function equals the conditional probability of correctness: $P(\text{correct} | x, y_{1:t})$.
• Architecture: A lightweight two-layer MLP probe trained on the hidden states of the LLM.
• Training: Uses Binary Cross-Entropy (BCE) loss. The probe is trained in parallel on all non-terminal token positions of completed trajectories. Theoretical proofs show that under realizability assumptions, minimizing BCE recovers the true value function.

3. Key Theoretical Contributions

The paper provides rigorous proofs establishing the superiority of dynamic value-thresholding:

1. Dominance over No Abstention: The dynamic policy strictly dominates any non-abstaining policy if there is a positive probability of abstention. It never performs worse and strictly improves when it stops unpromising traces.
2. Dominance over Fixed-Position Abstention: Dynamic value-thresholding dominates methods that decide to abstain at a fixed token $k$. It strictly outperforms them when low-value states cannot recover to high-value states by token $k$ (a common scenario in reasoning errors).
3. Optimality: Under ideal conditions (zero KL regularization and an optimal base policy), the dynamic value-thresholding rule is globally optimal among all abstention-enabled policies.
4. Improvement Bound: The magnitude of improvement is bounded linearly by the abstention rate and the Lipschitz constant of the value function, providing a theoretical guarantee on the gains.

4. Experimental Results

The framework was evaluated on Qwen2.5-7B and Phi-3-mini across two mathematical reasoning benchmarks: GSM8K (grade-school math) and OlympiadBench (competition-level math).

• Selective Accuracy:

  • Dynamic abstention achieved the highest selective accuracy across all abstention rates compared to input-processing (prompt-only) and fixed-position baselines.
  • OlympiadBench Results: At a 90% abstention rate, dynamic abstention reached 64% accuracy on Phi-3 (vs. 34% for baselines) and 91% accuracy on Qwen (vs. 75% for baselines).
  • The advantage is most pronounced on harder problems where reasoning traces are longer and more prone to early errors.

• Efficiency:

  • The method retains 63–86% of the token savings of input-only methods at low abstention rates, rising to 92–95% at high abstention rates.
  • Abstention typically occurs in the first half of the generation, even at low rates, validating the efficiency of early termination.

• Generalization & Robustness:

  • Cross-Dataset Transfer: A probe trained on GSM8K generalized well to OlympiadBench without retraining, suggesting the hidden states encode general reasoning quality.
  • Robustness: The method is invariant to monotone transformations of the value estimator (ranking suffices) and degrades gracefully under noise.
  • Toxicity Avoidance: The framework also successfully applied to toxicity avoidance on RealToxicityPrompts, achieving perfect non-toxic rates at moderate abstention levels.

5. Significance and Impact

• Principled Decision Making: Moves beyond heuristic or empirical rules for stopping generation, providing a mathematically grounded framework for dynamic compute allocation.
• Efficiency at Scale: By terminating unpromising traces early, this method significantly reduces the energy consumption and environmental footprint of LLM inference.
• Reliability: Enables models to "know when they don't know," deferring difficult or uncertain tasks to human experts or stronger models, which is critical for high-stakes domains like medicine, law, and finance.
• Practical Deployment: The method adds negligible inference overhead (a single MLP forward pass per token) and requires only a small held-out set for calibration, making it highly deployable.

In summary, the paper establishes that dynamic value-thresholding is the optimal strategy for balancing the cost of computation with the quality of LLM outputs, outperforming all existing static or fixed-position abstention methods.