OmniSparse: Training-Aware Fine-Grained Sparse Attention for Long-Video MLLMs

TL;DR

OmniSparse introduces a training-aware, fine-grained sparse attention method for long-video multimodal language models, enabling dynamic token selection and significant speed and memory improvements while maintaining full attention performance.

Contribution

It proposes a novel framework with adaptive mechanisms for token and head selection that operate during both training and inference, bridging the training-inference gap in sparse attention methods.

Findings

Achieves up to 2.7x speedup during prefill

Reduces memory usage by 2.4x during decoding

Maintains full attention performance

Abstract

Existing sparse attention methods primarily target inference-time acceleration by selecting critical tokens under predefined sparsity patterns. However, they often fail to bridge the training-inference gap and lack the capacity for fine-grained token selection across multiple dimensions such as queries, key-values (KV), and heads, leading to suboptimal performance and limited acceleration gains. In this paper, we introduce OmniSparse, a training-aware fine-grained sparse attention framework for long-video MLLMs, which operates in both training and inference with dynamic token budget allocation. Specifically, OmniSparse contains three adaptive and complementary mechanisms: (1) query selection via lazy-active classification, retaining active queries that capture broad semantic similarity while discarding most lazy ones that focus on limited local context and exhibit high functional redundancy; (2) KV selection with head-level dynamic budget allocation, where a shared budget is determined based on the flattest head and applied uniformly across all heads to ensure attention recall; and (3) KV cache slimming to reduce head-level redundancy by selectively fetching visual KV cache according to the head-level decoding query pattern. Experimental results show that OmniSparse matches the performance of full attention while achieving up to 2.7x speedup during prefill and 2.4x memory reduction during decoding.