Search papers, labs, and topics across Lattice.
The paper introduces Prefix Sliding, a method that optimizes test-time scaling in language models by selectively discarding less important intermediate reasoning tokens, thereby reducing memory usage while maintaining performance. This approach leverages the observation that many intermediate tokens lose significance as reasoning progresses, allowing models to operate efficiently even with extended reasoning traces. Experimental results demonstrate that Prefix Sliding can enhance existing models' speed by three times without retraining and can further improve performance when combined with reinforcement learning for long-horizon reasoning tasks.
Discarding irrelevant reasoning tokens can make language models three times faster at test time without sacrificing performance.
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, hard tasks that need long thinking can be prohibitively expensive. However, we find most intermediate reasoning tokens lose importance as the model continues reasoning. 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 instructions and tools available to the model, while the most recent tokens are the current reasoning the model is working on. This caps the total memory requirement regardless of how long the model reasons, 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 can achieve better performance by enabling scaling to reasoning traces beyond a hundred thousand tokens. Ablations show Prefix Sliding outperforms summarizing intermediate tokens or vanilla sliding window. Our code is at https://github.com/Muennighoff/prefix-sliding