Learning When to Think: Shaping Adaptive Reasoning in R1-Style Models via Multi-Stage RL

Songjun Tu; Jiahao Lin; Qichao Zhang; Xiangyu Tian; Linjing Li; Xiangyuan Lan; Dongbin Zhao

Source-linked AI summary

Learning When to Think: Shaping Adaptive Reasoning in R1-Style Models via Multi-Stage RL

Songjun Tu, Jiahao Lin, Qichao Zhang, Xiangyu Tian, Linjing Li, Xiangyuan Lan, Dongbin Zhao

arXiv:2505.10832v3May 16, 2025cs.CLcs.AI

TL;DR

Overthinking makes explicit reasoning costly, motivating models that decide when reasoning is necessary. AutoThink combines ellipsis-triggered dual-mode behavior with multi-stage RL reward shaping to learn adaptive reasoning. It reports favorable accuracy–efficiency trade-offs across mathematical benchmarks, while retaining limitations in behavioral separation, length control, and training-data selection.

Problem
Excessive reasoning increases computational cost and latency, while existing concise-reasoning approaches may enforce uniform brevity rather than adapt effort to problem complexity.
Method
AutoThink uses an ellipsis prompt and multi-stage RL with stage-wise reward shaping to train R1-style models to select thinking or no-thinking behavior adaptively.
Results
AutoThink achieves favorable accuracy–efficiency trade-offs and outperforms prompting and RL baselines without compromising performance.
Takeaways & Limitations
AutoThink offers a scalable and controllable approach to efficient reasoning in R1-style language models.
Takeaways & Limitations
AutoThink may leave reasoning-related tokens in no-thinking mode, cannot control overall response length, and uses unfiltered training data.

Abstract

from arXiv · show

Large reasoning models (LRMs) are proficient at generating explicit, step-by-step reasoning sequences before producing final answers. However, such detailed reasoning can introduce substantial computational overhead and latency, particularly for simple problems. To address this over-thinking problem, we explore how to equip LRMs with adaptive thinking capabilities: enabling them to dynamically decide whether or not to engage in explicit reasoning based on problem complexity. Building on R1-style distilled models, we observe that inserting a simple ellipsis ("...") into the prompt can stochastically trigger either a thinking or no-thinking mode, revealing a latent controllability in the reasoning behavior. Leveraging this property, we propose AutoThink, a multi-stage reinforcement learning (RL) framework that progressively optimizes reasoning policies via stage-wise reward shaping. AutoThink learns to invoke explicit reasoning only when necessary, while defaulting to succinct responses for simpler tasks. Experiments on five mainstream mathematical benchmarks demonstrate that AutoThink achieves favorable accuracy-efficiency trade-offs compared to recent prompting and RL-based pruning methods. It can be seamlessly integrated into any R1-style model, including both distilled and further fine-tuned variants. Notably, AutoThink improves relative accuracy by 6.4 percent while reducing token usage by 52 percent on DeepSeek-R1-Distill-Qwen-1.5B, establishing a scalable and adaptive reasoning paradigm for LRMs. Project Page: https://github.com/ScienceOne-AI/AutoThink.

1 Introduction

R1-style models use explicit chain-of-thought reasoning to improve accuracy, but excessive reasoning creates computational cost and latency. AutoThink addresses this overthinking problem by learning when to reason and when to respond concisely.

Explicit step-by-step reasoning supports self-reflection, backtracking, and validation in R1-style models.
AutoThink uses an ellipsis prompt and three-stage RL training to learn adaptive thinking behavior.The stages stabilize dual-mode behavior, reinforce accurate reasoning, and prune redundant reasoning through length-aware rewards.
The ellipsis prompt stochastically switches R1-style models between thinking and no-thinking modes.
Experiments on mathematical benchmarks report better accuracy–efficiency trade-offs than existing pruning and compression methods without sacrificing performance.

2 An Ellipsis Unlocks Random Thinking in R1-Style Models

A single ellipsis can make R1-style models either reason explicitly or skip reasoning, reducing token use but not initially aligning thinking with problem difficulty. Multi-stage RL is introduced to learn that missing difficulty-aware behavior.

Uniformly concise reasoning can shorten responses but may degrade performance on harder instances.
The prompt <think>\n...\n can trigger either immediate termination or a full derivation without additional training.
Ellipsis prompting reduces token usage while providing a more balanced accuracy–efficiency trade-off than no-thinking prompting.
Without additional training, ellipsis prompting shows no clear relationship between no-thinking rate and problem difficulty.The flat distribution across eight MATH500 difficulty levels indicates that the switching behavior is unguided.
Multi-stage RL produces difficulty-aware reasoning by rewarding appropriate effort relative to problem difficulty.

3 Guiding When to Think via Multi-Stage Reinforcement Learning

AutoThink uses three-stage GRPO training with progressively refined rewards to stabilize dual-mode reasoning, improve task performance, and reduce unnecessary reasoning length. The stages combine batch-level mode balancing, direct correctness optimization, and length-aware reward modulation.

Overall framework: AutoThink applies GRPO with a token-level policy-gradient objective across three training stages.The objective uses sampled outputs, token-level importance weights, and estimated token-level advantages.
Stage 1: Batch reward balance: Stage 1 assigns naive rewards from thinking status and correctness, then uses batch-level balancing to prevent collapse into either mode.The balance mechanism adjusts rewards according to the batch proportion of thinking trajectories and a target ratio.
Stage 1: Batch reward balance: Stage 1’s soft penalty modulation reduces rewards for the overrepresented behavior, gently pushing the policy toward a target thinking balance.The target ratio is γ, while λ controls the penalty slope.
Stage 2: Reliable dual-mode behavior: Stage 2 removes batch-level balancing and directly optimizes the naive correctness reward to improve reasoning quality and accuracy within both modes.A larger training context budget allows longer responses when needed, while Stage 1 regularization helps the thinking proportion remain balanced.
Stage 3: Length-aware reward: Stage 3 introduces length-aware reward modulation that encourages brevity in no-thinking mode and rewards elaboration when warranted.Response length is standardized within query groups, with α and β controlling shaping sensitivity; the reward decays with length for correct responses and grows for incorrect ones.

4 Experiments

Experiments evaluate AutoThink across five mathematical benchmarks and three R1-style models, showing staged gains in adaptive reasoning, accuracy, and token efficiency.

Setup: Experiments use 40K mathematical problems for training and evaluate MATH, Minerva, Olympiad, AIME24, and AMC23 on three R1-style models.The models are Distill-R1-1.5B/7B and DeepScaleR-Preview-1.5B.
Setup: AutoThink is compared with standard, no-thinking, and ellipsis prompts, RL-based pruning methods, and contextual open-source reference models.The RL baselines include Concise-RL, ShorterBetter, and ThinkPrune.
Evaluation: E-F1 jointly measures accuracy and token usage, with positive values indicating improvement over the standard baseline on both dimensions.Its subscripts distinguish standard and no-thinking baselines.
Main Results: Prompt-only baselines reduce response length but can sacrifice accuracy, while RL baselines provide limited Distill-R1 gains and sometimes lower DeepScaleR accuracy.Some open-source models produce outputs up to twice as long as AutoThink at the same model size.
Main Results: AutoThink’s Stage 2 improves accuracy across backbones, and Stage 3 further reduces token usage while minimizing performance degradation.On Distill-R1-1.5B, Stage 3 reaches 51.7% accuracy with half the standard-prompt token usage; on DeepScaleR, Stage 2 improves performance by 0.6 and cuts usage by 10%.
Ablation Study: Batch reward balancing prevents Stage 1 collapse into either all-thinking or all-no-thinking behavior.Naive rewards caused thinking-mode collapse, whereas length-aware rewards encouraged degenerate no-thinking behavior.
Ablation Study: Skipping Stage 2 leaves accuracy stagnant and response length rebounds, yielding 47.6% accuracy versus 51.7% after full training.The results support reinforcement before pruning for stable, discriminative reasoning behavior.
In-Depth Behavioral and Efficiency Analysis: Thinking frequency and response length increase with task difficulty, while stronger models reason less frequently yet outperform weaker models.Accuracy serves as the proxy for dataset difficulty.

5 Related Works

Related work addresses LLM overthinking through reinforcement fine-tuning, prompting, supervised short reasoning, length-aware rewards, and model-assisted acceleration.

RL-based Post-Training for LLMs: Reinforcement fine-tuning improves LLM reasoning but can induce verbose reasoning with limited additional benefit.Recent RL work also targets training efficiency and effectiveness through sampling, value-bias mitigation, reweighting, and ensembling.
Mitigating Overthinking for LLMs: Prompting, manual difficulty control, and DPO with short chains of thought suppress or shorten reasoning but can reduce performance or enforce fixed behavior.Other approaches include terminating on unsolvable problems and inserting pseudo-thinking cues.
Mitigating Overthinking for LLMs: Alternative efficiency methods use short-chain supervised fine-tuning, response-length-aware RL rewards, or smaller models to guide larger models toward faster reasoning.These methods motivate adaptive rather than uniformly concise reasoning.

6 Conclusion & Limitations

AutoThink combines minimal prompting with multi-stage reinforcement learning to enable adaptive, task-aware reasoning in R1-style models. It improves efficiency while maintaining performance, but limitations remain in behavioral separation, response-length control, and training-data selection.

Conclusion: AutoThink pairs an ellipsis prompt with three-stage reinforcement learning to stabilize reasoning patterns, reinforce effective behaviors, and prune unnecessary steps.The framework learns a dual-mode policy for thinking and no-thinking behavior.
Conclusion: Experiments report favorable accuracy–efficiency trade-offs, outperforming prompting and RL baselines without compromising performance.
Limitations: Limitations include reward hacking after </think>, uncontrolled overall response length, and unfiltered training data without task-difficulty selection.
Conclusion: No-thinking behavior is defined as immediately closing the <think> tag without substantive reasoning before producing the final answer.
Additional findings: Prompting alone with Think By Difficulty did not meaningfully produce selective thinking and slightly reduced both accuracy and token consumption.

B.1 Additional Results on Skywork-OR1-Math-7B

On Skywork-OR1-Math-7B, the ellipsis prompt alone has limited effect, but staged AutoThink training amplifies latent no-thinking behavior and substantially reduces reasoning length while largely preserving accuracy. Training dynamics show that balanced mode allocation is associated with shorter responses and improved overall accuracy.

Prompt sensitivity: The ellipsis prompt induces only a marginal proportion of no-thinking responses in Skywork-OR1-Math-7B, indicating limited prompt sensitivity.The model is described as having a deterministic reasoning policy.
Adaptive training: Stage 1 batch-level contrastive training captures and amplifies latent no-thinking behavior, while Stages 2 and 3 progressively refine the resulting reasoning patterns.The three stages use 600, 500, and 30 training steps, respectively.
Efficiency: Nearly 60% fewer reasoning tokens, from 9053 to 3974, are achieved in the final stage while task accuracy remains within less than 2% degradation versus the standard prompting baseline.
Training dynamics: Stage 1 training shifts responses toward no-thinking behavior between steps 100 and 200, reducing average response length as the target balance ratio is set to γ = 0.5.
Training dynamics: Overall accuracy continues improving even as thinking-response accuracy slightly decreases, consistent with skipping unnecessary reasoning on simpler problems.

B.3 Hyperparameter Sensitivity

AutoThink uses simple, interpretable stage-specific reward shaping with limited tuning. Sensitivity analyses indicate that the hyperparameters regulate mode balance and response brevity as intended, while overall training trends remain consistent.

Stage 1: Stage 1 balances thinking and no-thinking behavior through a batch-wise reward balance to prevent mode collapse.
Stage 2: Stage 2 reinforces accuracy within each reasoning mode without adding further reward shaping.
Stage 3: Stage 3 uses length-aware shaping to encourage brevity for correct responses and elaboration for incorrect ones.
Hyperparameters: Stage 1 hyperparameters γ and λ control modal balance, with increasing γ encouraging more thinking and larger λ enforcing stricter adherence to the target balance.
Robustness: Sensitivity results show consistent overall training trends, supporting AutoThink’s robustness without extensive hyperparameter tuning.
Hyperparameters: Stage 3 parameters α and β regulate reward decay or growth with response length, while larger α accelerates length decay for correct responses.

B.4 Training Cost Comparison

AutoThink’s three-stage training has a cost comparable to concise-reasoning baselines while using longer contexts and shortening responses through Stage 3. It can be completed within one day for 1.5B models and 2.5 days for 7B models on the stated H100 setup.

Training schedule: AutoThink uses a three-stage schedule with increasing context lengths and 500 total steps, comparable to ThinkPrune’s 540 steps.
Cost comparison: Unlike methods that compress by reducing context length, AutoThink expands context but prunes through shorter Stage 3 responses, resulting in comparable training cost.
Runtime: On H100 clusters with four nodes, all stages complete within one day for 1.5B models and 2.5 days for 7B models.
Cost comparison: AutoThink operates within the same order of GPU-hour compute as ThinkPrune and ShorterBetter while achieving notably stronger performance.
Cost comparison: DeepScaleR requires more than 3× higher compute because of its longer context length and increased reinforcement-learning iterations.

C Addressing Potential Challenges

AutoThink’s remaining challenges concern compute-budget control, dataset noise, and reward hacking. The paper presents these as complementary mechanisms for future integration rather than current components of the framework.

Token-Budget Control: Token-budget constraints could be integrated with AutoThink to enforce global compute budgets and further improve efficiency.Prior budget-aware rewards penalize excessively long completions.
Dataset Noise: Noise in large-scale reasoning datasets can hinder training efficiency, while curriculum learning or correctness- and difficulty-based filtering may improve learning quality.These strategies are described as orthogonal to AutoThink’s reward design.
Reward Hacking: Reward hacking can occur when the model continues reasoning after the </think> tag.The paper suggests penalizing reasoning patterns outside the <think> span or rewarding clean thought-answer separation.
Future Integration: The paper views these solutions as complementary and composable with AutoThink, with tighter integration left for future work.The authors frame these mechanisms as directions toward more comprehensive reasoning control.

D Case Study

The case study shows AutoThink varying its reasoning mode with problem difficulty: fast no-thinking responses for easy problems, mixed behavior for medium problems, and deeper verification for hard problems.

Easy Case: For easy problems, AutoThink produces correct answers without explicit reasoning, using few tokens and reflecting effective fast thinking.Figure 11 identifies this behavior as the no-thinking mode.
Medium Case: For medium problems, AutoThink may activate both thinking and no-thinking modes, so the two response types can coexist.Figure 12 presents this mixed-mode behavior.
Hard Case: For hard problems, AutoThink engages in deeper reasoning with iterative understanding and self-verification before reaching the correct solution.Figure 13 characterizes this as thinking mode with repeated verification.