향상된 테스트 시간 스케일링을 위한 연속 사고 토큰 학습

초록

테스트 시간 스케일링은 추론 시 추가적인 계산 자원을 활용하여 언어 모델의 성능을 향상시키는 효과적인 접근법으로 부상했다. 최근 연구들은 사고 종료 토큰을 재정의함으로써(예: "</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.

향상된 테스트 시간 스케일링을 위한 연속 사고 토큰 학습

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

초록

Support