PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Ke Wan, Chen Chen

arXiv:2609.27373v2cs.CLcs.LG

TL;DR

Recurrent-depth inference repeatedly recomputes global attention even as routing may stabilize before representations finish refining. WISE discovers a sparse working set during early unrestricted attention and reuses its support while keeping refinement dynamic. Across multi-hop QA benchmarks, it largely preserves full-attention behavior and achieves up to a 1.76× late-step attention speedup over native FlashAttention at 4K context.

  • Problem

    Recurrent inference repeats full-context global attention at every step, and it is unclear whether later steps still need unrestricted routing after earlier steps identify relevant context.

    Source: §1 · pp. 1–2

  • Method

    WISE uses early unrestricted attention to discover a sparse block-structured working set, then reuses only its routing support while recurrent representations and within-support attention remain dynamic.

    Source: §1 · p. 2

  • Results

    Across multi-hop QA benchmarks, WISE largely preserves full-attention behavior and achieves up to a 1.76× late-step attention speedup over native FlashAttention at 4K context.

    Source: §1 · p. 3

  • Takeaways & Limitations

    Routing support can become computationally reusable while late representation refinement remains valuable, with quality–efficiency tradeoffs improving as support becomes sparser at longer contexts.

    Source: §1 · p. 3

    The theoretical analysis explains finite-depth emergence of reusable routing structure but does not predict the empirical switching depth td.

    Source: §6 · p. 6

Abstract

from arXiv · show

Recurrent-depth language models, such as looped Transformers, repeatedly apply shared network blocks to refine latent representations without generating explicit intermediate reasoning tokens. However, each step recomputes full attention over the entire context, repeating costly global routing. We study how attention routing evolves across recurrent depth and find a consistent separation in convergence timescales: attention support and distributions stabilize substantially earlier than hidden states and attention outputs. This suggests two stages of recurrent inference: early discovery of a sparse working set, followed by representation refinement over largely stable routing support. Motivated by this finding, we introduce WISE (Working-set Inference with Support Exploitation), a training-free method that uses unrestricted attention during early recurrent steps to discover a block-structured working set, then reuses its support in later steps while keeping attention weights and recurrent refinement dynamic. Controlled interventions show that multi-step discovery yields more effective working sets than first-step selection, and that support reuse better preserves model behavior than more restrictive forms of attention reuse. Across multi-hop QA benchmarks, WISE largely preserves full-attention performance. Matched context-scaling experiments reveal an increasingly favorable quality-efficiency tradeoff as routing support becomes sparser with longer contexts. A sparse-attention implementation achieves up to a 1.76x late-step attention speedup over native FlashAttention at 4K context. Code: https://github.com/tbn5pj/WISE_code.

1 INTRODUCTION

Recurrent-depth models repeatedly refine latent representations but recompute global attention at every step. The paper finds that routing stabilizes before representations, motivating WISE, which reuses discovered routing support while preserving dynamic refinement and largely maintaining full-attention behavior across multi-hop QA benchmarks.

1 INTRODUCTION

  • Existing recurrent inference recomputes full-context attention at every step, repeating quadratic-cost global routing while it remains unclear whether later steps need to search the full context.Recurrent depth provides additional test-time computation without generating explicit intermediate reasoning tokens.Source: §1 · pp. 1–2
  • Attention support and distributions stabilize substantially earlier than hidden states and attention outputs across benchmarks, convergence thresholds, and two recurrent-depth models.Which of support or distribution stabilizes first varies across tasks and models.Source: §1 · pp. 2–3
  • WISE uses unrestricted attention in early recurrent steps to discover a sparse block-structured working set, then reuses only its routing support while within-support computation remains dynamic.It is training-free and preserves recurrent depth, hidden representations, queries, keys, values, and attention weights as they evolve.Source: §1 · pp. 2–3
  • WISE largely preserves full-attention behavior across multi-hop QA benchmarks, with up to a 1.76× late-step attention speedup over native FlashAttention at 4K context.Matched context-scaling experiments show an increasingly favorable quality–efficiency tradeoff as routing support becomes sparser with longer contexts.Source: §1 · p. 3
  • Multi-step working-set discovery is more effective than first-step selection, and support-only reuse preserves model behavior better than more restrictive attention reuse.The paper also reports that delayed discovery is most useful when early routing continues to reorganize, and that quality is largely preserved through 2K context.Source: §1 · p. 3

2 RELATED WORK

Prior work studies recurrent refinement for test-time computation, adaptive allocation of computation, and efficient or sparse attention. WISE instead keeps recurrent depth and dynamic within-support computation while reusing routing support in recurrent-depth models.

Recurrent and Latent Test-Time Computation

  • Recurrent architectures repeatedly apply shared blocks to increase effective depth, and recent language-model work uses recurrent depth for additional test-time computation and multi-step reasoning.WISE shifts the question to whether every recurrent step needs unrestricted global routing.Source: §2 · p. 3

Adaptive Transformer Computation

  • Adaptive computation methods reduce how much computation inputs or tokens receive, whereas WISE preserves recurrent depth and narrows only late global routing.Attention weights and hidden-state updates remain dynamic.Source: §2 · p. 3

Efficient and Sparse Attention

  • Sparse-attention and attention-reuse methods address efficient attention, while WISE reuses routing support after stabilization and keeps within-support attention and representations dynamic.The paper describes this support-level sparsity as complementary to efficient kernels such as FlashAttention.Source: §2 · p. 3

3 PROBLEM SETUP

The setup represents recurrent inference as repeated application of a shared transformation to a sequence of hidden states. At each step, attention is recomputed from the current representations, and working sets describe sparse routing supports for later reuse.

3 PROBLEM SETUP

  • A sequence of length n has hidden states H(t) in R^(n×d) after recurrent step t, with shared parameters θ applied over recurrent depth T.The formulation abstracts away model-specific components outside the recurrent core.Source: §3 · p. 4
  • At each recurrent step, queries, keys, and values are projected from current hidden states, and attention produces an attention matrix and output.The setup defines Q(t), K(t), and V(t) using shared projection matrices.Source: §3 · p. 4
  • For each query, the attention support is the smallest admissible key set whose cumulative attention mass reaches threshold η.The setup also defines a working set as sparse routing support discovered during early recurrent computation and reused later.Source: §3 · p. 4
  • Standard recurrent inference recomputes global attention over all admissible query–key pairs at every recurrent step, even after routing structure has stabilized.For structured execution, the setup denotes active key-block support for each query block.Source: §3 · p. 4

4 COMPUTATIONAL STRUCTURE OF RECURRENT INFERENCE

Repeated full attention is the context-quadratic component of recurrent inference. If early steps identify a sparse working set, later routing can depend on that set while recurrent representations and attention computations continue to evolve.

4 COMPUTATIONAL STRUCTURE OF RECURRENT INFERENCE

  • Self-attention costs O(n^2d) per step, while linear projections and feed-forward layers scale as O(nd^2), making repeated global routing the context-quadratic component.Over T recurrent steps, standard inference repeatedly incurs this attention cost.Source: §4 · p. 4
  • WISE exploits earlier routing stabilization by discovering globally early and reusing the working set late, while Q(t), K(t), V(t), and within-support attention weights continue evolving.This targets repeated global routing rather than recurrent representation refinement.Source: §4 · p. 4
  • When early computation identifies an average of m ≪ n relevant keys per query, late attention routing can reduce from O(n^2d) to O(nmd).With working-set density ρ = m/n, the late-stage cost is O(ρn^2d); it is subquadratic if m grows more slowly than n and linear when m is bounded.Source: §4 · p. 4

5 WISE: DISCOVER EARLY, REUSE LATE

WISE discovers a block-structured routing support during early unrestricted attention, then reuses it while recurrent representations and within-support attention remain dynamic. This changes late routing cost to depend on the working-set size rather than the full context.

5 WISE: DISCOVER EARLY, REUSE LATE

  • WISE separates inference into unrestricted support discovery and later reuse of a block-structured working set, trading some fine-grained sparsity for GPU-executable regular computation.The support-reuse principle is agnostic to routing granularity, but the method instantiates it with blocks for efficient sparse execution.Source: §5 · p. 5

5.1 BLOCK-STRUCTURED WORKING-SET DISCOVERY

  • During discovery, WISE selects the smallest set of admissible key blocks whose cumulative mass reaches η and unions direct supports across recurrent steps.The default uses T = 32, td = 12, B = 32, and η = 0.95, combining supports from steps 9–12 independently by recurrent block, head, and query block.Source: §5.1 · p. 5

5.2 WORKING-SET REUSE

  • After discovery, the block support stays fixed, while hidden representations, queries, keys, values, and within-support attention weights continue to change.WISE fixes where attention may route, not how retained information is weighted and transformed.Source: §5.2 · p. 5

5.3 COMPUTATIONAL COMPLEXITY

  • Late-stage routing cost changes from O(n^2d) to O(nmd), with larger gains when the working set is sparser relative to context.Here m is the average number of active key positions per query; the stated comparison ignores block-boundary and kernel overhead.Source: §5.3 · pp. 5–6

6 THEORETICAL ANALYSIS

The analysis formalizes why block-level routing support can stabilize at finite recurrent depth while representations continue to evolve. It also distinguishes support stabilization from representation convergence and does not predict the empirical switching depth.

6 THEORETICAL ANALYSIS

  • The analysis formalizes finite-depth stabilization of the block-level cumulative-mass support, but does not predict WISE’s empirical switching depth.Under the stated convergence and nondegeneracy conditions, supports become stable after finite depth; a temporal-union window beyond that depth equals the limiting support, while an earlier window buffers local routing fluctuations.Source: §6 · p. 6, §6.1 · p. 6

6.1 FINITE-TIME IDENTIFICATION OF THE WORKING SET

  • With a positive margin from ranking and cumulative-mass boundaries, continuous changes cannot alter the selected block support after finite depth.Taking the maximum identification time across the finite set of recurrent blocks, heads, and query blocks yields simultaneous stabilization of all nondegenerate supports.Source: §6.1 · p. 6

6.2 SUPPORT STABILIZATION AND REPRESENTATION CONVERGENCE

  • Finite-time support stabilization does not imply representation convergence: recurrent states can keep evolving within the same routing region.WISE can therefore reuse the block support while hidden states and within-support attention continue to change; recurrence need not stop when support stabilizes.Source: §6.2 · p. 7

7 EXPERIMENTS

Experiments find that attention routing stabilizes before representation refinement, supporting WISE’s reuse of discovered support while recurrent computation continues. On Huginn, context scaling shows increasing late-reuse speedups alongside a measurable quality cost at 4K.

7.1 EXPERIMENTAL SETUP

  • The primary experiments use Huginn at recurrent depth T = 32, with Recurrent-Llama-T32 as a replication; benchmark comparisons use matched 100-example cohorts and mechanism diagnostics use fixed 30-example cohorts.WISE uses B = 32, discovery depth 12, mass threshold 0.95, and a four-step discovery window spanning steps 9–12, without benchmark-specific tuning.Source: §7.1 · p. 7

7.2 MECHANISM DISCOVERY

  • Routing-related quantities stabilize earlier than hidden states and attention outputs across HotpotQA and GSM8K, and the separation persists across criteria and on Recurrent-Llama-T32.The deployed steps-9–12 block-support union enters a high-stability regime near discovery depth 12.Source: §7.2 · pp. 7–8

7.3 CAUSAL VALIDATION OF SUPPORT-ONLY REUSE

  • Controlled interventions indicate that continuing recurrence matters and that reusing only routing support preserves Full-model behavior better than freezing attention distributions or using early-static controls.FREEZE-A@12 shares WISE’s support but freezes within-support attention weights; truncating recurrence at step 12 worsens teacher-forced gold-answer loss.Source: §7.3 · p. 8

7.5 QUALITY–EFFICIENCY SCALING WITH CONTEXT LENGTH

  • 1.76× setup-inclusive late-reuse speedup over native FlashAttention at 4K accompanies no detected average F1 loss through 2K, but a 3.3-point reduction at 4K.Working-set density falls from roughly 57% to 37% across 512-to-4K contexts while retaining about 96–98% of future Full-attention mass.Source: Figure 2 · p. 9, §7.5 · p. 9

7.4 PRACTICAL GPU EFFICIENCY

  • 1.36× complete T = 32 attention-trajectory speedup over native FlashAttention remains after including unrestricted discovery, while a matched exact-B = 32 backend comparison yields 2.50× speedup.The implementation uses a Triton kernel that skips inactive key blocks while preserving WISE support and attention semantics.Source: §7.4 · pp. 8–9, Table 2 · p. 9

7.5 QUALITY–EFFICIENCY SCALING WITH CONTEXT LENGTH

  • B = 16 yields the sparsest support, but B = 32 achieves the highest setup-inclusive speedup at 2K and 4K; no tested granularity consistently dominates quality across both benchmarks.The authors use B = 32 as the systems-facing knee between sparsity and executable regularity.Source: §7.5 · p. 9

8 CONCLUSION

The conclusion presents WISE as reusing early-discovered sparse routing support while preserving recurrent refinement. The accompanying proof establishes eventual support stabilization under stated convergence and ranking conditions.

8 CONCLUSION

  • WISE preserves recurrent depth while reusing sparse routing support discovered during early global attention, and the conclusion identifies kernel co-design as a possible route to further efficiency.Across two recurrent backbones, results support routing–representation separation; delayed discovery is most useful when early routing remains predictive of later attention only after several steps.Source: §8 · p. 10
  • The conclusion frames support reuse as compatible with ongoing representation refinement rather than as a replacement for recurrent computation.It also describes WISE’s block-sparse routing as complementary to FlashAttention-style kernel optimization.Source: §8 · p. 10

A.1 PROOF OF THEOREM 6.1

  • For a finite collection of routing units, the proof gives a finite recurrent depth after which all nondegenerate block supports are simultaneously identified, under its local assumptions.The argument uses convergence of recurrent representations and a strict ranking boundary for the limiting block-mass support.Source: §A.1 · pp. 13–14

A.2 PROOF OF PROPOSITION 1

The proposition proof distinguishes convergence of representations from stabilization of routing support: representation convergence entails eventual support stabilization, but support stability alone does not entail representation convergence.

A.2 PROOF OF PROPOSITION 1

  • Convergence to a nondegenerate limiting representation implies eventual stabilization of the block support.The proof states this consequence for the routing support associated with query block q.Source: §A.2 · p. 14
  • Support stabilization does not imply representation convergence, because distinct representations can lie in a region that induces the same block support.Without additional constraints on recurrent dynamics, the trajectory may continue moving while preserving identical routing support.Source: §A.2 · p. 14

B EXPERIMENTAL DETAILS

The experimental details define matched evaluation cohorts, a shared WISE configuration, intervention controls, and separate diagnostics for routing dynamics. Context-quality scaling and GPU timing use different cohorts and settings.

B EXPERIMENTAL DETAILS

  • Experiments use Huginn at T = 32 and Recurrent-Llama-T32, with fixed matched inputs and recurrent initialization; behavioral comparisons use 100-example cohorts and mechanism diagnostics use 30-example cohorts.Recurrent-Llama uses a fixed plain-completion prompt, and no model training or parameter update is performed.Source: B EXPERIMENTAL DETAILS · pp. 14, 16
  • WISE uses B = 32, discovery depth td = 12, mass threshold η = 0.95, and the union of direct block-level 95%-mass supports from steps 9–12.After discovery, the support remains fixed while hidden states, queries, keys, values, and within-support attention weights continue to be recomputed.Source: B EXPERIMENTAL DETAILS · p. 15
  • Intervention controls compare WISE with three early-static controls using the same block representation and intervention timing; STATIC-95 freezes the direct 95%-mass support identified at the first recurrent step.Source: B EXPERIMENTAL DETAILS · p. 15
  • Evaluation reports EM, token-level F1, paired F1 differences, answer changes, adjusted-loss changes, block density, and retained future attention mass.Retained future attention mass measures the fraction of unrestricted Full-attention mass after discovery that falls inside the frozen support.Source: B EXPERIMENTAL DETAILS · p. 16
  • Mechanism diagnostics track token support, attention distributions, hidden representations, and pre-projection attention outputs, using fixed thresholds across both backbones.These diagnostics characterize dynamics but do not determine WISE’s switching depth or working-set configuration; the deployed block support is analyzed separately.Source: B EXPERIMENTAL DETAILS · p. 16
  • Quality scaling uses the same 100-example HotpotQA cohort at approximately 512, 1K, 2K, and 4K, while systems timing uses a distinct 30-example Huginn cohort on RTX A6000 GPUs.The checkpoint supports at most 4096 positions, so the study reports no model-level 8K quality result; timing compares setup-inclusive reuse, the full T = 32 trajectory, and a same-backend support comparison.Source: B EXPERIMENTAL DETAILS · pp. 16–17

C.1 SENSITIVITY TO CONVERGENCE CRITERIA

Across matched threshold settings, routing stabilizes before representations, and the deployed working set becomes highly stable near its discovery depth. Cross-backbone results suggest delayed discovery is most useful when early routing continues to reorganize.

C.1 SENSITIVITY TO CONVERGENCE CRITERIA

  • Across five matched thresholds, G stays positive in 30/30 examples on both tasks, with mean gaps of 4.33–6.67 steps on HotpotQA and 7.80–17.50 on GSM8K.Under deliberately asymmetric strict-routing/loose-representation calibration, the mean HotpotQA gap reverses to G = −1.57; the ordering is not invariant to arbitrary cross-quantity calibration.Source: §C.1 · pp. 17–18, Figure 3 · p. 18

C.2 STABILITY OF THE DEPLOYED WORKING SET

  • At t = 12, mean adjacent rolling-set Jaccard is 0.9674 and all 30 examples meet the 0.90 criterion; all meet 0.95 by t = 13, placing td = 12 at the onset of high stability.WISE freezes the working set after discovery; this is not an exact convergence boundary, and slow cumulative drift can remain.Source: §C.2 · p. 19, Figure 4 · p. 19

C.3 CROSS-BACKBONE REPLICATION

  • On Recurrent-Llama-T32, all 30 examples on both tasks show positive routing–representation gaps, preserving the separation despite support stabilizing before attention distributions.Attention-output convergence on GSM8K is right-censored: 27/30 trajectories do not converge by T = 32.Source: §C.3 · pp. 19–20, Table 4 · p. 20

C.4 WHEN DOES DELAYED DISCOVERY HELP?

  • Recurrent discovery adds about 2 percentage points of late-mass coverage on Huginn but only about 0.4 on Recurrent-Llama-T32, where early-static support already retains 96.13% and 96.41%.The authors suggest delayed discovery helps most when routing reorganizes during early recurrence, while noting that saved trajectories cannot isolate routing stability from broader differences in recurrent dynamics.Source: §C.4 · p. 20

C.5 WISE HYPERPARAMETER SENSITIVITY

Local parameter sweeps show tradeoffs rather than a universally optimal WISE setting. B = 32 provides the strongest tested execution speed, while the default threshold balances sparsity and retained attention mass.

C.5 WISE HYPERPARAMETER SENSITIVITY

  • No tested neighboring configuration clearly dominates the frozen default on the native-context HotpotQA cohort; the sweep supports local robustness, not hyperparameter insensitivity or optimality.The analysis varies one parameter at a time and does not establish robustness to long contexts or joint parameter interactions.Source: §C.5 · pp. 20–21, Table 7 · p. 21

C.6 BLOCK-SIZE TRADEOFF

  • 1.769× at 4K: B = 32 reaches the highest speedup, while the sweep identifies no block size that consistently dominates downstream F1 across both benchmarks.The selected B = 32 operating point balances routing sparsity with executable regularity; finer sparsity alone does not guarantee faster execution.Source: §C.6 · p. 22, Table 8 · p. 22

C.7 SENSITIVITY TO THE WORKING-SET MASS THRESHOLD

  • η = 0.90 yields sparser working sets but lower retained future mass and more behavioral change; η = 0.99 retains nearly all future mass at roughly 80% density without consistent quality gains.The authors describe η = 0.95 as a reasonable quality–sparsity operating point, not a narrowly tuned F1 optimum.Source: §C.7 · p. 22, Table 9 · p. 22

Related papers

Loading 2609.27373v2…