Video Mixture-of-Block Attention (VMoBA)

• VMoBA is a sparse attention mechanism for video diffusion models that integrates spatio-temporal priors using cyclic block partitioning.
• It employs a recurrent 1D-2D-3D block partitioning strategy and global block selection to allocate computational resources to high-affinity tokens.
• VMoBA achieves up to 3× FLOPs reduction and faster inference while maintaining or improving generation quality in high-resolution video tasks.

Video Mixture-of-Block Attention (VMoBA) is a sparse attention mechanism tailored for Video Diffusion Models (VDMs) operating over long-duration, high-resolution video sequences. It generalizes the Mixture-of-Block Attention (MoBA) paradigm—originally proposed for LLMs—by integrating spatio-temporal priors and global block selection dynamics specific to video data, achieving substantial reductions in computational overhead while maintaining or exceeding the generation quality of dense, full attention (Wu et al., 30 Jun 2025).

1. Empirical Motivation and Analysis

Analysis of attention maps in pre-trained video diffusion transformers reveals three persistent patterns:

1. Spatio-temporal locality: Distinct attention heads exhibit localized patterns—some heads specialize in 1D (temporal) attention across frames, others in 2D (spatial) within frames, and yet others in full 3D (spatio-temporal) neighborhoods.
2. Query importance variation: The sum of a query’s top-k key similarities varies significantly; some queries exhibit high concentration over a few blocks (e.g., objects of interest), while others have distributed or low affinities. Fixed per-query block allocation can under-provision "important" queries.
3. Head-specific concentration levels: Distributions of query–block similarities differ markedly between heads: some drop steeply (concentrated), others display long-tailed, diffuse profiles, making a uniform top-k per head suboptimal.

These findings motivate the core VMoBA design: adapting block partitioning (to spatio-temporal structure), allowing global competition in block assignment, and using adaptive thresholds for head-specific sparsity control (Wu et al., 30 Jun 2025).

2. Block Partitioning and Cyclic Layerwise Scheme

VMoBA introduces a recurrent 1D–2D–3D block partition across transformer layers to match emergent attention behaviors in video models. At each layer $l$, keys are partitioned into blocks as follows:

• 1D (temporal) partition ($l \bmod 3 = 0$): Blocks cover consecutive frames, shape $(s_b^T \times H \times W)$, yielding $N_b^T = T/s_b^T$ blocks.
• 2D (spatial) partition ($l \bmod 3 = 1$): Blocks cover patches within each frame, shape $(T \times s_b^H \times s_b^W)$, total blocks $N_b^H \times N_b^W = (H/s_b^H) \cdot (W/s_b^W)$.
• 3D (spatiotemporal) partition ($l \bmod 3 = 2$): Cubic spatio-temporal blocks, shape $(s_b^T \times s_b^H \times s_b^W)$.

For each block, a mean-pooled block centroid is computed. This cyclic block partitioning ensures that different locality structures are captured at each layer, aligning with empirical head behaviors in VDMs.

Pseudocode (see (Wu et al., 30 Jun 2025)):

def partition_keys(K, layer): if layer % 3 == 0: block_shape = (s_b_T, H, W) # Temporal elif layer % 3 == 1: block_shape = (T, s_b_H, s_b_W) # Spatial else: block_shape = (s_b_T, s_b_H, s_b_W) # Spatio-temporal K_blocks = rearrange(K, block_shape) block_centroids = K_blocks.mean(axis=1) return block_centroids, K_blocks

3. Global Block Selection and Adaptive Thresholding

Global Block Selection

Departing from MoBA’s per-query top-k block masking, VMoBA selects the $k$ largest query–block affinities globally within each attention head. Let $S_i \in \mathbb{R}^{s \times N_b}$ be the similarity matrix between $s$ queries and $N_b$ blocks for head $i$. All $s \cdot N_b$ entries in $S_i$ are considered, the top $k$ are selected, and a binary mask $M_i \in \{0,1\}^{s \times N_b}$ is constructed. This approach disproportionately allocates block capacity to the most salient queries, as empirically required by high-affinity tokens.

Threshold-Based Block Selection

To account for head-specific sparsity, a cumulative-sum threshold $\tau$ is applied:

1. Flatten and normalize $S_i$ so its entries sum to 1.
2. Sort (descending) and select the minimal $k$ such that the cumulative sum $\geq \tau$.
3. Mask the $k$ largest interactions.

Pseudocode (see (Wu et al., 30 Jun 2025)):

def select_blocks_threshold(S_i, tau): flat = S_i.flatten() normalized = flat / flat.sum() sorted_vals = np.sort(normalized)[::-1] cumsum = 0.0 for j, val in enumerate(sorted_vals, 1): cumsum += val if cumsum >= tau: k = j break mask = TopkMask(S_i, k) return mask

This dual-stage selection resolves both query and head-level variation in concentration, dynamically adjusting the compute allocation to the most relevant regions.

4. Computational Complexity and Acceleration

Let $s = T \cdot H \cdot W$, $d$ hidden dimension, $N_b$ number of blocks, and $s_b$ block size.

Method	FLOPs per head	Dominant Terms	Latency Growth
Full Attn	$O(s^2 \cdot d)$	pairwise (token, token)	$\propto s^2$
VMoBA	$$O( s \cdot d \cdot (N_b + k_{\mathrm{avg}} \cdot s_b) )$$	block scoring plus local attention	$\sim O(s^{1.5})$

With typical choices ($s_b \approx 64$, $k_{\mathrm{avg}} \approx 18$), VMoBA reduces FLOPs by approximately $3\times$, closely matching empirical speedups (Wu et al., 30 Jun 2025).

5. Empirical Evaluation

Training-Based Long-Sequence Fine-Tuning

For videos mapped to $55$K tokens ($93 \times 576 \times 1024$):

Method	FLOPs (T)	Latency (GPU-h)	TextConsis (%)	Dynamic (%)	BGConsis (%)	ImageQual (%)	SubConsist (%)
FullAttn	705 (1x)	276 (1x)	24.61	61.58	94.69	69.49	90.86
VMoBA	248.7 (2.83x less)	187 (1.48x faster)	25.88	56.91	96.76	67.45	94.72

Training-Free High-Resolution Inference

On $76$K-token ($81 \times 720 \times 1280$) videos:

Method	PSNR	FLOPs (T)	Latency (s)	TextConsis (%)	BGConsis (%)
FullAttn	–	1246.8 (1x)	406 (1x)	27.99	93.74
VMoBA	18.80	519.8 (2.4x less)	300 (1.35x faster)	28.06	92.85

VMoBA achieves comparable or higher text consistency and background consistency, despite significantly lower resource consumption.

6. Implementation Considerations

VMoBA is implemented as a drop-in module in diffusion transformer stacks. The layerwise cyclic partitioning requires dynamic key/value rearrangement per layer, but selected blocks are contiguous, supporting efficient token gathering. Since block selection and masking are per-head, the method is compatible with parallelism and standard deep learning pipelines. VMoBA is applicable both during training (for end-to-end differentiable attention) and inference (for acceleration), without reliance on training-free approximations.

7. Context and Impact

VMoBA introduces a set of strategies—cyclic block locality, global block selection, and adaptive headwise thresholds—that specifically exploit the spatio-temporal structure and heterogeneity observed in video attention maps, diverging from the fixed top-k, per-query routing present in MoBA (Lu et al., 18 Feb 2025). The result is a structured sparse mechanism that:

• Reduces drift from full attention for high-affinity tokens and heads,
• Minimizes redundant computation on low-relevance regions,
• Achieves 2.92× FLOPs and 1.48× latency speedup in training, and 2.40× FLOPs and 1.35× latency speedup in inference, matching or surpassing full attention generation quality.

These design principles may generalize to other domains characterized by block-wise locality and variable query salience (Wu et al., 30 Jun 2025).