PaperlayerView exact revision on arXiv

Source-linked AI summary

Trainable Log-linear Sparse Attention for Efficient Diffusion Transformers

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

arXiv:2512.16615v2cs.CV

TL;DR

DiT self-attention and existing Top-K selection limit scaling to long token sequences, with single-level designs also relying on larger K to retain global context. LLSA uses hierarchical selection and KV enrichment to reduce attention complexity; on 256 × 256 pixel-token sequences, it improves training efficiency while maintaining generation quality.

  • Problem

    Full self-attention has quadratic cost, while single-level Top-K methods retain quadratic selection cost on compressed tokens and use larger K as sequences grow.

    Source: §1 · pp. 1–2, §3.2 · p. 3

  • Method

    LLSA applies coarse-to-fine hierarchical Top-K selection and enriches attention with coarse key/value representations to preserve global context at smaller K.

    Source: §1 · p. 2, §4.1 · p. 4, §4.2 · p. 5

  • Results

    On 256 × 256 pixel-token sequences, LLSA improves full-attention DiT training efficiency by 6.09× while maintaining generation quality.

    Source: §1 · p. 2

  • Takeaways & Limitations

    At K = 8, LLSA outperforms prior Top-K methods using K = 20 or K = 32 in quality and efficiency.

    Source: §1 · p. 2, §5.2 · p. 7

    The complexity analysis focuses on similarity-matrix computation; Top-K cost depends on implementation, though it remains O(N) when B and K are constant.

    Source: §4.3 · p. 5

Abstract

from arXiv · show

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

1. Introduction

Full self-attention and single-level Top-K selection limit DiT scaling to long sequences. LLSA uses hierarchical selection and KV enrichment to reduce attention cost while retaining global context, with an index-based GPU implementation and high-resolution pixel-space evaluation.

1. Introduction

  • Existing single-level Top-K attention retains quadratic selection cost on compressed tokens and requires larger K for longer sequences to maintain global context.The paper attributes both scaling limitations to a single coarse-grained view that cannot represent long-sequence global structure adequately.Source: §1 · p. 2
  • LLSA combines coarse-to-fine hierarchical Top-K selection with Hierarchical KV Enrichment to reduce attention complexity from quadratic to log-linear while preserving global context.Enriched coarse key/value representations let it use smaller K and lower attention cost.Source: §1 · p. 2
  • The GPU implementation operates directly on sparse indices in forward and backward passes, avoiding dense attention-mask construction and preserving end-to-end log-linear training complexity.It gathers selected key blocks in the forward pass and uses a sparse-index transpose kernel for backward lookup.Source: §1 · p. 2
  • 6.09× training-efficiency gain over full-attention DiTs on 256 × 256 pixel tokens, while maintaining generation quality.LLSA also reports higher quality and training throughput than prior Top-K methods, including at K = 8 versus their K = 20 or K = 32.Source: §1 · p. 2

2. Related Works

Prior DiT sparse-attention work uses learnable single-level block selection, while log-linear attention includes hierarchical and other key-selection strategies. Pixel-space DiTs address high-resolution generation through several architectural approaches, and the paper positions LLSA as an efficient alternative for this setting.

2.1. Sparse Attention in DiTs

  • Learnable DiT sparse-attention methods such as VMoBA, VSA, and SLA use single-level block selection; this work extends that approach to multiple hierarchical levels.Earlier strategies also used predefined patterns or dynamic search rules.Source: §2.1 · p. 2

2.2. Log-linear Attention

  • Log-linear attention methods use static positional rules or dynamic key selection; Multi-resolution Attention is the closest prior work, but lacks a high-performance GPU implementation validated on extremely long DiT sequences.LLSA instead uses block sparse attention and evaluates it in long-sequence DiT training.Source: §2.2 · pp. 2–3

2.3. Pixel-space DiTs

  • Pixel-space DiTs have used staged resolutions, coordinate-based decoding, or U-shaped downsampling; the paper evaluates LLSA on a lightweight pixel-space DiT without input downsampling.The paper states that prior work had not trained a pure pixel-space DiT with state-of-the-art high-resolution performance without input downsampling.Source: §2.3 · p. 3

3. Background

Block-sparse attention can skip query–key block pairs selected by a sparsity mask, while tiling avoids storing the full intermediate attention matrix. In Top-K sparse attention, compressed block representations guide selection, but the quadratic selection stage remains the bottleneck as sequences grow.

3.1. Block Sparse FlashAttention

  • FlashAttention uses tiling to compute output block by block rather than storing the full N × N intermediate matrix, and its sparse form skips block pairs marked invalid by a binary mask.Top-K selection is one approach to determining which query–key blocks can be skipped.Source: §3.1 · p. 3

3.2. Complexity of Top-K Sparse Attention

  • Top-K selection mean-pools query and key blocks, computes scores between compressed representations, then retains the K highest-scoring key blocks for each query block.The method separates selection on compressed tokens from attention on selected blocks.Source: §3.2 · p. 3
  • The Top-K selection stage has quadratic cost in sequence length, so it dominates the otherwise linear-in-N sparse-attention stage as sequences grow.The supplied analysis therefore identifies selection, not selected-block attention, as the scaling bottleneck in prior Top-K methods.Source: §3.2 · p. 3

4. Log-linear Sparse Attention Mechanism

LLSA targets reducing Top-K sparse attention from O(N^2) to O(N log N) through hierarchical selection and sparse GPU kernels. It progressively selects keys across levels and enriches fine-level attention with coarse tokens to preserve global context.

4.1. Log-linear Sparse Attention

  • Coarse-to-fine Top-K selection uses recursively pooled query and key representations, then adds selected coarse KV tokens to fine sparse attention to preserve global context.KV reweighting scales each coarse token's weight by its block size, W^(l) = B^l.Source: §4.1 · pp. 4–5

4.2. Complexity Analysis

  • With constant K, LLSA has overall O(N log N) complexity, versus O(N^2) for full attention and prior Top-K methods.Selection costs O(NK); attention over fine and enriched coarse candidates costs O(NK log N).Source: §4.2 · p. 5

4.3. Efficient Kernel Implementation

  • Sparse-index transposition avoids dense backward masks by mapping each key to its relevant queries, enabling sparse–dense accumulation for key and value gradients.The implementation stores query indices in a flat vector and uses cumulative key offsets to delimit each key's range.Source: §4.3 · pp. 5–6
  • Index ordering places similar pixels at neighboring 1D indices, unlike raster ordering, improving their grouping during 1D pooling.Source: Figure 2 · p. 6

5. Experiments

The experiments evaluate LLSA on pixel-space image generation, attention efficiency, and design ablations. LLSA outperforms the compared sparse-attention methods on generation benchmarks, while hierarchical selection and sparse-index backward processing support efficiency at long sequence lengths.

5.1. Pixel DiT Implementation

  • The pixel-space DiT setup reorders neighboring pixels into adjacent sequence tokens, rescales noise for images above 64 × 64, and initializes high-resolution models from low-resolution weights.Quality is evaluated with FID on 10,000 samples generated in 20 diffusion steps; throughput is measured on one H200 GPU.Source: §5.1 · p. 6

5.2. Ablation Study

  • The ablations find that KV Enrichment improves FID, correct coarse-token importance weighting further improves quality, and multi-level attention raises throughput with a minor quality drop.At K = 8, LLSA outperforms the single-level K = 32 baseline in quality and efficiency; larger blocks are faster but yield worse quality than smaller blocks.Source: §5.2 · pp. 6–7

5.3. Image Generation Benchmark

  • Across FFHQ-128 and FFHQ-256, LLSA achieves the best FID and highest training throughput; on ImageNet-256, it outperforms VSA and SLA in FID and Inception Score.Table 2 reports FID and throughput; Table 3 reports FID, Inception Score, and throughput.Source: §5.3 · p. 7, Table 2 · p. 7, Table 3 · p. 7

5.4. Efficiency Benchmark

  • For long sequences, switching from L = 1 to L = 2 improves efficiency, and LLSA consistently outperforms VSA and SLA across the tested settings.Figure 3 compares training and inference acceleration against PyTorch Attention (FlashAttention2) across sequence lengths and block sizes on an H200 GPU.Source: §5.4 · pp. 7–8, Figure 3 · p. 8
  • LLSA’s sparse key–value backward kernel maintains nearly constant throughput across sequence lengths, while the mask-based baseline degrades as sequences grow.The baseline constructs and scans a dense block mask, an O(N^2) operation.Source: §5.4 · p. 8, Figure 4 · p. 8

6. Conclusion

LLSA reduces self-attention complexity from O(N^2) to O(N log N) through hierarchical search and Hierarchical KV Enrichment. In high-resolution pixel-space DiT experiments, it achieves better FID and higher training throughput than prior sparse-attention methods.

6. Conclusion

  • LLSA reduces self-attention complexity from O(N^2) to O(N log N) using hierarchical search and Hierarchical KV Enrichment to preserve global context and model quality under long-sequence sparsification.Its GPU implementation uses block-sparse attention and sparse-index backward processing without dense masks.Source: §6 · p. 8

A.1. Clarification of Kernel Implementation

The appendix contrasts sparse-index backward processing with mask-based alternatives and describes kernel-efficiency comparisons and training configurations.

A.1. Clarification of Kernel Implementation

  • SLA uses mask-based sparse backward, while VSA still constructs a binary mask; the authors reimplement both with sparse-index transpose for a conservative throughput comparison.Under this setting, LLSA remains more efficient than SLA and VSA for an equal number of effective sparse tokens.Source: §A.1 · p. 10
  • Figure 3 compares LLSA’s kernel efficiency with official SLA and VSA implementations, with block-size configurations adjusted to support B = 16 inference.The appendix provides detailed FFHQ and ImageNet training configurations in Table 4.Source: §A.1 · p. 10, §A.2 · p. 10

A.2. Training Configuration

  • Detailed training configurations for FFHQ and ImageNet are provided in Table 4.These settings are identified separately from the kernel-implementation clarification.Source: §A.2 · p. 10

B.1. Additional Ablation Results

The additional ablations examine enrichment depth, scaling to 512 × 512, SNR adjustment, index ordering, and low-resolution pretraining. They report improved quality or throughput for the tested choices, with the 512 × 512 throughput trend following LLSA’s O(N log N) complexity.

More enrichment levels

  • More KV enrichment levels increase effective token count and generation quality while slightly reducing throughput.The comparison trains three two-level LLSA DiTs on FFHQ-128 with different enrichment-layer counts.Source: §B.1 · p. 11

512 × 512 pixel token sequence generation

  • 375.34 to 323.29 per-token throughput from 256 × 256 to 512 × 512 closely follows LLSA’s O(N log N) complexity.At 512 × 512, L = 1 fails to converge within a reasonable time budget; L = 2 improves throughput, and L = 3 adds speed gains.Source: §B.1 · p. 11

Noise rescaling

  • Noise rescaling yields the best generation quality among the tested SNR-adjustment methods.The comparison includes timestep shift, logit-normal sampling, and noise rescaling.Source: §B.1 · p. 11

Index reordering

  • Index reordering improves model quality; without it, the model converges but achieves worse FID.The comparison uses the default raster-scan ordering for 2D pixels.Source: §B.1 · p. 11

Pretraining

  • Low-resolution pretraining substantially reduces high-resolution pixel-space DiT training cost, with the pretrained model converging within the first epoch.Figure 5 compares FID curves for training from scratch and initializing from a lower-resolution pretrained checkpoint.Source: §B.1 · p. 11, Figure 5 · p. 11

B.2. Training Curves

On PixelFlow ImageNet-256, LLSA’s generation quality is consistently better than baseline attention approaches throughout training, as tracked by FID and Inception Score.

Training curves

  • LLSA’s generation quality is consistently better than baseline attention approaches throughout training.The curves report FID and Inception Score on PixelFlow ImageNet-256.Source: §B.2 · p. 11
  • The comparison tracks FID and Inception Score on the PixelFlow ImageNet-256 benchmark.Source: §B.2 · p. 11
  • LLSA maintains its reported quality advantage across the training process, rather than only at a single point.Source: §B.2 · p. 11

C. Qualitative Results

The qualitative results show LLSA samples at three FFHQ resolutions and compare ImageNet-256 samples from LLSA with SLA, VSA, and a well-trained PixelFlow reference. The FFHQ-512 model was trained for only two epochs.

FFHQ samples

ImageNet-256 samples

Related papers

Loading 2512.16615v2…