Atención Dispersa Log-Lineal Entrenable para Transformadores de Difusión 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

Resumen

Los Transformadores de Difusión (DiTs) establecen el estado del arte en generación visual, pero su coste cuadrático de autoatención limita fundamentalmente la escalabilidad a secuencias largas de tokens. Los enfoques recientes de atención dispersa Top-K reducen la computación de los DiTs comprimiendo los tokens en representaciones por bloques y seleccionando un pequeño conjunto de bloques clave relevantes, pero aún adolecen de (i) un coste cuadrático de selección sobre los tokens comprimidos y (ii) la necesidad de aumentar K para mantener la calidad del modelo a medida que las secuencias crecen. Identificamos que su ineficiencia se debe al diseño de un solo nivel, ya que un único nivel grueso es insuficiente para representar la estructura global. En este artículo, presentamos la Atención Dispersa Log-Lineal (LLSA), un mecanismo de atención dispersa entrenable para secuencias de tokens extremadamente largas que reduce tanto los costes de selección como de atención de una complejidad cuadrática a log-lineal mediante la utilización de una estructura jerárquica. LLSA realiza una selección Top-K jerárquica, adoptando progresivamente una selección Top-K dispersa con los índices encontrados en el nivel anterior, e introduce un mecanismo de Enriquecimiento KV Jerárquico que preserva el contexto global mientras utiliza menos tokens de diferente granularidad durante el cálculo de la atención. Para soportar un entrenamiento eficiente, desarrollamos una implementación de GPU de alto rendimiento que utiliza únicamente índices dispersos tanto para las pasadas hacia adelante como hacia atrás, eliminando la necesidad de máscaras de atención densas. Evaluamos LLSA en la generación de imágenes de alta resolución en espacio de píxeles sin utilizar `patchification` ni codificación VAE. LLSA acelera la inferencia de atención en 28.27x y el entrenamiento de DiT en 6.09x en secuencias de tokens de 256x256 píxeles, manteniendo la calidad de la generación. Los resultados demuestran que LLSA ofrece una dirección prometedora para entrenar DiTs de secuencias largas de manera eficiente. El código está disponible en: 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