Attention Log-Linéaire Entraînable et Creuse pour des Transformers de Diffusion Efficaces

Trainable Log-linear Sparse Attention for Efficient Diffusion Transformers

December 18, 2025

papers.authors: Yifan Zhou, Zeqi Xiao, Tianyi Wei, Shuai Yang, Xingang Pan

cs.AI

papers.abstract

Les Transformers de Diffusion (DiTs) établissent l'état de l'art en génération visuelle, mais leur coût quadratique d'auto-attention limite fondamentalement la mise à l'échelle vers de longues séquences de tokens. Les approches récentes d'attention creuse Top-K réduisent le calcul des DiTs en compressant les tokens en une représentation par blocs et en sélectionnant un petit ensemble de blocs clés pertinents, mais souffrent toujours (i) d'un coût de sélection quadratique sur les tokens compressés et (ii) d'un K requis croissant pour maintenir la qualité du modèle à mesure que les séquences s'allongent. Nous identifions que leur inefficacité est due à la conception à niveau unique, un niveau grossier unique étant insuffisant pour représenter la structure globale. Dans cet article, nous présentons Log-linear Sparse Attention (LLSA), un mécanisme d'attention creuse entraînable pour des séquences de tokens extrêmement longues qui réduit les coûts de sélection et d'attention d'une complexité quadratique à log-linéaire en utilisant une structure hiérarchique. LLSA effectue une sélection Top-K hiérarchique, adoptant progressivement une sélection Top-K creuse avec les indices trouvés au niveau précédent, et introduit un mécanisme d'Enrichissement Hiérarchique KV qui préserve le contexte global tout en utilisant moins de tokens de granularité différente pendant le calcul d'attention. Pour supporter un entraînement efficace, nous développons une implémentation GPU haute performance qui n'utilise que des indices creux pour les passes avant et arrière, éliminant le besoin de masques d'attention denses. Nous évaluons LLSA sur la génération d'images en haute résolution dans l'espace pixel sans utiliser de patchification ni d'encodage VAE. LLSA accélère l'inférence d'attention par 28.27x et l'entraînement des DiT par 6.09x sur des séquences de tokens d'images 256x256 pixels, tout en maintenant la qualité de génération. Les résultats démontrent que LLSA offre une direction prometteuse pour entraîner efficacement des DiT à longues séquences. Le code est disponible à l'adresse : 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