Aprendendo um Token de Pensamento Contínuo para Melhorar a Escalabilidade em Tempo de Teste

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

June 12, 2025

Autores: Liran Ringel, Elad Tolochinsky, Yaniv Romano

cs.AI

Resumo

A escala em tempo de teste surgiu como uma abordagem eficaz para melhorar o desempenho de modelos de linguagem ao utilizar computação adicional durante a inferência. Estudos recentes demonstraram que substituir tokens de fim de pensamento (por exemplo, trocar "</think>" por "Wait") pode estender os passos de raciocínio e melhorar a precisão. Neste trabalho, exploramos se um token dedicado para continuar o pensamento pode ser aprendido para desencadear um raciocínio estendido. Aumentamos uma versão destilada do DeepSeek-R1 com um único token aprendido "<|continue-thinking|>", treinando apenas sua incorporação por meio de aprendizado por reforço enquanto mantemos os pesos do modelo congelados. Nossos experimentos mostram que esse token aprendido alcança uma precisão melhorada em benchmarks matemáticos padrão em comparação tanto com o modelo base quanto com uma abordagem de escala em tempo de teste que usa um token fixo (por exemplo, "Wait") para forçar o orçamento. Em particular, observamos que, nos casos em que a abordagem de token fixo melhora a precisão do modelo base, nosso método alcança uma melhoria significativamente maior. Por exemplo, no benchmark GSM8K, a abordagem de token fixo resulta em uma melhoria absoluta de 1,3% na precisão, enquanto nosso método de token aprendido alcança uma melhoria de 4,2% sobre o modelo base que não utiliza forçamento de orçamento.

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.