Structured-Sparse Attention for Entity Tracking with Subquadratic Sequence Complexity

2026-05-21

AI Summary

This work addresses the challenges of efficiently maintaining entity states over long sequences and the high computational cost of dense attention mechanisms. To this end, the authors propose a block-sparse attention scheme that exploits the block-diagonal dominant structure inherent in attention distributions for entity tracking tasks. By introducing a resolvent-based block approximation method, the approach enables multi-hop state propagation within a single layer. Notably, it achieves sub-quadratic sequence complexity—specifically O(n^{4/3}d)—while preserving accuracy on par with dense attention. Experimental results demonstrate that the method reduces runtime by 12–29% compared to standard dense Transformers, with peak speedups reaching 2.4×.

Abstract

Entity tracking requires maintaining and updating latent states for entities and attributes over long sequences. Recent task-specific attention operators can compress deep Transformer stacks into a few layers by performing multi-hop state propagation within a single layer, but their dense evaluation remains expensive. We show that in this setting, learned attention is strongly structured: most mass concentrates in local block-diagonal neighborhoods with a light cross-block residue. Exploiting this, we derive a blockwise evaluation of a resolvent-style operator that keeps within-block interactions exact and routes cross-block interactions through a reduced system. The resulting evaluation is subquadratic in sequence length $O(n^{4/3}d)$ (and $O(n^{7/3})$ when $d\approx n$). On controlled tracking benchmarks, our method matches the dense operator's accuracy while reducing wall-clock time by $12-29\%$ under a standardized measurement protocol, and is up to $2.4 \times$ faster than a compact dense Transformer at comparable exact-match accuracy. We further provide ablations over block size and model capacity, and identify a limitation: performance collapses when the number of simultaneously evolving properties exceeds the number of attention heads.

Problem

• entity tracking
• structured-sparse attention
• subquadratic complexity
• long sequences
• attention efficiency

Innovation

• structured-sparse attention
• subquadratic complexity
• entity tracking
• blockwise evaluation
• resolvent-style operator