STS: Sparse Attention Mechanism Combined with Speculative Decoding, Achieving 2.67x Speedup at 90% Sparsity

STS (Sparse Attention Mechanism Combined with Speculative Decoding) dynamically constructs sparse masks using attention scores from a draft model. Without retraining large language models, it achieves 90% sparsity and a 2.67x inference speedup, providing a new solution to break through the quadratic complexity bottleneck of attention in long-context processing.

The computational complexity of self-attention mechanisms in large language models is proportional to the square of the sequence length, which becomes a memory and computational bottleneck when processing long texts. Traditional sparse attention schemes require retraining or fixed patterns and cannot dynamically adapt to inputs, so there is an urgent need for an efficient, flexible sparse attention mechanism that does not require retraining.

The key insight of STS is that important tokens identified by small draft models are also predictive for large target models. It innovatively extends the speculative decoding framework, reusing attention scores from the draft model to guide the target model to focus on tokens. Its advantages include no need for retraining, dynamic adaptation, and computational reuse.

Steps include: 1. Extract and aggregate cross-head and cross-layer attention scores from the draft model to obtain the global importance of tokens; 2. Construct dynamic sparse masks based on scores (supporting head-level sparsification); 3. Seamlessly integrate with speculative decoding, hiding the mask construction overhead in the draft generation phase.

In the NarrativeQA benchmark test: 2.67x speedup (vs dense attention), approximately 90% sparsity, and negligible accuracy loss; outperforms existing methods in the sparsity-accuracy trade-off; the speedup effect becomes more significant as the sequence length increases, making it suitable for ultra-long sequence scenarios.

Suitable scenarios: 1. Agents and long-context applications (code assistants, research assistants, multimodal understanding); 2. Edge device deployment (reducing energy consumption and latency); 3. Real-time interactive systems (improving user experience).

Limitations: Dependence on the quality of the draft model, task specificity (needs verification on more tasks), and storage overhead of dynamic masks. Future directions: Optimize attention prediction with lightweight draft models, compress and cache sparse masks, and extend to other attention variants.

STS is an important progress in sparse attention research, proving that high-quality dynamic sparse attention can be achieved using speculative decoding draft models. With a 2.67x speedup at 90% sparsity and minimal accuracy loss, it provides a powerful tool for deploying large models in long-context scenarios, representing a new idea of model collaboration and computational reuse to break through bottlenecks.