Attenzione Sparsa Log-Lineare Addestrabile per Trasformatori di Diffusione Efficienti

Trainable Log-linear Sparse Attention for Efficient Diffusion Transformers

December 18, 2025

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

cs.AI

Abstract

I Diffusion Transformer (DiT) stabiliscono lo stato dell'arte nella generazione visiva, tuttavia il loro costo quadratico di self-attention limita fondamentalmente la scalabilità a sequenze lunghe di token. I recenti approcci di attenzione sparsa Top-K riducono il calcolo dei DiT comprimendo i token in rappresentazioni basate su blocchi e selezionando un piccolo insieme di blocchi chiave rilevanti, ma soffrono ancora di (i) un costo di selezione quadratico sui token compressi e (ii) la necessità di aumentare K per mantenere la qualità del modello man mano che le sequenze crescono. Identifichiamo che la loro inefficienza è dovuta al design a singolo livello, poiché un unico livello grossolano è insufficiente per rappresentare la struttura globale. In questo articolo, introduciamo la Log-linear Sparse Attention (LLSA), un meccanismo di attenzione sparsa addestrabile per sequenze di token estremamente lunghe che riduce sia i costi di selezione che di attenzione da una complessità quadratica a una log-lineare, utilizzando una struttura gerarchica. LLSA esegue una selezione Top-K gerarchica, adottando progressivamente una selezione Top-K sparsa con gli indici trovati al livello precedente, e introduce un meccanismo di Hierarchical KV Enrichment che preserva il contesto globale utilizzando meno token di diversa granularità durante il calcolo dell'attenzione. Per supportare un addestramento efficiente, sviluppiamo un'implementazione GPU ad alte prestazioni che utilizza solo indici sparsi sia per i passi in avanti che all'indietro, eliminando la necessità di maschere di attenzione dense. Valutiamo LLSA sulla generazione di immagini ad alta risoluzione nello spazio dei pixel senza utilizzare patchificazione e codifica VAE. LLSA accelera l'inferenza dell'attenzione di 28.27x e l'addestramento del DiT di 6.09x su sequenze di token di pixel 256x256, mantenendo la qualità di generazione. I risultati dimostrano che LLSA offre una direzione promettente per addestrare efficientemente DiT a sequenza lunga. Il codice è disponibile all'indirizzo: 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