Prefix sliding for efficient test-time scaling
Abstract
Test-time scaling uses extra test-time compute to improve performance, such as letting language models reason longer when solving a problem. As models keep the entire reasoning trace in memory via full attention, it is prohibitively expensive for hard tasks that need long thinking. However, we find most intermediate reasoning tokens lose their importance quickly. This calls into question whether retaining them is worth the cost. Based on this insight, we propose Prefix Sliding which discards tokens during reasoning that are not part of the prefix or the window of the last few thousand tokens. The prefix has key information about instructions to the model and tools available to it, while the most recent tokens are the current reasoning trace the model is working on. This caps the total memory requirement regardless of how long the model reasons, thus allowing for efficient long-horizon test-time scaling. Without training, Prefix Sliding can make existing models 3x faster while maintaining performance. Training with Prefix Sliding using reinforcement learning allows reaching better performance under the same memory usage by scaling to reasoning traces beyond 100,000 tokens. Ablation studies show Prefix Sliding is more effective than summarizing intermediate tokens or vanilla sliding window. We seek to make our code and data publicly available.