Het leren van een door-denken-token voor verbeterde schaling tijdens testtijd

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

June 12, 2025

Auteurs: Liran Ringel, Elad Tolochinsky, Yaniv Romano

cs.AI

Samenvatting

Test-time schaling is naar voren gekomen als een effectieve aanpak om de prestaties van taalmodelen te verbeteren door extra rekenkracht te benutten tijdens de inferentie. Recente studies hebben aangetoond dat het overschrijven van eind-van-denken-tokens (bijvoorbeeld het vervangen van "</think>" door "Wacht") de redeneerstappen kan verlengen en de nauwkeurigheid kan verbeteren. In dit werk onderzoeken we of een speciaal voort-denken-token kan worden aangeleerd om uitgebreid redeneren te activeren. We breiden een gedistilleerde versie van DeepSeek-R1 uit met een enkel aangeleerd "<|continue-thinking|>" token, waarbij we alleen de embedding ervan trainen via reinforcement learning terwijl de modelgewichten bevroren blijven. Onze experimenten tonen aan dat dit aangeleerde token een verbeterde nauwkeurigheid bereikt op standaard wiskundige benchmarks in vergelijking met zowel het basismodel als een test-time schaling-benadering die een vast token (bijvoorbeeld "Wacht") gebruikt voor budgetforcing. In het bijzonder observeren we dat in gevallen waar de vaste-token-benadering de nauwkeurigheid van het basismodel verbetert, onze methode een aanzienlijk grotere verbetering bereikt. Op de GSM8K benchmark levert de vaste-token-benadering bijvoorbeeld een absolute verbetering van 1,3% in nauwkeurigheid op, terwijl onze aangeleerde-token-methode een verbetering van 4,2% bereikt ten opzichte van het basismodel dat geen budgetforcing gebruikt.

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.