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

• URL: http://arxiv.org/abs/2506.11274v1
• Date: Thu, 12 Jun 2025 20:28:54 GMT
• Title: Learning a Continue-Thinking Token for Enhanced Test-Time Scaling
• Abstract summary: Test-time scaling has emerged as an effective approach for improving language model performance by utilizing additional compute at inference time.<n>Recent studies have shown that overriding end-of-thinking tokens can extend reasoning steps and improve accuracy.<n>In this work, we explore whether a dedicated continue-thinking token can be learned to trigger extended reasoning.
• Score: 15.231981094907386
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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 "</think>" 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.

Related papers

• Prefix Sliding for efficient test-time scaling [148.8911065058742]
  Test-time scaling uses extra test-time compute to improve performance.<n>We find most intermediate reasoning tokens lose importance as the model continues reasoning.<n>We propose Prefix Sliding, which discards tokens during reasoning that are not part of the prefix or the window of the last few thousand tokens.
  arXiv  Detail & Related papers  (2026-08-26T17:37:15Z)
• Understanding and Mitigating Premature Confidence for Better LLM Reasoning [76.16007941549857]
  Long chains of thought (CoT) from current language models frequently contain logical gaps and unjustified leaps.<n>We find such a signal in how the model's confidence evolves during reasoning.<n>We exploit this in progressive confidence shaping, a reinforcement learning objective that trains models to update their confidence as they reason rather than commit early.
  arXiv  Detail & Related papers  (2026-05-23T04:42:45Z)
• Unlearning What Matters: Token-Level Attribution for Precise Language Model Unlearning [5.454773103061359]
  TokenUnlearn is a token-level attribution framework that identifies and selectively targets critical tokens.<n>Our approach combines knowledge-aware signals via masking, and entropy-aware signals to yield importance scores for precise token selection.
  arXiv  Detail & Related papers  (2026-05-01T02:59:03Z)
• On the Direction of RLVR Updates for LLM Reasoning: Identification and Exploitation [81.29776444949539]
  We argue that the direction of updates is a more critical lens for understanding RLVR's effects.<n>Our work establishes the direction of change as a key principle for analyzing and improving RLVR.
  arXiv  Detail & Related papers  (2026-03-23T15:42:24Z)
• CarBoN: Calibrated Best-of-N Sampling Improves Test-time Reasoning [62.56541355300587]
  We introduce a general test-time calibration framework that adaptively modifies the model toward high-reward reasoning paths.<n>Within this framework, we propose CarBoN, a two-phase method that first explores the solution space and then learns a calibration of the logits.<n>Experiments on MATH-500 and AIME-2024 show that CarBoN improves efficiency, with up to $4times$ fewer rollouts to reach the same accuracy.
  arXiv  Detail & Related papers  (2025-10-17T14:04:37Z)
• LaSeR: Reinforcement Learning with Last-Token Self-Rewarding [54.72617309922891]
  Reinforcement Learning with Verifiable Rewards (RLVR) has emerged as a core paradigm for enhancing the reasoning capabilities of Large Language Models (LLMs)<n>Previous practice requires the LLM to sequentially generate solutions and self-verifications using two separate prompt templates, which significantly reduces efficiency.<n>We propose LaSeR (Reinforcement Learning with Last-Token Self-Rewarding), an algorithm that simply augments the original RLVR loss with a MSE loss.
  arXiv  Detail & Related papers  (2025-10-16T17:55:11Z)
• Broken Tokens? Your Language Model can Secretly Handle Non-Canonical Tokenizations [83.93566096400723]
  We find that instruction-tuned models retain up to 93.4% of their original performance when given a randomly sampled tokenization.<n>Character-level segmentation improves string manipulation and code understanding tasks by up to +14%.<n>Right-aligned digit grouping enhances large-number arithmetic by +33%.
  arXiv  Detail & Related papers  (2025-06-23T18:02:26Z)
• Multipole Attention for Efficient Long Context Reasoning [64.94673641704289]
  Large Reasoning Models (LRMs) have shown promising accuracy improvements on complex problem-solving tasks.<n>LRMs need to generate long chain-of-thought reasoning in order to think before answering.<n>We introduce Multipole Attention, which accelerates autoregressive reasoning by only computing exact attention for the most important tokens.
  arXiv  Detail & Related papers  (2025-06-16T03:00:40Z)
• Learning to Insert [PAUSE] Tokens for Better Reasoning [6.823521786512908]
  We introduce a novel approach termed Dynamic Inserting Tokens Training (DIT)<n>Our method identifies positions within sequences where model confidence is lowest according to token log-likelihood.<n>With this simple yet effective method, we achieve accuracy gains of up to 4.7%p on GSM8K, 3.23%p on AQUA-RAT, and pass@1 improvements of up to 3.4%p on MBPP datasets.
  arXiv  Detail & Related papers  (2025-06-04T06:48:41Z)
• Fast Quiet-STaR: Thinking Without Thought Tokens [51.79231070632772]
  Fast Quiet STaR is a more efficient reasoning framework that preserves the benefits of token-level reasoning while reducing computational cost.<n>Our method introduces a curriculum learning based training strategy that gradually reduces the number of thought tokens.<n>Experiments on four benchmark datasets with Mistral 7B and Qwen2.5 7B demonstrate that Fast Quiet-STaR consistently outperforms Quiet-STaR in terms of average accuracy.
  arXiv  Detail & Related papers  (2025-05-23T11:14:12Z)
• KTAE: A Model-Free Algorithm to Key-Tokens Advantage Estimation in Mathematical Reasoning [19.25257653416883]
  Key-token Advantage Estimation (KTAE) is a novel algorithm that estimates fine-grained, token-level advantages without introducing additional models.<n>We show that models trained with GRPO+KTAE and DAPO+KTAE outperform baseline methods across five mathematical reasoning benchmarks.
  arXiv  Detail & Related papers  (2025-05-22T16:00:33Z)
• S$^2$R: Teaching LLMs to Self-verify and Self-correct via Reinforcement Learning [51.84977135926156]
  We introduce S$2$R, an efficient framework that enhances LLM reasoning by teaching models to self-verify and self-correct during inference.<n>Our results demonstrate that Qwen2.5-math-7B achieves an accuracy improvement from 51.0% to 81.6%, outperforming models trained on an equivalent amount of long-CoT distilled data.
  arXiv  Detail & Related papers  (2025-02-18T13:40:22Z)
• Rethinking Fine-Tuning when Scaling Test-Time Compute: Limiting Confidence Improves Mathematical Reasoning [32.45574194957491]
  We show that training with cross-entropy loss can be misaligned with pass@N in that pass@N accuracy $it decreases$ with longer training.<n>We suggest a principled, modified training loss that is better aligned to pass@N by limiting model confidence and rescuing pass@N test performance.
  arXiv  Detail & Related papers  (2025-02-11T00:33:31Z)
• T1: Advancing Language Model Reasoning through Reinforcement Learning and Inference Scaling [52.34735382627312]
  Large language models (LLMs) have demonstrated remarkable capabilities in complex reasoning tasks.<n>Existing approaches mainly rely on imitation learning and struggle to achieve effective test-time scaling.<n>We present T1 to scale reinforcement learning by encouraging exploration and understand inference scaling.
  arXiv  Detail & Related papers  (2025-01-20T18:33:33Z)
• Exact Byte-Level Probabilities from Tokenized Language Models for FIM-Tasks and Model Ensembles [23.134664392314264]
  Tokenization is associated with many poorly understood shortcomings in language models (LMs)<n>This work studies how tokenization impacts model performance by analyzing and comparing models with their byte-level counterparts.<n>We introduce the Byte-Token Representation Lemma, a framework that establishes a mapping between the learned token distribution and its equivalent byte-level distribution.
  arXiv  Detail & Related papers  (2024-10-11T23:30:42Z)
• Masked Thought: Simply Masking Partial Reasoning Steps Can Improve Mathematical Reasoning Learning of Language Models [102.72940700598055]
  In reasoning tasks, even a minor error can cascade into inaccurate results. We develop a method that avoids introducing external resources, relying instead on perturbations to the input. Our training approach randomly masks certain tokens within the chain of thought, a technique we found to be particularly effective for reasoning tasks.
  arXiv  Detail & Related papers  (2024-03-04T16:21:54Z)
• One-bit Supervision for Image Classification: Problem, Solution, and Beyond [114.95815360508395]
  This paper presents one-bit supervision, a novel setting of learning with fewer labels, for image classification. We propose a multi-stage training paradigm and incorporate negative label suppression into an off-the-shelf semi-supervised learning algorithm. In multiple benchmarks, the learning efficiency of the proposed approach surpasses that using full-bit, semi-supervised supervision.
  arXiv  Detail & Related papers  (2023-11-26T07:39:00Z)
• The Right Tool for the Job: Matching Model and Instance Complexities [62.95183777679024]
  As NLP models become larger, executing a trained model requires significant computational resources incurring monetary and environmental costs. We propose a modification to contextual representation fine-tuning which, during inference, allows for an early (and fast) "exit" We test our proposed modification on five different datasets in two tasks: three text classification datasets and two natural language inference benchmarks.
  arXiv  Detail & Related papers  (2020-04-16T04:28:08Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.