Why We Need a Thinking Budget
Have you ever watched an LLM, mid-stream in deep thought, spiral into an endless loop of self-doubt?
Chain of thought (CoT) does improve model accuracy to some degree, and reasoning mode has steadily become a must-have feature across major open-source models. But a new problem has followed close behind: runaway thinking. LLMs frequently exhibit overthinking — excessive redundant speculation, constant self-negation — all of it driving up the token budget. These issues are especially pronounced in smaller reasoning models obtained through distillation.
So a natural idea presents itself: give the model a thinking budget. Use a small budget for simple problems, and a larger one for hard problems.
Reasoning Control from Mainstream Model Providers
Across mainstream model providers, you’ll find that APIs and application services generally include mechanisms for controlling reasoning length.
OpenAI: provides parameters like reasoning.effort to hint at the model’s thinking depth (for o1 and o3 series models)
Anthropic: allows triggering different levels of thinking through budget_tokens or specific prompts (e.g., “think” < “think hard” < “think harder” < “ultrathink.”)
Gemini & Qwen: offer an explicit thinking_budget parameter directly in their APIs
Open-Source Implementation: Budget Forcing
So how should developers — especially those deploying LLMs locally — implement thinking-budget?
The most intuitive approach is to use a prompt, but the results are underwhelming. Unless the model was specifically trained for it during training, explicitly telling it to “think with fewer tokens” generally does little. In contrast, the Qwen3 series of hybrid thinking models incorporated /think and /no_think markers during training, so adding /no_think to the system prompt or user prompt at inference time can switch the model out of thinking mode.
(Section 4.3) “To better integrate the two modes and enable users to dynamically switch the model’s thinking process, we design chat templates for Qwen3, as shown in Table 9. Specifically, for samples in thinking mode and non-thinking mode, we introduce
/thinkand/no thinkflags in the user query or system message, respectively. This allows the model to follow the user’s input and select the appropriate thinking mode accordingly.” …(Section 4.4): “Format Following: In addition to explicit instructions, we expect the model to adhere to specific formatting conventions. For instance, it should respond appropriately to the
/thinkand/no thinkflags by switching between thinking and non-thinking modes…” … (Section 4.7): “ThinkFollow: Involves multi-turn dialogues with randomly inserted/thinkand/no thinkflags to test whether the model can correctly switch thinking modes based on user queries.”
While reading through Qwen3’s technical report and related documentation, I came across the approach Qwen officially provides the open-source community for thinking-budget.
It proposes a two-step implementation:
Step 1: Generate. This is no different from a standard generation — you simply set max_new_tokens to thinking_budget.
# first generation until thinking budget
generated_ids = model.generate(
**model_inputs,
max_new_tokens=thinking_budget
)
output_ids = generated_ids[0][input_length:].tolist()
Step 2: Check for </think>.
That is, check whether the token ID for </think> (151668) appears in the output.
- If it does: The model completed its thinking within budget — no intervention needed.
- If it doesn’t: The thinking process was cut off by the budget. Inject the
early_stopping_text, concatenate it with the previous output, and feed it as a new round of input for a secondmodel.generatecall.
# check if the thinking process has finished (151668 is </think>)
if 151668 not in output_ids:
print("thinking budget is reached")
# Prepare the early-stopping prompt
early_stopping_text = "\n\nConsidering the limited time by the user, I have to give the solution based on the thinking directly now.\n</think>\n\n"
early_stopping_ids = tokenizer([early_stopping_text], ...).input_ids.to(model.device)
# Concatenate the prompt to the end of the first generation
input_ids = torch.cat([generated_ids, early_stopping_ids], dim=-1)
# Run the second generation
generated_ids = model.generate(
input_ids=input_ids,
...
)
But this approach requires two rounds of generation. Is there a more elegant way to do it?
The answer is yes — by using a LogitsProcessor.
The idea comes from this blog post and a GitHub issue discussion.
The core principle: create a custom LogitsProcessor that inspects the state at every token generation step:
- Determine whether the model is currently between
<think>and</think> - Count the number of tokens generated since
<think>began - When the count reaches
thinking_budget, forcibly modify the probability distribution (logits) for the next token: set the probability of</think>to 1.0 and all other tokens to negative infinity
This implements thinking-budget through constrained decoding.
from typing import List, Optional
import torch
class ThinkLogitsProcessor:
"""A logits processor that limits the number of thinking tokens."""
def __init__(self, think_start_token, think_end_token, num_think_tokens: int = 100):
"""
Initialize the think logits processor.
Args:
tokenizer: The tokenizer used for the model
num_think_tokens: Maximum number of tokens allowed in thinking section
"""
self.num_think_tokens = num_think_tokens
self.think_start_token = think_start_token
self.think_end_token = think_end_token
def __call__(
self,
input_ids: List[int],
logits: torch.Tensor,
) -> torch.Tensor:
"""
Process the logits to enforce </think> token when needed.
Args:
input_ids: List of input token IDs.
logits: Tensor of logits for the next token.
Returns:
Processed logits tensor.
"""
# Check if we're in a thinking section
if self.think_start_token in input_ids and self.think_end_token not in input_ids:
# Find the position of the last <think> token
think_start_pos = len(input_ids) - 1 - input_ids[::-1].index(self.think_start_token)
# Calculate number of tokens since <think>
tokens_since_think = len(input_ids) - think_start_pos - 1
# If we've reached the maximum thinking length, force </think>
if tokens_since_think >= self.num_think_tokens:
# Set all other logits to -inf except for </think>
logits = torch.full_like(logits, float('-inf'))
logits[self.think_end_token] = 1.0
return logits
Here’s how to use it with vLLM:
(Note: after vLLM upgraded to the V1 engine, Python-based LogitsProcessors are no longer supported. The following code requires a pre-V1 version of vLLM, such as 0.7.3.)
from transformers import AutoTokenizer
from vllm import LLM, SamplingParams
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-4B")
think_start_token = tokenizer.convert_tokens_to_ids("<think>")
think_end_token = tokenizer.convert_tokens_to_ids("</think>")
sampling_params = SamplingParams(
temperature=temperature,
max_tokens=max_tokens,
logits_processors=[ThinkLogitsProcessor(think_start_token, think_end_token, num_think_tokens)]
)
llm = LLM(model="Qwen/Qwen3-4B")
conversations = ["your-conversations"]
outputs = llm.chat(messages=conversations,
sampling_params=sampling_params,
use_tqdm=False)
All of the above approaches use a hard-cutoff strategy — forcibly inserting </think>. The result can feel rather abrupt, and the model often continues a stretch of thinking in its response anyway, following its own momentum.
Refinements to this approach, such as a soft cutoff (linearly or polynomially increasing the output probability of </think> as the budget limit approaches), can improve things to some degree.
Beyond that, inserting special tokens or directly injecting natural-language prompts are also viable options.
Difficulty-Based Truncation
But a new problem arises: a one-size-fits-all budget is far from perfect. For simple questions, the budget is too long; for hard questions, it’s too short. Manually setting a budget isn’t practical during synthetic data generation and rollout. The most intuitive approach is to use a classification model to predict an optimal budget.
In TALE, the authors use a zero-shot prompt that asks the model to analyze the question and estimate the minimum number of tokens required to generate a complete and accurate response, then embed the budget value into the prompt, like:
Task: Analyze the given question and estimate the
minimum number of tokens required to generate a
complete and accurate response. Please Give the
response by strictly following this format: [[budget]],
for example, Budget: [[12]].
In practice, when examining model outputs, you’ll notice that the model easily falls into patterns of guessing, uncertainty, and self-doubt, such as:
Or, maybe XXX?
…
Maybe the answer is XX?
…
Maybe I made a calculation error?
…
Maybe I got something wrong?
…
In particularly long CoT traces — especially toward the tail end — the model essentially gets trapped in a doom loop of self-doubt and circular reasoning. One could perhaps capture these uncertainty markers and, when they appear frequently, boost the probability of </think> to trigger early stopping. However, this approach is really only suitable for small models — especially distilled reasoning models. The formulaic patterns of “self-doubt” and “uncertainty” seen in these small models likely originate from the teacher model, but the small model itself lacks (or struggles to develop) genuine self-correction capability. As a result, all those “wait…”, “But wait…”, “Alternatively…” — these may not actually contribute much to reasoning (this needs experimental corroboration).
Code
import torch
from typing import List, Dict, Tuple, Set, Optional
from collections import defaultdict
class UncertaintyThinkLogitsProcessor:
"""
A logits processor that limits thinking based on detected uncertainty patterns.
Optimized for efficiency and memory usage.
"""
def __init__(
self,
tokenizer,
think_start_token: int,
think_end_token: int,
uncertainty_phrases: Dict[str, float],
uncertainty_threshold: float = 5.0,
growth_strategy: str = "linear",
growth_factor: float = 1.0,
max_think_tokens: int = 1000,
debug: bool = False
):
"""Initialize the uncertainty-based think logits processor."""
self.tokenizer = tokenizer
self.think_start_token = think_start_token
self.think_end_token = think_end_token
self.uncertainty_threshold = uncertainty_threshold
self.growth_strategy = growth_strategy
self.growth_factor = growth_factor
self.max_think_tokens = max_think_tokens
self.debug = debug
# Process uncertainty phrases and build efficient lookup structures
self.uncertainty_phrases = {}
self.phrase_by_length = defaultdict(list) # Group phrases by length for efficient lookup
for phrase, score in uncertainty_phrases.items():
token_ids = tuple(tokenizer.encode(phrase, add_special_tokens=False))
self.uncertainty_phrases[token_ids] = score
self.phrase_by_length[len(token_ids)].append(token_ids)
# Sort by length for more efficient searching (shorter phrases first)
self.sorted_lengths = sorted(self.phrase_by_length.keys())
self.max_phrase_length = max(self.sorted_lengths) if self.sorted_lengths else 0
# State tracking - optimized data structures
self.accumulated_uncertainty = 0.0
self.last_processed_position = 0 # Track last processed position instead of full input length
self.detected_phrase_positions = set() # Use set for O(1) lookup, store positions only
# Pre-allocate tensor for logits modification to avoid repeated allocation
self._force_think_end_logits = None
if self.debug:
print(f"Initialized with {len(uncertainty_phrases)} uncertainty phrases")
print(f"Max phrase length: {self.max_phrase_length}")
print(f"Uncertainty threshold: {uncertainty_threshold}")
def _find_uncertainty_phrases(self, tokens: List[int], start_scan_pos: int = 0) -> List[Tuple[int, int, float]]:
"""
Efficiently find uncertainty phrases in tokens using optimized algorithm.
Returns list of (start_pos, end_pos, uncertainty_score) tuples.
"""
found_phrases = []
tokens_len = len(tokens)
# Only scan from the specified start position
for pos in range(start_scan_pos, tokens_len):
# Check phrases of each length, starting from shortest
for phrase_len in self.sorted_lengths:
if pos + phrase_len > tokens_len:
break # Can't fit this phrase length
# Extract candidate phrase
candidate = tuple(tokens[pos:pos + phrase_len])
# Check if this phrase exists in our lookup for this length
if candidate in self.phrase_by_length[phrase_len]:
# Verify it's actually in our main dictionary (should always be true)
if candidate in self.uncertainty_phrases:
found_phrases.append((pos, pos + phrase_len - 1, self.uncertainty_phrases[candidate]))
return found_phrases
def __call__(
self,
input_ids: List[int],
logits: torch.Tensor
) -> torch.Tensor:
"""Process logits based on uncertainty in thinking."""
# Convert input_ids to list if it's a tensor (avoid repeated conversion)
if isinstance(input_ids, torch.Tensor):
input_ids = input_ids.tolist()
# Early exit if no uncertainty phrases configured
if not self.uncertainty_phrases:
return logits
# Check if we're in a thinking section
try:
think_start_pos = len(input_ids) - 1 - input_ids[::-1].index(self.think_start_token)
in_thinking = self.think_end_token not in input_ids[think_start_pos:]
except ValueError:
# No think_start_token found
return logits
if in_thinking:
thinking_tokens = input_ids[think_start_pos + 1:]
tokens_since_think = len(thinking_tokens)
# Only process new tokens since last call
new_tokens_start = max(0, self.last_processed_position)
if tokens_since_think > new_tokens_start:
# Calculate scan start position to catch phrases spanning the boundary
scan_start = max(0, new_tokens_start - self.max_phrase_length + 1)
# Find new uncertainty phrases
new_phrases = self._find_uncertainty_phrases(thinking_tokens, scan_start)
# Process only truly new phrases (not already detected)
for start_pos, end_pos, base_score in new_phrases:
phrase_id = (start_pos, end_pos)
if phrase_id not in self.detected_phrase_positions:
# Apply growth strategy
if self.growth_strategy == "poly":
uncertainty_increase = base_score * (tokens_since_think ** self.growth_factor)
else: # linear or default
uncertainty_increase = base_score
self.accumulated_uncertainty += uncertainty_increase
self.detected_phrase_positions.add(phrase_id)
if self.debug:
phrase_tokens = thinking_tokens[start_pos:end_pos + 1]
phrase_text = self.tokenizer.decode(phrase_tokens)
print(f"Detected: '{phrase_text}' pos:{start_pos}-{end_pos} "
f"score:+{uncertainty_increase:.2f} total:{self.accumulated_uncertainty:.2f}")
# Update last processed position
self.last_processed_position = tokens_since_think
# Check if we should force </think>
should_force_end = (
self.accumulated_uncertainty >= self.uncertainty_threshold or
tokens_since_think >= self.max_think_tokens
)
if should_force_end:
if self.debug:
reason = ("uncertainty threshold" if self.accumulated_uncertainty >= self.uncertainty_threshold
else "max tokens")
print(f"Forcing </think> due to {reason}")
# Efficiently create force-end logits (reuse tensor if possible)
if self._force_think_end_logits is None or self._force_think_end_logits.shape != logits.shape:
self._force_think_end_logits = torch.full_like(logits, float('-inf'))
else:
self._force_think_end_logits.fill_(float('-inf'))
self._force_think_end_logits[self.think_end_token] = 1.0
return self._force_think_end_logits
else:
# Reset state when not in thinking mode or when thinking has ended
self._reset_state()
return logits
def _reset_state(self):
"""Reset the processor state efficiently."""
self.accumulated_uncertainty = 0.0
self.last_processed_position = 0
self.detected_phrase_positions.clear() # More efficient than creating new set
def reset(self):
"""Reset the processor state."""
self._reset_state()
# Also clear the cached tensor
self._force_think_end_logits = None
Usage
from vllm import LLM, SamplingParams
from transformers import AutoTokenizer
# Define uncertainty phrases
uncertainty_phrases = {
# Chinese uncertainty phrases
"可能我哪里算错了": 2.0,
"我不确定": 1.5,
"让我想想": 1.0,
"可能": 0.5,
"也许": 0.5,
"或许": 0.5,
"等等": 0.8,
"我需要重新思考": 1.2,
"有点困惑": 1.0,
"我搞错了": 1.5,
"我可能理解错了": 1.8,
"这个有点棘手": 1.0,
"需要再想想": 1.0,
"可能我的分析有误？": 2.0,
"可能我的理解有误？":2.0,
"？":1.0,
"正确？":1.0,
"或者": 1.0,
"这似乎不符":0.8,
# English uncertainty phrases
"I'm not sure": 1.5,
"I need to reconsider": 1.2,
"Let me think": 1.0,
"maybe": 0.5,
"perhaps": 0.5,
"wait": 0.8,
"hmm": 0.7,
"let me check again": 1.2,
"I might be wrong": 1.5,
"I'm confused": 1.2,
"I'm not certain": 1.5,
"?":1.0,
"That can't be right":1.0,
"Wait":0.7,
"right?":0.5
}
# Setup
llm = LLM(model="Qwen/Qwen3-4B")
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-4B")
# Get token IDs
think_start_token = tokenizer.convert_tokens_to_ids("<think>")
think_end_token = tokenizer.convert_tokens_to_ids("</think>")
uncertainty_processor = UncertaintyThinkLogitsProcessor(
tokenizer=tokenizer,
think_start_token=think_start_token,
think_end_token=think_end_token,
uncertainty_phrases=uncertainty_phrases,
uncertainty_threshold=5.0,
growth_strategy="linear",
max_think_tokens=1000,
debug=False
)
sampling_params = SamplingParams(
n=5,
temperature=0.7,
max_tokens=5000,
# top_p = 0.95,
# top_k = 20,
logits_processors=[uncertainty_processor]
)
Budget Guidance
The paper Steering LLM Thinking with Budget Guidance introduces the concept of Budget Guidance.
Drawing on the idea of Classifier Guidance from Diffusion Models, this approach does not force the model to output specific content or a particular format. Instead, it “softly” guides the generation process so that it naturally converges toward a target length. It does this through a lightweight predictor module (Auxiliary Predictor) that, at each token the LLM generates, predicts the length distribution of the remaining thinking process.
Based on Bayes’ theorem, the budget-constrained generation problem p(next token | context, remaining length <= budget) is transformed into the form p(next token | context) * P(remaining length <= budget | context, next token):
$$ p(Y_t|X,Y_{<t}, L_t\leq \bar{l}-t) \to p(Y_t|X,Y_{<t})\cdot Pr(L_t \leq \bar{l}-t|X,Y_{<t},Y_t) $$
Here p(Y_t|X,Y_{<t}) is the original LLM’s output probability, and Pr(L_t \leq \bar{l}-t|X,Y_{<t},Y_t) is the guidance score provided by the Auxiliary Predictor.
In this paper, the Auxiliary Predictor uses a BERT-base model. Rather than outputting a single length prediction, it models the logarithm of the remaining length as a Gamma distribution. BERT takes as input all hidden states from the LLM’s final layer at each generation step, and outputs two parameters — the Gamma distribution’s shape parameter lambda and rate parameter alpha. Using the Gamma distribution’s CDF, it computes a probability value based on the current remaining budget (target total budget - length generated so far) — this is the guidance score.
Predictor training: On OpenR1-Math-220k, a complete thinking process is truncated at different positions, and the predictor learns to estimate the length required from the truncation point to the end. The training objective is to maximize the log-likelihood of the predicted Gamma distribution, making the predicted distribution as close as possible to the true remaining length.
Is Longer Thinking Really Better?
In Lost at the Beginning of Reasoning, the authors discovered that the model’s reasoning process has an extremely strong dependency on the initial reasoning step. Once the first step goes wrong, the model struggles to correct the error in subsequent steps, which means the final answer is highly likely to be wrong too — hence “Lost at the Beginning.”
Building on this finding, the authors proposed a sampling strategy called Early Pruning. The core idea: first generate multiple short candidates for the first reasoning step, then use a reward model to quickly evaluate their quality, keep only the top few, and continue generating the full chain of thought exclusively along those promising paths.
They also created a new benchmark called LaBoR (Lost at the Beginning of Reasoning) — a benchmark specifically designed to measure how well long chain-of-thought models can self-correct when the opening move is flawed. Each sample consists of a question paired with a deliberately flawed first reasoning step.
Closing Thoughts
This article has focused on how to effectively control an LLM’s thinking budget, though it’s worth noting that the discussion has primarily centered on inference-side solutions. The core advantage of these methods lies in their flexibility and “plug-and-play” nature — they require no modifications to the model itself, making them ideal for local deployment and rapid validation.
This, however, is also the limitation of this article. Inference-side control is more like an “external constraint,” which inevitably feels somewhat crude or lacks robustness. A more comprehensive, more fundamental thinking-budget solution necessarily requires the interplay between inference and training. In particular, during post-training, markers for different reasoning intensities can be introduced alongside specialized training, so that on the inference side, more efficient thinking can be activated through instructions.
References
Claude Code: Best practices for agentic coding
Gemini API: Set a thinking budget
Token-Budget-Aware LLM Reasoning