Atenção Esparsa Log-linear Treinável para Transformadores de Difusão Eficientes

Trainable Log-linear Sparse Attention for Efficient Diffusion Transformers

December 18, 2025

Autores: Yifan Zhou, Zeqi Xiao, Tianyi Wei, Shuai Yang, Xingang Pan

cs.AI

Resumo

Os Diffusion Transformers (DiTs) estabelecem o estado da arte na geração visual, mas o seu custo quadrático de self-attention limita fundamentalmente a escalabilidade para sequências longas de tokens. Abordagens recentes de atenção esparsa Top-K reduzem a computação dos DiTs ao comprimir tokens em representações por blocos e selecionar um pequeno conjunto de blocos-chave relevantes, mas ainda sofrem com (i) o custo quadrático de seleção em tokens comprimidos e (ii) o K crescente necessário para manter a qualidade do modelo à medida que as sequências aumentam. Identificamos que a sua ineficiência se deve ao design de nível único, pois um único nível grosso é insuficiente para representar a estrutura global. Neste artigo, introduzimos a Log-linear Sparse Attention (LLSA), um mecanismo de atenção esparsa treinável para sequências extremamente longas de tokens que reduz os custos de seleção e de atenção de uma complexidade quadrática para log-linear, utilizando uma estrutura hierárquica. A LLSA realiza uma seleção Top-K hierárquica, adotando progressivamente uma seleção Top-K esparsa com os índices encontrados no nível anterior, e introduz um mecanismo de Hierarchical KV Enrichment que preserva o contexto global enquanto utiliza menos tokens de diferentes granularidades durante o cálculo da atenção. Para suportar um treino eficiente, desenvolvemos uma implementação de alto desempenho para GPU que utiliza apenas índices esparsos tanto para a passagem direta quanto para a retropassagem, eliminando a necessidade de máscaras de atenção densas. Avaliamos a LLSA na geração de imagens de alta resolução no espaço de pixels sem usar "patchification" nem codificação VAE. A LLSA acelera a inferência de atenção em 28.27x e o treino de DiT em 6.09x em sequências de tokens de 256x256 pixels, mantendo a qualidade de geração. Os resultados demonstram que a LLSA oferece uma direção promissora para treinar DiTs de longa sequência de forma eficiente. O código está disponível em: 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