Learning a Continue-Thinking Token for Enhanced Test-Time Scaling

2025-06-12

AI Summary

This work addresses the challenge of adaptively controlling reasoning depth in language models. We propose a learnable “continue-thinking” token (<|continue-thinking|>), whose embedding is optimized end-to-end via reinforcement learning while keeping the backbone parameters frozen. Unlike hand-crafted, fixed trigger tokens (e.g., “Wait”), this mechanism enables differentiable, dynamic control over the number of reasoning steps during inference. Evaluated on a distilled version of DeepSeek-R1, our approach achieves a 4.2% absolute accuracy gain on GSM8K over the baseline—outperforming fixed-token baselines by +1.3%—and demonstrates superior generalization across multiple mathematical reasoning benchmarks. Our key contribution is the first lightweight, differentiable, and end-to-end learnable reasoning trigger token, effectively balancing inference efficiency with controllability.

Technology Category

• Cognitive Modeling & Cognitive Systems: Conceptual Inference and Reasoning
• Knowledge Representation and Reasoning: Reasoning with Beliefs
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

Test-time scaling has emerged as an effective approach for improving language model performance by utilizing additional compute at inference time. Recent studies have shown that overriding end-of-thinking tokens (e.g., replacing""with"Wait") can extend reasoning steps and improve accuracy. In this work, we explore whether a dedicated continue-thinking token can be learned to trigger extended reasoning. We augment a distilled version of DeepSeek-R1 with a single learned"<|continue-thinking|>"token, training only its embedding via reinforcement learning while keeping the model weights frozen. Our experiments show that this learned token achieves improved accuracy on standard math benchmarks compared to both the baseline model and a test-time scaling approach that uses a fixed token (e.g.,"Wait") for budget forcing. In particular, we observe that in cases where the fixed-token approach enhances the base model's accuracy, our method achieves a markedly greater improvement. For example, on the GSM8K benchmark, the fixed-token approach yields a 1.3% absolute improvement in accuracy, whereas our learned-token method achieves a 4.2% improvement over the base model that does not use budget forcing.

Problem

• Learning a token to extend reasoning steps
• Improving model accuracy with continue-thinking token
• Enhancing test-time scaling via reinforcement learning

Innovation

• Learned continue-thinking token for extended reasoning
• Reinforcement learning to train token embedding
• Frozen model weights with single token augmentation