Trainierbare log-lineare Sparse-Attention für effiziente Diffusion-Transformer

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

Diffusion Transformer (DiTs) setzen den Maßstab in der visuellen Generierung, doch ihre quadratischen Self-Attention-Kosten begrenzen grundlegend die Skalierbarkeit auf lange Token-Sequenzen. Jüngste Top-K-Sparse-Attention-Ansätze reduzieren den Berechnungsaufwand von DiTs, indem sie Token in blockweise Darstellungen komprimieren und eine kleine Menge relevanter Schlüsselblöcke auswählen, leiden aber weiterhin unter (i) quadratischen Auswahlkosten für komprimierte Token und (ii) einem mit wachsenden Sequenzen steigenden K-Wert, um die Modellqualität zu erhalten. Wir identifizieren, dass diese Ineffizienz auf den einstufigen Aufbau zurückzuführen ist, da eine einzelne grobe Ebene unzureichend ist, um die globale Struktur abzubilden. In diesem Artikel stellen wir Log-lineare Sparse Attention (LLSA) vor, einen trainierbaren Sparse-Attention-Mechanismus für extrem lange Token-Sequenzen, der sowohl Auswahl- als auch Attention-Kosten durch Nutzung einer hierarchischen Struktur von quadratischer auf log-lineare Komplexität reduziert. LLSA führt eine hierarchische Top-K-Auswahl durch, die schrittweise eine sparse Top-K-Auswahl mit den auf der vorherigen Ebene gefundenen Indizes anwendet, und führt einen Hierarchical-KV-Enrichment-Mechanismus ein, der den globalen Kontext bewahrt, während während der Attention-Berechnung weniger Token unterschiedlicher Granularität verwendet werden. Um effizientes Training zu unterstützen, entwickeln wir eine hochperformante GPU-Implementierung, die für Vorwärts- und Rückwärtsdurchläufe ausschließlich sparse Indizes verwendet und auf dichte Attention-Masken verzichtet. Wir evaluieren LLSA für die Bildgenerierung im hochauflösenden Pixelraum ohne Verwendung von Patchifizierung und VAE-Codierung. LLSA beschleunigt die Attention-Inferenz um das 28,27-fache und das DiT-Training um das 6,09-fache auf 256x256 Pixel Token-Sequenzen, bei gleichzeitiger Beibehaltung der Generierungsqualität. Die Ergebnisse demonstrieren, dass LLSA eine vielversprechende Richtung für das effiziente Training von DiTs mit langen Sequenzen bietet. Der Code ist verfügbar unter: 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