PaperlayerView exact revision on arXiv

Source-linked AI summary

Prefix Sliding for efficient test-time scaling

Niklas Muennighoff, Zhengyang Wang, Zeyi Chen, Weijia Shi, Binyuan Hui, John Yang, Dapeng Jiang, Mika Senghaas, Fares Obeid, Johannes Hagemann, Sami Jaghouar, Ludwig Schmidt, Percy Liang, Jason Wei, Andrew Y. Ng, Luke Zettlemoyer, Yejin Choi, Mike Lewis

arXiv:2608.26070v1cs.CLcs.AIcs.LG

TL;DR

Long-horizon test-time scaling is costly because full attention retains the entire reasoning trace and grows more expensive with each generated token. Prefix Sliding retains the prefix and a recent-token window, achieving comparable performance at 3× higher speed without training and supporting reinforcement-learning rollouts beyond 100,000 tokens. The paper also reports limitations involving information loss and restricted comparisons.

  • Problem

    Full attention makes long-horizon reasoning prohibitively expensive because every new token attends to the entire accumulated reasoning trace.

  • Method

    Prefix Sliding retains the instruction prefix and a sliding window of recent reasoning tokens, and applies the approach during reinforcement-learning training.

  • Results

    3× faster without training while matching full-attention performance, with reinforcement-learning rollouts extending beyond 100,000 tokens.

  • Takeaways & Limitations

    Prefix Sliding supports efficient long-horizon reasoning on existing models and can outperform alternative bounded-cost methods.

  • Takeaways & Limitations

    Information loss can require a larger sliding window to match full attention, as observed on LiveCodeBench when code moves outside the window.

Abstract

from arXiv · show

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

1 Introduction

Prefix Sliding addresses the rising cost of long-horizon reasoning by retaining the prefix and recent tokens while discarding older intermediate reasoning. It maintains performance while making existing models 3× faster without training.

  • Full attention makes long-horizon reasoning prohibitively expensive because each new token attends to an increasingly long reasoning trace.The per-token cost grows linearly with the number of previously generated tokens.
  • Intermediate reasoning tokens often lose importance, whereas the prefix and most recent tokens remain important for generation.The prefix carries instructions, tools, and task information; recent tokens capture the model’s current reasoning.
  • Prefix Sliding keeps the prefix and a moving window of recent tokens, capping memory and making generation cost independent of total reasoning length.A 100-token prefix with a 4096-token window keeps at most 4196 tokens in memory.
  • 3× faster without training while matching full-attention performance, with reinforcement-learning rollouts extending beyond 100,000 tokens.These are the reported empirical outcomes of Prefix Sliding.

2 Prefix Sliding

Prefix Sliding retains the prefix while sliding over recent reasoning tokens, with position-embedding and training procedures designed to support efficient long traces. Its training variant uses the limited receptive field of sliding windows to reduce backpropagation cost.

  • Motivation: Prefix Sliding focuses attention on the prefix and recent reasoning tokens because intermediate tokens can become less important during generation.The prefix also functions as an attention sink, while recent tokens represent the model’s current work.
  • Prefix Sliding without training: Prefix Sliding retains the prefix and advances a sliding window as generation proceeds, removing older intermediate tokens from memory.The method preserves the model instructions while bounding the active reasoning context.
  • Prefix Sliding without training: Continue PE reuses cached representations and is more efficient than Reset PE because it avoids reapplying position embeddings to the same token.Both options address position embeddings for models trained with schemes such as RoPE.
  • Prefix Sliding with training: Sliding-window receptive fields allow training to pass only a bounded context around the tokens being backpropagated.The theoretical receptive field is W × L, but practical information flow is closer to 1.5×W.
  • Prefix Sliding with training: Truncated backpropagation can compute loss on a final 2048-token window while passing 8192 preceding tokens and context to the trainer.For a 100,000-token reasoning trace, the earlier 6144 tokens provide context and the final 2048 tokens receive the token-level RL loss.
  • Prefix Sliding kernel implementation: The Prefix Sliding kernel combines elementwise masking for partially overlapping tiles with skipping for tiles outside the prefix or window.This preserves mathematical correctness while avoiding redundant loads and computations.

3 Setup

The study evaluates Prefix Sliding with standard language-model inference and reinforcement-learning setups on reasoning benchmarks, measuring user-facing generation speed and accuracy.

  • Modeling: Experiments primarily use Qwen3-1.7B with vLLM and FlashAttention across sliding-window sizes from 512 to 16,384 tokens.Custom Nvidia Hopper kernels enable Prefix Sliding with FlashAttention.
  • Training: Reinforcement-learning experiments use GRPO with synchronous and asynchronous implementations, backpropagating either entire generations or only the last sliding window.When truncating backpropagation, four times as many recent tokens are passed to the trainer.
  • Evaluation: Evaluation covers GPQA, MATH500, and AIME25, averaging results across 64 runs and using budget forcing for specified thinking lengths.Answers are verified with simpleverify, and generation speed is measured as average thinking time per sample.

4 Results

Prefix Sliding maintains bounded generation cost and can improve efficiency and performance relative to full attention under comparable memory budgets. Its training setup also reduces the context needed for backpropagation while preserving low log-probability mismatch.

  • Prefix Sliding without training: Prefix Sliding’s generation speed stabilizes after the window fills, while full attention continually slows as the attention context grows.The Prefix Sliding kernel reaches approximately regular sliding-window speed, with a small overhead from retaining the prefix.
  • Prefix Sliding without training: Prefix Sliding is faster than full attention in experiments generating 1024 sequences with a 4096-token sliding window on one 80GB H100 GPU.The comparison uses vLLM, FlashAttention, and automatic batch sizing.
  • Prefix Sliding with training: Under near-equal 8,192-token memory budgets, Prefix Sliding can improve performance relative to full attention.The full-attention condition uses 8,192 maximum tokens, while Prefix Sliding uses an 8,192-token window.
  • Prefix Sliding with training: KL divergence is lowest when more context is passed to the trainer, while passing only the final 2K-token window produces a mismatch above 0.1.Passing 4K substantially lowers the mismatch, and 8K is slightly lower and roughly comparable to 16K.
  • Comparisons: AIME25 comparisons use generation lengths up to 262,144 tokens and a 4096-token local window for sliding, restart, and summary methods.The restart and summary alternatives use 256-token retained or summary lengths.

5 Ablations

The ablations compare Prefix Sliding with last-k, summary, and vanilla sliding-window approaches, finding that Prefix Sliding offers the strongest efficiency–performance tradeoff with less overhead.

  • Last k: Last-k methods delete all but the most recent tokens after a threshold, trading speed for token reprocessing when k is large or information loss when k is small.The retained tokens are processed twice when the context changes.
  • Summary: Summary methods periodically compress prior reasoning into a new context, but add summarization overhead, reprocessing, memory volatility, and several hyperparameters.The extra choices include threshold, summary length, prompt, summarizer model, and summary placement.
  • Sliding window: Pure sliding-window attention can lose critical task information once the window is reached, causing performance to flatten during long reasoning.Unlike Prefix Sliding, it does not retain the prefix.
  • Prefix Sliding provides the best performance-efficiency tradeoff among the compared methods.It adds only one hyperparameter: sliding-window size.
  • Prefix Sliding: Prefix Sliding retains the prefix while advancing a window, avoiding the extra summary step and required token reprocessing of competing approaches.Its design preserves task information while bounding the active context.

6 Related Work

Prior test-time scaling methods use sequential or parallel computation, while context-extension methods differ in whether their cost per new token remains bounded as sequences grow.

  • Test-time scaling: Parallel test-time scaling can increase compute by launching more processes, but it faces stark diminishing returns.Majority voting is given as an example of parallel scaling.
  • Context extension: Bounded context-extension methods have asymptotically constant cost per new token, unlike unbounded methods whose per-token cost grows in the limit.Unbounded methods may still have subquadratic complexity.

7 Conclusion

The conclusion presents Prefix Sliding as a way to support extremely long reasoning horizons while improving efficiency over full attention and outperforming alternative bounded-cost approaches.

  • Prefix Sliding is proposed to enable language models to reason for extremely long horizons.The method applies both without further training and during reinforcement-learning training.
  • Prefix Sliding is more efficient than full attention even at reasoning horizons of only thousands of tokens.The conclusion also reports that it outperforms alternatives supporting infinite test-time scaling.
  • The method can be applied to existing language models without further training and can also be used during reinforcement-learning training.

Limitations

The paper’s evidence is bounded by its comparison scope, information-loss cases, short-generation regimes, context-flooding risks, and experiments using 7-billion-parameter models and hundreds of thousands of tokens.

  • Limited comparisons: Comparisons are limited to methods that work out of the box on existing pretrained transformers and have bounded cost per new token.Alternative architectures and several other method classes are excluded.
  • Information loss: LiveCodeBench requires a Prefix Sliding window of at least 16384 tokens to match full attention.The evaluation uses a maximum of 262144 tokens.
  • Limited benefit for short generations: Prefix Sliding benefits short generations less because the window may not slide before generation ends.HealthBench averages 2086 tokens while the evaluated window is 2048, leaving little room for speed-up.
  • System outputs and multi-turn: In agentic tasks, large website or file outputs can flood the context and cause important content to fall outside the sliding window.The paper discusses stepwise reading and output guardrails as possible mitigations.
  • Scale: The experiments scale Prefix Sliding to hundreds of thousands of thinking tokens and 7-billion-parameter models, while further scaling remains future work.

Reproducibility Statement

The paper provides public code and situates Prefix Sliding among context-efficiency methods, with evaluations also considering larger sliding-window sizes.

  • Reproducibility Statement: The authors release code for Prefix Sliding to support reproduction of the key results.The paper states that the method is simple and likely reproducible from its description.
  • Reproducibility Statement: Related work on efficient context handling includes methods targeting the KV cache during prompt processing, generation, or both.
  • Reproducibility Statement: Efficient-thinking research spans train-time, test-time, and combined approaches, including methods targeting reinforcement learning or infrastructure.
  • Reproducibility Statement: Figure 13 evaluates Prefix Sliding with sliding-window sizes of 8192 and 16384.

C Tabular results

The tabular-results section reports comparisons with full attention and includes experiments on reinforcement-learning training and positional-encoding choices.

  • C Tabular results: Table 1 reports that Prefix Sliding achieves performance comparable to full attention at significantly higher speeds.
  • C Tabular results: Figure 15 shows that Prefix Sliding with truncated backpropagation can match full attention.
  • C Tabular results: In a short reinforcement-learning experiment, 32768 tokens reach the trainer while only 8192 Prefix Sliding tokens are backpropagated with a multiplier of 4.The experiment uses a window size of 8192 and finds comparable performance, while noting that more experiments are needed.

F Training Dataset

This section describes dataset filtering and alternative context-management baselines, including fixed-token carryover, summaries, and prompt formats for tool and context use.

  • F Training Dataset: The reinforcement-learning dataset combines public sources including SkyWork and s1.
  • F Training Dataset: Problem filtering uses guessability, verifiability, and difficulty criteria.Samples are removed when small models solve them without thinking, answers are difficult to verify objectively, or models solve them consistently.
  • F Training Dataset: The carryover baseline selects k = 256 tokens from the end of a generation and prepends them to the model’s next generation.The procedure does not preserve sentence boundaries, so carried-over tokens may be cut off.
  • F Training Dataset: Table 2 selects k = 256 because it performs best while keeping the carried context small.The ablation reports avg@64 with a 2048-token context and one pass.
  • F Training Dataset: The summary baseline forces a tool call when k tokens remain, uses the model itself for summarization, and sets k = 256.Prompt variants are ablated, and adding an example of tool and context use raises performance.
  • F Training Dataset: The prompt templates provide function availability, function signatures in XML tags, and examples of using the pass tool.
  • F Training Dataset: Additional figures show prompt formats for framing summaries as tools, inserting context directly, and supplying tool-usage or context-usage examples.
  • F Training Dataset: H2O retains recent tokens and attention-based heavy hitters, whereas Prefix Sliding preserves prefix tokens that may matter later.The paper notes that combining the approaches could be promising but potentially challenging for reinforcement learning.
Loading 2608.26070v1…