STS: Efficient Sparse Attention with Speculative Token Sparsity

2026-05-15

AI Summary

This work addresses the memory and computational bottlenecks imposed by the quadratic complexity of attention mechanisms in large language models during long-sequence reasoning. The authors propose STS, a sparse attention mechanism that leverages attention scores from a small draft model to predict critical tokens for a larger target model. Integrated within a speculative decoding framework, STS dynamically generates token-wise and head-wise sparse masks without requiring any retraining, achieving approximately 90% sparsity. Evaluated on the NarrativeQA benchmark, the method attains a 2.67× speedup with negligible accuracy degradation, substantially outperforming existing sparse attention approaches and establishing a new state of the art in the sparsity–accuracy trade-off.

Abstract

The quadratic complexity of attention imposes severe memory and computational bottlenecks on Large Language Model (LLM) inference. This challenge is particularly acute for emerging agentic applications that require processing multi-million token sequences. We propose STS, a sparse attention mechanism that requires no model retraining. STS leverages the key insight that tokens identified as important by a smaller draft model are highly predictive of important tokens for a larger target model. By integrating into speculative decoding frameworks, STS repurposes the draft model's attention scores to dynamically construct a token-and-head-wise sparsity mask. This mask effectively prunes the expensive attention computation in the target LLM. Our evaluation shows that STS achieves a 2.67x speedup operating at approximately 90% sparsity on representative benchmark NarrativeQA, maintaining negligible accuracy degradation compared to dense attention. STS establishes a new state-of-the-art on the sparsity-accuracy trade-off, outperforming prior techniques by enabling higher sparsity levels for a given accuracy budget.

Problem

• sparse attention
• Large Language Model
• quadratic complexity
• inference bottleneck
• long-context

Innovation

• sparse attention
• speculative decoding
• token sparsity
• large language models
• attention pruning