Think or Not? Selective Reasoning via Reinforcement Learning for Vision-Language Models
Proposes TON framework with Thought Dropout and GRPO for selective reasoning, reducing 90% inference length while maintaining accuracy.
Key Findings
Methodology
TON employs a two-stage training process: first, supervised fine-tuning (SFT) with Thought Dropout, randomly replacing reasoning traces with empty thoughts to train the model to recognize when to skip reasoning; second, reinforcement learning via Group Relative Policy Optimization (GRPO), enabling the model to explore when to think or not, optimizing for task-specific rewards. The Thought Dropout introduces a format that encourages the model to learn skipping unnecessary reasoning steps. During GRPO, multiple candidate responses are sampled, and the model adjusts its policy based on relative advantages computed from reward signals, promoting diversity and efficiency. Extensive experiments on tasks like GSM8K, CLEVR, GeoQA, and AITZ demonstrate that TON reduces inference length by up to 90%, with performance comparable or better than baseline methods.
Key Results
- On CLEVR and GeoQA, TON reduces reasoning length by 87% and 65%, respectively, while improving accuracy by 17% on GeoQA. In AITZ navigation tasks, output length drops from 3.6K to 0.9K tokens, saving 70% of reasoning tokens. Across different model sizes (3B, 7B), TON maintains high accuracy and demonstrates strong generalization, with skip ratios increasing during training, indicating adaptive learning of when to think. Ablation studies confirm the correlation between skipping behavior and reward improvements, validating the effectiveness of the approach.
- Compared to vanilla GRPO, TON achieves significant inference speedups without sacrificing accuracy. Notably, in GeoQA, shorter reasoning trajectories outperform longer ones, exemplifying a 'free-lunch' effect. Multi-domain and out-of-distribution tests show that TON generalizes well, maintaining performance while drastically reducing output length. The approach offers a scalable solution for efficient reasoning in large models, with potential applications in real-world AI systems.
- Overall, the method not only enhances inference efficiency but also improves reasoning quality by avoiding unnecessary steps, paving the way for more human-like, resource-efficient AI systems.
Significance
This work addresses a fundamental bottleneck in current vision-language reasoning models: the inefficiency caused by uniform full-length reasoning. By enabling models to decide when to think, the approach mimics human cognitive strategies, leading to substantial resource savings. The ability to dynamically skip reasoning steps without performance loss has broad implications for deploying large models in real-time applications, reducing computational costs and latency. It also opens new research avenues in adaptive reasoning, where models learn to balance accuracy and efficiency based on task complexity. This paradigm shift could accelerate AI adoption in resource-constrained environments and inspire further innovations in intelligent decision-making systems.
Technical Contribution
The core technical innovation lies in integrating Thought Dropout with the GRPO reinforcement learning framework to facilitate adaptive, selective reasoning. Thought Dropout introduces stochastic masking of reasoning traces during supervised training, enabling the model to learn formats that support skipping. The GRPO algorithm then allows the model to explore and optimize when to invoke reasoning, guided by task-specific rewards. This combination results in a policy that dynamically adjusts reasoning depth, reducing unnecessary computation. The approach also introduces a novel reward design that encourages skipping for easy tasks while maintaining accuracy on complex ones. The method is validated across multiple benchmarks, demonstrating its robustness and scalability.
Novelty
This is the first work explicitly modeling 'when to think' in vision-language models through a combination of Thought Dropout and reinforcement learning. Unlike prior methods that enforce full reasoning chains or rely on external control mechanisms, this approach enables the model to autonomously decide the necessity of reasoning steps, aligning more closely with human cognitive strategies. The integration of format-based training with self-exploration via GRPO represents a significant departure from existing paradigms, offering a new direction for improving inference efficiency without compromising accuracy.
Limitations
- The current approach relies on the assumption that the model can accurately distinguish between simple and complex tasks during training, which may not hold in highly ambiguous scenarios. The effectiveness of Thought Dropout depends on the quality of the initial training data; noisy or biased data could impair learning. Additionally, the exploration process during reinforcement learning incurs computational overhead, especially for large models and complex tasks. The method's performance in real-world, dynamic environments with evolving task distributions remains to be validated. Future work should focus on improving reward design, robustness, and extending applicability to broader domains.
Future Work
Future research will explore more sophisticated reward mechanisms, including human-in-the-loop feedback, to better guide the model's reasoning decisions. Integrating multi-task learning and transfer learning could further enhance adaptability across diverse scenarios. Additionally, extending the framework to multimodal tasks involving audio, video, or sensor data can broaden its impact. Investigating explainability and interpretability of the reasoning skipping decisions will be crucial for practical deployment. Finally, optimizing training efficiency and scaling to larger models will be vital for real-world applications, such as autonomous agents and real-time AI assistants.
AI Executive Summary
In the rapidly evolving field of vision-language modeling, a persistent challenge has been balancing reasoning depth with computational efficiency. Traditional models, such as those based on Chain-of-Thought prompting, tend to perform exhaustive reasoning regardless of task complexity, leading to significant resource consumption and latency. This inefficiency becomes particularly problematic when deploying large-scale models in real-time or resource-constrained environments. Recognizing this, researchers have sought methods to enable models to adaptively decide when to reason, akin to human cognitive strategies.
This paper introduces the Think-or-Not (TON) framework, a novel approach that empowers models with the ability to selectively perform reasoning steps. The key innovation lies in a two-stage training process. First, during supervised fine-tuning (SFT), the model is trained with a Thought Dropout mechanism, where reasoning traces are randomly replaced with empty placeholders. This trains the model to recognize situations where reasoning can be skipped. Next, in the reinforcement learning phase, the model employs Group Relative Policy Optimization (GRPO) to explore and learn when to invoke reasoning, guided by task-specific rewards that favor shorter, accurate responses.
The experimental results are compelling. On benchmarks like GSM8K, CLEVR, GeoQA, and AITZ, TON reduces inference length by up to 90% without sacrificing accuracy. In some cases, shorter reasoning trajectories even outperform longer ones, demonstrating a 'free-lunch' effect. The model also generalizes well across different tasks, model sizes, and out-of-distribution scenarios, confirming its robustness and scalability.
This work marks a significant step toward human-like reasoning in AI systems. By enabling models to decide when to think, it addresses a core inefficiency in current approaches, paving the way for faster, more resource-efficient intelligent agents. Despite some limitations, such as reliance on training data quality and exploration costs, the framework opens promising avenues for future research in adaptive, efficient reasoning strategies across diverse AI applications.
Deep Dive
Abstract
Reinforcement Learning (RL) has proven to be an effective post-training strategy for enhancing reasoning in vision-language models (VLMs). Group Relative Policy Optimization (GRPO) is a recent prominent method that encourages models to generate complete reasoning traces before answering, leading to increased token usage and computational cost. Inspired by the human-like thinking process-where people skip reasoning for easy questions but think carefully when needed-we explore how to enable VLMs to first decide when reasoning is necessary. To realize this, we propose TON, a two-stage training strategy: (i) a supervised fine-tuning (SFT) stage with a simple yet effective 'thought dropout' operation, where reasoning traces are randomly replaced with empty thoughts. This introduces a think-or-not format that serves as a cold start for selective reasoning; (ii) a GRPO stage that enables the model to freely explore when to think or not, while maximizing task-aware outcome rewards. Experimental results show that TON can reduce the completion length by up to 90% compared to vanilla GRPO, without sacrificing performance or even improving it. Further evaluations across LLM (GSM8K), VLM (CLEVR, Super-CLEVR, GeoQA), and Agentic (AITZ) tasks-covering a range of reasoning difficulties under both 3B and 7B models-consistently reveal that the model progressively learns to bypass unnecessary reasoning steps as training advances. These findings shed light on the path toward human-like reasoning patterns in RL approaches. Our code is available at https://github.com/kokolerk/TON.