Attention Routing Stabilizes Early: Working-Set Inference for Recurrent Language Models

Ke Wan, Chen Chen
9/23/2026
cs.CLcs.LG

Official code: tbn5pj/WISE_code (3 stars, MIT, last commit 2026-09-24).

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.

View on arXivView PDF

Code Implementations(1)

tbn5pj/WISE_codeOfficial100%
MIT
30PythonSep 23, 20263 weeks agoMIT

Similar Papers

Cite this paper

@article{wan2026attention,
  title  = {Attention Routing Stabilizes Early: Working-Set Inference for Recurrent Language Models},
  author = {Ke Wan and Chen Chen},
  year   = {2026},
  eprint = {2609.27373},
  archivePrefix = {arXiv},
  url    = {http://arxiv.org/abs/2609.27373v1}
}

Add a badge to your README

Shows whether official code is linked for this paper.

Code of Paper badge
[![Code of Paper](https://codeofpaper.com/badge/2609.27373.svg)](https://codeofpaper.com/paper/2609.27373)

Discussion