Attention routing method speeds up long context language model inference
Attention Routing Stabilizes Early: Working-Set Inference for Recurrent Language Models
Computation and LanguageMachine Learning
Summary
Recurrent language models process text repeatedly, but usually recalculate attention over the whole context each time, which is slow. The authors found that early on in processing, the model focuses on a small, stable set of important context parts. They designed WISE, a method that uses full attention early but then reuses this smaller attention focus later, speeding up processing without much loss in quality. This approach works well especially for tasks like question-answering and scales efficiently to very long text inputs.
What this means in practice
- •For natural language processing engineers: Speed up inference in recurrent language models on long text contexts by reusing discovered attention supports to reduce computation.
- •For document retrieval teams: Improve efficiency of multi-hop question answering systems by maintaining high accuracy while processing very long context efficiently.
Authors
Ke Wan, Chen Chen
Abstract
Recurrent language models repeatedly apply shared network blocks to refine latent representations, but standard inference recomputes global attention at every recurrent step. We study attention dynamics across recurrent depth and find that attention support and distributions stabilize substantially earlier than hidden states and attention outputs. This suggests a two-stage structure: early steps discover a sparse working set of relevant context, while later steps refine representations over largely the same routing support. Motivated by this structure, we introduce WISE (Working-set Inference with Support Exploitation), a training-free method that uses unrestricted global attention during early recurrence and later reuses directly discovered block-structured support while keeping recurrent depth and within-support attention computation dynamic. Controlled interventions show that recurrent discovery is important and that support-only reuse better preserves model behavior than more restrictive attention-reuse alternatives. Across multi-hop QA benchmarks, WISE largely preserves full-attention performance, while context scaling reveals increasingly sparse working sets and greater efficiency gains. Quality is largely preserved through 2K context, with a measurable loss at 4K. An optimized sparse-attention implementation achieves up to a 1.76x attention speedup over native FlashAttention at 4K and a 1.36x speedup for the full 32-step attention trajectory. Our code is available at https://github.com/tbn5pj/WISE_code.