Video diffusion transformers run faster with smarter attention retrieval
PARK: Accurate Block Retrieval for Sparse Attention in Video Diffusion Transformers
Computer Vision and Pattern Recognition
Summary
Processing videos using transformers is slow because the method looks at every part of the video in detail, which takes a long time. The authors found that grouping video parts for this process often mixes up important details, which hurts quality or wastes effort. They designed PARK, a way to group parts more carefully by keeping the unique information from each query and grouping keys based on how well they actually match. This method speeds up video generation while keeping good video quality without needing extra training.
What this means in practice
- •For video processing engineers: Speed up high-quality video generation models by more accurately selecting which video parts to attend to during inference.
- •For machine learning deployment teams: Reduce computational costs and latency when deploying video generation systems on GPU hardware by using an efficient block retrieval method.
- •For augmented reality developers: Enable faster and more efficient real-time video synthesis by improving transformer attention retrieval accuracy in video generation.$Commercial implications: Allows developers to build faster, real-time AR video synthesis apps that require efficient transformer-based video generation.
Authors
Yun Dai, Jiarui Wen, Huiping Zhuang, Cen Chen, Ziqian Zeng
Abstract
Diffusion Transformers (DiTs) have become a dominant architecture for video generation, but their efficiency is limited by the quadratic complexity of full attention. Sparse attention reduces this cost by retrieving important blocks and computing attention only within them, but inaccurate retrieval can either degrade generation quality or yield unnecessary computation. We identify two retrieval mismatches in methods that retrieve blocks using the averaged representations of query and key blocks: (i) query-side aggregation mismatch, where averaging queries before Softmax fails to preserve their individual attention preferences, and (ii) key-side clustering metric mismatch, where standard Euclidean clustering in the original key space can group keys with dissimilar QK scores under the current query, so their average representation may not accurately represent how the current query scores individual keys. These mismatches can lead to inaccurate block retrieval. To address these mismatches, we propose PARK, a training-free sparse attention method for accurate block retrieval. PARK retains every original query, independently normalizes its attention over key blocks, and then averages these distributions within each query block. It also uses information from the current queries to transform keys before clustering, so that keys receiving similar QK scores are grouped together. A fused GPU kernel further reduces the overhead of block retrieval. Experiments on HunyuanVideo and Wan demonstrate that PARK improves block retrieval accuracy and preserves generation quality while accelerating inference, achieving the best quality-efficiency trade-off among the compared sparse attention methods.