훈련 가능한 희소 주의력을 통한 더 빠른 비디오 확산

Faster Video Diffusion with Trainable Sparse Attention

May 19, 2025

저자: Peiyuan Zhang, Haofeng Huang, Yongqi Chen, Will Lin, Zhengzhong Liu, Ion Stoica, Eric P. Xing, Hao Zhang

cs.AI

초록

비디오 확산 트랜스포머(DiTs)의 확장은 3차원 어텐션의 이차적 복잡도로 인해 제한되며, 대부분의 어텐션 가중치는 소수의 위치에 집중된다는 점이 관찰되었습니다. 우리는 이러한 관찰을 바탕으로 VSA(Video Sparse Attention)를 제안합니다. VSA는 학습 가능하고 하드웨어 효율적인 희소 어텐션으로, 학습과 추론 모두에서 전체 어텐션을 대체합니다. VSA에서는 경량화된 coarse 단계에서 토큰을 타일로 풀링하고 높은 가중치를 가진 중요한 토큰을 식별합니다. 이후 fine 단계에서는 해당 타일 내에서만 토큰 수준의 어텐션을 계산하며, 블록 컴퓨팅 레이아웃을 적용하여 하드웨어 효율성을 보장합니다. 이를 통해 단일 미분 가능 커널이 엔드투엔드로 학습되며, 사후 프로파일링이 필요 없고 FlashAttention3 MFU의 85%를 유지합니다. 우리는 60M에서 1.4B 파라미터까지 DiTs를 사전 학습하며 대규모 ablation 연구와 스케일링 법칙 실험을 수행했습니다. VSA는 확산 손실의 감소 없이 학습 FLOPS를 2.53배 절감하는 파레토 포인트에 도달했습니다. 오픈소스 Wan-2.1 모델에 VSA를 적용한 결과, 어텐션 시간이 6배 빨라졌고, 엔드투엔드 생성 시간이 31초에서 18초로 단축되었으며 품질은 유지되었습니다. 이러한 결과는 학습 가능한 희소 어텐션이 전체 어텐션의 실용적인 대안이자 비디오 확산 모델의 추가 확장을 위한 핵심 요소임을 입증합니다.

English

Scaling video diffusion transformers (DiTs) is limited by their quadratic 3D attention, even though most of the attention mass concentrates on a small subset of positions. We turn this observation into VSA, a trainable, hardware-efficient sparse attention that replaces full attention at both training and inference. In VSA, a lightweight coarse stage pools tokens into tiles and identifies high-weight critical tokens; a fine stage computes token-level attention only inside those tiles subjecting to block computing layout to ensure hard efficiency. This leads to a single differentiable kernel that trains end-to-end, requires no post-hoc profiling, and sustains 85\% of FlashAttention3 MFU. We perform a large sweep of ablation studies and scaling-law experiments by pretraining DiTs from 60M to 1.4B parameters. VSA reaches a Pareto point that cuts training FLOPS by 2.53times with no drop in diffusion loss. Retrofitting the open-source Wan-2.1 model speeds up attention time by 6times and lowers end-to-end generation time from 31s to 18s with comparable quality. These results establish trainable sparse attention as a practical alternative to full attention and a key enabler for further scaling of video diffusion models.