AB-Attention Brings More Accurate Sparse Attention to Video Diffusion

Video diffusion Transformers produce increasingly capable results, but their inference cost rises rapidly with frame count, spatial resolution, and model depth. A major bottleneck is 3D full attention, whose quadratic token-to-token interactions become especially expensive for long or high-resolution clips. Dynamic block-sparse attention reduces this burden by selecting only some key/value blocks, yet aggressive pruning can create a substantial gap from dense attention.

AB-Attention addresses that gap by treating block selection as an approximation problem rather than a simple keep-or-drop decision. Existing approaches may score blocks through coarse summaries or rely on expensive global clustering. They can also discard an entire tile even when only a small number of keys inside it are relevant to the current queries. Once a block is pruned, it no longer participates in the softmax, which can distort the resulting attention distribution.

The first component, the Intra-Block Clustering Selector (IBCS), clusters queries locally within each query block. It does not permute tokens and avoids the cost of constructing a global clustering structure. This local view provides a more detailed estimate of which key/value blocks matter to different query patterns inside the same block. Compared with a single coarse score, it is intended to preserve more of the variation that drives attention.

The second component is Sampling with Inclusion-Probability Weighting (SIP). Rather than deterministically truncating all blocks below a threshold, the method samples blocks according to their estimated importance. It then reweights sampled blocks using the inverse of their inclusion probability. In this way, one sparse pass serves as an estimator of dense attention instead of treating unselected blocks as permanently irrelevant. The correction is particularly important when the target sparsity is high and selection errors would otherwise accumulate.

AB-Attention also includes a hardware-oriented kernel design. Key/value blocks that are smaller than a standard tile can be gathered into full Tensor Core tiles. This enables finer selection granularity while keeping the matrix operations compatible with Tensor Core execution, without adding Tensor Core cost according to the paper. The result is a co-design of the approximation rule and the GPU execution path, rather than an optimization based only on theoretical FLOP reduction.

Experiments on Wan2.1 and HunyuanVideo show that AB-Attention skips 89%–92% of attention computation. In the reported comparisons, the baselines perform 1.8–3.6 times as much attention computation. The paper also reports an improvement of up to 4.4 dB in PSNR over the strongest baseline and up to 2.23× end-to-end acceleration.

The broader significance is that sparse video diffusion attention can be made less dependent on crude block truncation. Importance estimation, unbiased-style probability correction, and GPU-friendly gathering are combined into one inference-time method, with no additional training described in the supplied material. The results are promising, but they come from Wan2.1 and HunyuanVideo under specific experimental settings. Actual gains will still depend on hardware, sequence length, sparsity targets, and the workload. Further validation across models and generation conditions will be important before treating the approach as a general solution.

Source: Hugging Face Daily Papers