VMoBA: Mixture-of-Block Attention for Video Diffusion Models

TL;DR

VMoBA tackles the quadratic bottleneck of full attention in video diffusion models by introducing a sparse, video-tuned attention mechanism. It combines layer-wise recurrent 1D-2D-3D block partitioning, global block selection, and threshold-based block selection to focus computation on the most informative query–block interactions, achieving significant training-time speedups with comparable or improved video quality. In training-free inference, VMoBA also delivers notable FLOPs and latency reductions, demonstrating strong practical impact for long-sequence, high-resolution video generation. The results show robust gains across multiple resolutions and prompts, highlighting VMoBA's effectiveness in adapting sparse attention to the spatio-temporal locality of video data.

Abstract

The quadratic complexity of full attention mechanisms poses a significant bottleneck for Video Diffusion Models (VDMs) aiming to generate long-duration, high-resolution videos. While various sparse attention methods have been proposed, many are designed as training-free inference accelerators or do not optimally capture the unique spatio-temporal characteristics inherent in video data when trained natively. This paper introduces Video Mixture of Block Attention (VMoBA), a novel sparse attention mechanism specifically adapted for VDMs. Motivated by an in-depth analysis of attention patterns within pre-trained video transformers, which revealed strong spatio-temporal locality, varying query importance, and head-specific concentration levels, VMoBA enhances the original MoBA framework with three key modifications: (1) a layer-wise recurrent block partition scheme (1D-2D-3D) to dynamically adapt to diverse spatio-temporal attention patterns and improve efficiency; (2) global block selection to prioritize the most salient query-key block interactions across an entire attention head; and (3) threshold-based block selection to dynamically determine the number of attended blocks based on their cumulative similarity. Extensive experiments demonstrate that VMoBA significantly accelerates the training of VDMs on longer sequences, achieving 2.92x FLOPs and 1.48x latency speedup, while attaining comparable or even superior generation quality to full attention. Furthermore, VMoBA exhibits competitive performance in training-free inference, offering 2.40x FLOPs and 1.35x latency speedup for high-res video generation.

Figures (7)

• Figure 1: (a) VMoBA performs better than Full Attention while reducing training time. (b) VMoBA is faster when the sequence length goes higher. The 'x' points represent tested latencies, and the lines are fitted with quadratic functions. (c) VMoBA is an efficient and effective adaptation of MoBA.
• Figure 2: The overall pipeline of VMoBA. We first partition key blocks with Layer-wise Recurrent BLock Partition, then select the blocks using Global Block Selection and Threshold-based Block Selection. Finally, the attention is computed only with the selected blocks.
• Figure 3: The query-key block attention map shows 1-2-3D distinct patterns.
• Figure 4: Summation of top 25% query-key similarities. Different queries have varying importance.
• Figure 5: Sorted query-key block similarity and top 30%/50% cutoff lines. The right parts of the cutoff lines contain corresponding cumulative summations of query-key block similarity.