Ускорение видео-диффузии с обучаемым разреженным вниманием

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) ограничено их квадратичным 3D-вниманием, несмотря на то, что большая часть внимания сосредоточена на небольшом подмножестве позиций. Мы превращаем это наблюдение в VSA — обучаемое, аппаратно-эффективное разреженное внимание, которое заменяет полное внимание как на этапе обучения, так и на этапе вывода. В VSA легковесный грубый этап объединяет токены в тайлы и идентифицирует высоковесные критические токены; тонкий этап вычисляет токен-уровневое внимание только внутри этих тайлов, подчиняясь блочной вычислительной структуре для обеспечения аппаратной эффективности. Это приводит к единому дифференцируемому ядру, которое обучается сквозным образом, не требует постобработки и поддерживает 85\% MFU FlashAttention3. Мы провели масштабные исследования абляции и эксперименты по законам масштабирования, предварительно обучая DiTs с параметрами от 60M до 1.4B. VSA достигает точки Парето, сокращая вычислительные затраты на обучение в 2.53 раза без потери качества диффузии. Адаптация открытой модели Wan-2.1 ускоряет время внимания в 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.