Aprendizaje de un Token de Pensamiento Continuo para una Escalabilidad Mejorada en el Momento de la Prueba

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

June 12, 2025

Autores: Liran Ringel, Elad Tolochinsky, Yaniv Romano

cs.AI

Resumen

El escalado en tiempo de prueba ha surgido como un enfoque efectivo para mejorar el rendimiento de los modelos de lenguaje al utilizar recursos computacionales adicionales durante la inferencia. Estudios recientes han demostrado que anular los tokens de fin de razonamiento (por ejemplo, reemplazar "</think>" con "Wait") puede extender los pasos de razonamiento y mejorar la precisión. En este trabajo, exploramos si se puede aprender un token dedicado de "continuar pensando" para desencadenar un razonamiento extendido. Aumentamos una versión destilada de DeepSeek-R1 con un único token aprendido "<|continuar-pensando|>", entrenando únicamente su incrustación mediante aprendizaje por refuerzo mientras mantenemos congelados los pesos del modelo. Nuestros experimentos muestran que este token aprendido logra una mayor precisión en benchmarks matemáticos estándar en comparación tanto con el modelo base como con un enfoque de escalado en tiempo de prueba que utiliza un token fijo (por ejemplo, "Wait") para forzar el presupuesto. En particular, observamos que en los casos en los que el enfoque de token fijo mejora la precisión del modelo base, nuestro método logra una mejora notablemente mayor. Por ejemplo, en el benchmark GSM8K, el enfoque de token fijo produce una mejora absoluta del 1.3% en la precisión, mientras que nuestro método de token aprendido alcanza una mejora del 4.2% sobre el modelo base que no utiliza forzado de presupuesto.

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.