A Unified Sparse Attention via Multi-Granularity Compression

Efficient long-context understanding and reasoning are increasingly vital for large language model (LLM) applications such as multi-turn dialogue and program analysis. However, the core self-attention mechanism scales quadratically with sequence length, creating a fundamental computational bottleneck. Existing sparse attention methods alleviate this issue but face trade-offs: training-based methods are costly and cannot be directly applied as acceleration plugins for other models, while inference-time methods often compromise efficiency or cross-modal generality. To address these limitations, we present UniSparse, a unified mechanism that introduces the notion of composite tokens–compact representations that aggregate multi-granularity contextual information. Building on this abstraction, UniSparse dynamically constructs sparse attention through multi-granularity compression and block-level selection, enabling efficient and hardware-friendly execution on GPU. Across multiple modalities and tasks ranging from synthetic benchmarks to real-world applications, UniSparse consistently surpasses state-of-the-art sparse attention methods (e.g., MInference, XAttention, FlexPrefill) in both accuracy and efficiency, achieving $\ge$ 99% of full-attention accuracy and up to 2.61$\times$ faster attention computation than FlashAttention.

💡 Research Summary

The paper tackles the quadratic cost of self‑attention in large language models (LLMs) when processing very long sequences, a problem that becomes a bottleneck for applications such as multi‑turn dialogue, code analysis, and multimodal reasoning. Existing sparse‑attention approaches fall into two categories: training‑time methods that embed sparsity patterns during pre‑training but lack plug‑and‑play flexibility, and inference‑time methods that either use static, input‑agnostic masks (fast but inflexible) or dynamic masks generated by costly proxy computations (accurate but heavyweight). The authors ask whether a proxy can be both hardware‑friendly and robust across modalities without any model retraining.

UniSparse introduces “composite tokens,” which are compact representations obtained by spatial pooling (average pooling by default) of neighboring fine‑grained tokens. The key hypothesis is that the relative importance of blocks can be reliably inferred from these compressed representations because local tokens often share semantic coherence. Empirical evidence on the HELMET benchmark shows a Spearman correlation above 0.98 between block importance rankings computed in the compressed space (compression ratio 8) and those from full‑resolution attention, confirming that ranking, not absolute scores, drives effective mask selection.

The method proceeds in three stages. First, multi‑granularity compression independently downsamples the query (Q) and key (K) matrices along the sequence dimension (and optionally the head dimension) using integer compression factors c_q and c_k. This yields compressed blocks ˜Q and ˜K of size S′ = S / c, dramatically reducing the number of token‑level operations. Second, attention scores are computed in this reduced space, producing an S′×S′ score matrix for each block pair. These scores are aggregated at the block level (e.g., mean or max) and a Top‑P selection identifies the most salient blocks, forming a binary mask M. Third, the mask is fed into a standard block‑wise attention kernel such as FlashAttention, preserving full‑attention quality while only computing on the selected block pairs.

Complexity analysis shows that the selection cost C_select drops from O(L²) in many prior dynamic methods to O(L·S / c), essentially negligible compared to the sparse computation itself. Consequently, UniSparse achieves up to 2.61× speed‑up in attention computation while retaining ≥99% of full‑attention accuracy, using 1.5–2× fewer attention blocks. Experiments span pure‑text LLMs (Llama‑3.1‑8B‑Instruct, Qwen‑2.5‑7B‑Instruct) and multimodal models (vision‑language, video, audio), consistently outperforming state‑of‑the‑art sparse baselines such as MInference, SeerAttention, XAttention, and FlexPrefill. The approach is modality‑agnostic because average pooling is a universal operation applicable to any sequential data representation.

The paper also discusses practical implications: the method integrates seamlessly with existing GPU‑optimized kernels, requires no additional training or model‑specific heuristics, and reduces memory bandwidth by operating on compressed tensors. Limitations include sensitivity to the chosen compression ratios; overly aggressive compression may discard fine‑grained cues essential for tasks like code reasoning. Future work may explore learnable compression modules or clustering‑based token aggregation to adaptively balance compression and fidelity.

In summary, UniSparse provides a unified, plug‑and‑play sparse attention mechanism that leverages compressed composite tokens to generate high‑quality dynamic masks with minimal overhead. It bridges the gap between accuracy and efficiency in long‑context LLM inference, offering a scalable solution for real‑time deployment across diverse modalities without retraining.