首页 > AI前沿 > Attention Routing Stabilizes Early: Working-Set Inference for Recurrent Language Models

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

arXiv自然语言 2026-09-23 13:21 6 阅读 查看原文

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.