Обучение токена непрерывного мышления для улучшенного масштабирования во время тестирования

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

June 12, 2025

Авторы: Liran Ringel, Elad Tolochinsky, Yaniv Romano

cs.AI

Аннотация

Масштабирование на этапе тестирования стало эффективным подходом для повышения производительности языковых моделей за счет использования дополнительных вычислительных ресурсов на этапе вывода. Недавние исследования показали, что переопределение токенов завершения мышления (например, замена "</think>" на "Wait") может увеличить количество шагов рассуждения и повысить точность. В данной работе мы исследуем, можно ли обучить специальный токен продолжения мышления для запуска расширенного рассуждения. Мы дополняем дистиллированную версию модели DeepSeek-R1 единственным обученным токеном "<|continue-thinking|>", обучая только его встраивание с помощью обучения с подкреплением, при этом оставляя веса модели замороженными. Наши эксперименты показывают, что этот обученный токен достигает улучшенной точности на стандартных математических тестах по сравнению как с базовой моделью, так и с подходом масштабирования на этапе тестирования, использующим фиксированный токен (например, "Wait") для принудительного ограничения бюджета. В частности, мы наблюдаем, что в случаях, когда подход с фиксированным токеном повышает точность базовой модели, наш метод достигает значительно большего улучшения. Например, на тесте GSM8K подход с фиксированным токеном дает абсолютное улучшение точности на 1,3%, тогда как наш метод с обученным токеном достигает улучшения на 4,2% по сравнению с базовой моделью, не использующей принудительное ограничение бюджета.

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.