Vmoba Sparse Video Attention

VMoBA: Efficient Video Attention Through Adaptive Spatio-Temporal Sparsity

Video diffusion models generate high-quality videos but suffer from quadratic complexity in attention—processing each frame pair explodes computational costs, making long-sequence training impractical. VMoBA solves this by analyzing real attention patterns in pre-trained video models, discovering that attention naturally clusters in spatio-temporal neighborhoods. By replacing full attention with adaptive block selection that respects these patterns, VMoBA achieves 2.92× FLOPs reduction and 1.48× wall-clock speedup while maintaining or improving generation quality.

The insight is that video attention doesn't need all frame pairs—it primarily focuses on temporal neighbors (consecutive frames), spatial regions (nearby pixels), and occasional global interactions. Naive block attention designed for text performs poorly on video; VMoBA introduces layer-wise cyclic partitioning and global selection mechanisms tailored to video's specific spatio-temporal structure.

Core Concept

VMoBA builds on the Mixture-of-Block Attention (MoBA) framework but adapts it fundamentally for video. The key innovations are:

1. Layer-wise Cyclic Block Partitioning: Alternates between 1D temporal blocking, 2D spatial blocking, and 3D spatio-temporal blocking across layers, allowing the model to learn diverse attention patterns.

2. Global Block Selection: Instead of selecting blocks per query, aggregates query-key similarities across all queries to identify top-scoring key blocks, reducing selection overhead.

3. Adaptive Threshold-based Selection: Dynamically determines block counts based on cumulative similarity exceeding a threshold, rather than fixed top-k, enabling variable capacity per layer.

These mechanisms work together to approximate full attention where needed while maintaining sparsity where possible, achieving speed without quality loss.

Architecture Overview

• 1D Block Partitioning (Temporal): Keys grouped by frame index, focusing attention on frame neighbors
• 2D Block Partitioning (Spatial): Keys grouped by spatial regions within frames, capturing local pixel patterns
• 3D Block Partitioning (Spatio-Temporal): Keys grouped by 3D volumes combining frame and spatial structure
• Layer-wise Cycling: Layers alternate (1D→2D→3D→1D...) ensuring diverse attention pattern learning
• Global Similarity Aggregation: Sums query-key similarities across queries to identify important blocks
• Threshold-based Adaptive Selection: Cumulative similarity scores determine actual block count dynamically

Implementation

The block partitioning strategy cycles through three patterns across layers:

import torch
import torch.nn as nn
from typing import List, Tuple
import math

class LayerWiseBlockPartitioner:
    """
    Cycles through 1D (temporal), 2D (spatial), and 3D (spatio-temporal)
    block partitioning schemes across diffusion layers.
    """
    def __init__(self, num_layers=24, frame_count=16, spatial_size=32):
        self.num_layers = num_layers
        self.frame_count = frame_count
        self.spatial_h = spatial_size
        self.spatial_w = spatial_size
        self.partition_cycle = ['1d', '2d', '3d']

    def get_partition_type(self, layer_idx: int) -> str:
        """Determine partition type for layer (cycles 1D→2D→3D)."""
        cycle_position = layer_idx % 3
        return self.partition_cycle[cycle_position]

    def partition_keys_1d(self, keys: torch.Tensor) -> List[torch.Tensor]:
        """
        Temporal partitioning: group keys by frame.
        Args:
            keys: (batch, seq_len, dim) where seq_len = frames * height * width

        Returns:
            blocks: List of (batch, block_size, dim) tensors, one per frame
        """
        seq_len = keys.shape[1]
        spatial_tokens_per_frame = seq_len // self.frame_count

        blocks = []
        for frame_idx in range(self.frame_count):
            start = frame_idx * spatial_tokens_per_frame
            end = start + spatial_tokens_per_frame
            frame_block = keys[:, start:end, :]
            blocks.append(frame_block)

        return blocks

    def partition_keys_2d(self, keys: torch.Tensor) -> List[torch.Tensor]:
        """
        Spatial partitioning: group keys by spatial regions within each frame.
        """
        batch_size, seq_len, dim = keys.shape
        spatial_tokens_per_frame = seq_len // self.frame_count

        blocks = []
        patch_size = 4  # Divide each frame into 4×4 patches

        for frame_idx in range(self.frame_count):
            frame_start = frame_idx * spatial_tokens_per_frame
            frame_keys = keys[:, frame_start:frame_start + spatial_tokens_per_frame, :]

            # Reshape to spatial grid
            spatial_grid = frame_keys.view(
                batch_size, self.spatial_h, self.spatial_w, dim
            )

            # Partition into spatial patches
            for patch_y in range(0, self.spatial_h, patch_size):
                for patch_x in range(0, self.spatial_w, patch_size):
                    patch = spatial_grid[
                        :, patch_y:patch_y+patch_size, patch_x:patch_x+patch_size, :
                    ]
                    # Flatten spatial dims
                    patch_flat = patch.reshape(batch_size, -1, dim)
                    blocks.append(patch_flat)

        return blocks

    def partition_keys_3d(self, keys: torch.Tensor) -> List[torch.Tensor]:
        """
        Spatio-temporal partitioning: 3D volumes combining frames and spatial.
        """
        batch_size, seq_len, dim = keys.shape
        spatial_tokens_per_frame = seq_len // self.frame_count

        blocks = []
        frame_block_size = 4  # Group 4 consecutive frames
        spatial_block_size = 4

        for frame_block_idx in range(0, self.frame_count, frame_block_size):
            frames_in_block = min(
                frame_block_size, self.frame_count - frame_block_idx
            )

            frame_end = frame_block_idx + frames_in_block
            temporal_keys = keys[
                :, frame_block_idx*spatial_tokens_per_frame:frame_end*spatial_tokens_per_frame, :
            ]

            # Partition temporal chunk into spatial patches
            for patch_y in range(0, self.spatial_h, spatial_block_size):
                for patch_x in range(0, self.spatial_w, spatial_block_size):
                    # Extract 3D volume
                    block_3d = temporal_keys  # Simplified; full implementation indexes spatially
                    blocks.append(block_3d)

        return blocks

Global block selection identifies important interactions:

class GlobalBlockSelector:
    """
    Selects important blocks based on aggregated query-key similarities
    across all queries, reducing per-query overhead.
    """
    def __init__(self, similarity_threshold=0.1):
        self.threshold = similarity_threshold

    def select_blocks(
        self,
        queries: torch.Tensor,  # (batch, num_queries, dim)
        blocks: List[torch.Tensor],  # List of (batch, block_size, dim)
        num_heads: int = 12
    ) -> Tuple[List[torch.Tensor], torch.Tensor]:
        """
        Select blocks by aggregating similarities across queries.

        Returns:
            selected_blocks: List of selected block tensors
            selection_mask: (batch, num_blocks) binary mask
        """
        batch_size = queries.shape[0]
        num_blocks = len(blocks)
        dim = queries.shape[-1]
        head_dim = dim // num_heads

        # Aggregate similarities across all queries
        similarities = torch.zeros(batch_size, num_blocks)

        for block_idx, block in enumerate(blocks):
            # Compute query-block similarity
            # Simplified: cosine similarity summed over queries
            block_sim = torch.einsum('bqd,bkd->bq', queries, block) / math.sqrt(head_dim)
            # Aggregate over queries
            similarities[:, block_idx] = block_sim.max(dim=1)[0]

        # Threshold-based selection: select blocks exceeding threshold
        selection_mask = (similarities > self.threshold).float()

        # Ensure at least some blocks selected (e.g., top-k backup)
        min_blocks = max(1, int(num_blocks * 0.2))  # At least 20% of blocks
        for b in range(batch_size):
            num_selected = selection_mask[b].sum()
            if num_selected < min_blocks:
                # Fallback: select top-k blocks
                top_k = torch.topk(similarities[b], min_blocks)[1]
                selection_mask[b, top_k] = 1.0

        selected_blocks = [
            block for block_idx, block in enumerate(blocks)
            if selection_mask[0, block_idx] > 0  # Simplified: use first batch
        ]

        return selected_blocks, selection_mask

VMoBA attention layer integrates layer-wise cycling and global selection:

class VMoBAAttentionLayer(nn.Module):
    """
    Sparse attention layer combining cyclic block partitioning
    and global block selection for video diffusion.
    """
    def __init__(self, dim=768, num_heads=12, layer_idx=0, num_layers=24):
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.layer_idx = layer_idx

        self.partitioner = LayerWiseBlockPartitioner(
            num_layers=num_layers, frame_count=16, spatial_size=32
        )
        self.selector = GlobalBlockSelector(similarity_threshold=0.1)

        # Query, key, value projections
        self.q_proj = nn.Linear(dim, dim)
        self.k_proj = nn.Linear(dim, dim)
        self.v_proj = nn.Linear(dim, dim)
        self.out_proj = nn.Linear(dim, dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Apply VMoBA sparse attention.
        Args:
            x: (batch, seq_len, dim)

        Returns:
            attended: (batch, seq_len, dim)
        """
        batch_size, seq_len, dim = x.shape

        # Project to Q, K, V
        queries = self.q_proj(x)
        keys = self.k_proj(x)
        values = self.v_proj(x)

        # Reshape for multi-head
        queries = queries.view(batch_size, seq_len, self.num_heads, dim // self.num_heads)
        queries = queries.transpose(1, 2)  # (batch, heads, seq_len, head_dim)

        # Select partition type for this layer
        partition_type = self.partitioner.get_partition_type(self.layer_idx)

        # Partition keys into blocks
        if partition_type == '1d':
            blocks = self.partitioner.partition_keys_1d(keys)
        elif partition_type == '2d':
            blocks = self.partitioner.partition_keys_2d(keys)
        else:  # '3d'
            blocks = self.partitioner.partition_keys_3d(keys)

        # Select important blocks globally
        selected_blocks, _ = self.selector.select_blocks(queries, blocks)

        # Compute attention only over selected blocks
        attended = self._sparse_attention(
            queries, selected_blocks, values, self.num_heads
        )

        # Reshape back and project
        attended = attended.transpose(1, 2).contiguous()
        attended = attended.view(batch_size, seq_len, dim)
        output = self.out_proj(attended)

        return output

    def _sparse_attention(self, queries, selected_blocks, values, num_heads):
        """Compute attention only over selected sparse blocks."""
        # Simplified: compute attention on concatenated selected blocks
        selected_keys = torch.cat(selected_blocks, dim=1)
        attn_weights = torch.softmax(
            torch.bmm(queries, selected_keys.transpose(1, 2)) / math.sqrt(self.dim),
            dim=-1
        )
        attended = torch.bmm(attn_weights, selected_keys)
        return attended

Practical Guidance

When to use:

• Training video diffusion models on memory-constrained hardware
• Processing long video sequences (32+ frames)
• Reducing training time without sacrificing quality
• Deploying video models with strict computational budgets
• Scaling to higher resolution video generation

When NOT to use:

• If you need guaranteed identical attention to full attention (sparsity introduces approximation)
• Short videos (2-4 frames) where overhead of block selection dominates
• Scenarios where attention weights must be interpretable (sparse patterns are complex)
• Fixed-block attention patterns without layer-wise adaptation
• Systems with abundant compute where efficiency doesn't matter

Common pitfalls:

• Threshold value too high, discarding important interactions
• Threshold too low, reducing speedup benefit (set empirically per model)
• Not cycling block partitioning types, causing one pattern to dominate
• Using identical similarity threshold across all layers (deeper layers need higher tolerance)
• Assuming sparse patterns transfer across model sizes/depths
• Ignoring spatial clustering in temporal sequences (1D blocking alone insufficient)

Reference

"VMoBA: Mixture-of-Block Attention for Video Diffusion Models", 2025. arxiv.org/abs/2506.23858