Обучаемое логарифмически-линейное разреженное внимание для эффективных диффузионных трансформаторов

Trainable Log-linear Sparse Attention for Efficient Diffusion Transformers

December 18, 2025

Авторы: Yifan Zhou, Zeqi Xiao, Tianyi Wei, Shuai Yang, Xingang Pan

cs.AI

Аннотация

Диффузионные трансформаторы (DiT) задают современный уровень в области визуальной генерации, однако их квадратичная вычислительная сложность, обусловленная механизмом самовнимания, фундаментально ограничивает масштабирование на длинные последовательности токенов. Недавние подходы с разреженным вниманием Top-K сокращают вычисления в DiT путем сжатия токенов в блочные представления и выбора небольшого набора релевантных ключевых блоков, но все еще страдают от (i) квадратичной стоимости выбора на сжатых токенах и (ii) необходимости увеличивать K для сохранения качества модели по мере роста последовательностей. Мы выявили, что их неэффективность связана с одноуровневым дизайном, поскольку единственный грубый уровень недостаточен для представления глобальной структуры. В данной статье мы представляем Log-linear Sparse Attention (LLSA) — обучаемый механизм разреженного внимания для чрезвычайно длинных последовательностей токенов, который сокращает как стоимость выбора, так и стоимость внимания с квадратичной до логарифмически-линейной сложности за счет использования иерархической структуры. LLSA выполняет иерархический выбор Top-K, постепенно применяя разреженный выбор Top-K с индексами, найденными на предыдущем уровне, и вводит механизм Hierarchical KV Enrichment, который сохраняет глобальный контекст, используя меньшее количество токенов разной гранулярности при вычислении внимания. Для обеспечения эффективного обучения мы разработали высокопроизводительную реализацию для GPU, которая использует только разреженные индексы как для прямого, так и для обратного прохода, устраняя необходимость в плотных масках внимания. Мы оцениваем LLSA на задаче генерации изображений в высоком разрешении в пиксельном пространстве без использования патчификации и кодирования VAE. LLSA ускоряет вывод внимания в 28.27 раз и обучение DiT в 6.09 раз на последовательностях токенов размером 256x256 пикселей, сохраняя при этом качество генерации. Результаты демонстрируют, что LLSA предлагает перспективное направление для эффективного обучения DiT на длинных последовательностях. Код доступен по адресу: https://github.com/SingleZombie/LLSA

English

Diffusion Transformers (DiTs) set the state of the art in visual generation, yet their quadratic self-attention cost fundamentally limits scaling to long token sequences. Recent Top-K sparse attention approaches reduce the computation of DiTs by compressing tokens into block-wise representation and selecting a small set of relevant key blocks, but still suffer from (i) quadratic selection cost on compressed tokens and (ii) increasing K required to maintain model quality as sequences grow. We identify that their inefficiency is due to the single-level design, as a single coarse level is insufficient to represent the global structure. In this paper, we introduce Log-linear Sparse Attention (LLSA), a trainable sparse attention mechanism for extremely long token sequences that reduces both selection and attention costs from quadratic to log-linear complexity by utilizing a hierarchical structure. LLSA performs hierarchical Top-K selection, progressively adopting sparse Top-K selection with the indices found at the previous level, and introduces a Hierarchical KV Enrichment mechanism that preserves global context while using fewer tokens of different granularity during attention computation. To support efficient training, we develop a high-performance GPU implementation that uses only sparse indices for both the forward and backward passes, eliminating the need for dense attention masks. We evaluate LLSA on high-resolution pixel-space image generation without using patchification and VAE encoding. LLSA accelerates attention inference by 28.27x and DiT training by 6.09x on 256x256 pixel token sequences, while maintaining generation quality. The results demonstrate that LLSA offers a promising direction for training long-sequence DiTs efficiently. Code is available at: https://github.com/SingleZombie/LLSA