Apprentissage d'un jeton de pensée continue pour une mise à l'échelle améliorée lors des tests

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

June 12, 2025

papers.authors: Liran Ringel, Elad Tolochinsky, Yaniv Romano

cs.AI

papers.abstract

Le redimensionnement au moment du test est apparu comme une approche efficace pour améliorer les performances des modèles de langage en utilisant des ressources de calcul supplémentaires lors de l'inférence. Des études récentes ont montré que le remplacement des jetons de fin de réflexion (par exemple, remplacer "</think>" par "Wait") peut prolonger les étapes de raisonnement et améliorer la précision. Dans ce travail, nous explorons si un jeton dédié de continuation de réflexion peut être appris pour déclencher un raisonnement prolongé. Nous enrichissons une version distillée de DeepSeek-R1 avec un seul jeton appris "<|continue-thinking|>", en entraînant uniquement son plongement via l'apprentissage par renforcement tout en gardant les poids du modèle gelés. Nos expériences montrent que ce jeton appris permet d'obtenir une précision améliorée sur des benchmarks mathématiques standards par rapport à la fois au modèle de base et à une approche de redimensionnement au moment du test qui utilise un jeton fixe (par exemple, "Wait") pour forcer le budget. En particulier, nous observons que dans les cas où l'approche à jeton fixe améliore la précision du modèle de base, notre méthode permet une amélioration nettement plus importante. Par exemple, sur le benchmark GSM8K, l'approche à jeton fixe entraîne une amélioration absolue de la précision de 1,3 %, tandis que notre méthode à jeton appris permet une amélioration de 4,2 % par rapport au modèle de base qui n'utilise pas de forçage de budget.

English

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.