PaperlayerView exact revision on arXiv

Source-linked AI summary

PISA: Piecewise Sparse Attention Is Wiser for Efficient Diffusion Transformers

Haopeng Li, Shitong Shao, Wenliang Zhong, Zikai Zhou, Lichen Bai, Hui Xiong, Zeke Xie

arXiv:2602.01077v2cs.CV

TL;DR

Attention’s quadratic cost limits efficient Diffusion Transformers, while dropping non-critical blocks can degrade quality. PISA instead computes critical blocks exactly and approximates the remainder, achieving 1.91× and 2.57× speedups while preserving state-of-the-art quality.

  • Problem

    Diffusion Transformers face inefficient inference as growing token sequences make self-attention’s quadratic complexity a significant bottleneck, while sparse attention can lose performance by discarding context.

  • Method

    PISA uses unified exact-or-approximate execution, computing critical blocks exactly and approximating remaining blocks to cover the full attention span with sub-quadratic complexity.

  • Results

    1.91× and 2.57× speedups on Wan2.1-14B and Hunyuan-Video, respectively, are achieved while preserving state-of-the-art quality.

  • Takeaways & Limitations

    PISA provides a training-free sparse-attention approach that maintains full attention coverage while balancing inference efficiency and accuracy.

  • Takeaways & Limitations

    The first-order approximation term can be severely bottlenecked by memory access in practical implementations.

Abstract

from arXiv · show

Diffusion Transformers are fundamental for video and image generation, but their efficiency is bottlenecked by the quadratic complexity of attention. While block sparse attention accelerates computation by attending only critical key-value blocks, it suffers from degradation at high sparsity by discarding context. In this work, we discover that attention scores of non-critical blocks exhibit distributional stability, allowing them to be approximated accurately and efficiently rather than discarded, which is essentially important for sparse attention design. Motivated by this key insight, we propose PISA, a training-free Piecewise Sparse Attention that covers the full attention span with sub-quadratic complexity. Unlike the conventional keep-or-drop paradigm that directly drop the non-critical block information, PISA introduces a novel exact-or-approximate strategy: it maintains exact computation for critical blocks while efficiently approximating the remainder through block-wise Taylor expansion. This design allows PISA to serve as a faithful proxy to full attention, effectively bridging the gap between speed and quality. Experimental results demonstrate that PISA achieves 1.91 times and 2.57 times speedups on Wan2.1-14B and Hunyuan-Video, respectively, while consistently maintaining the highest quality among sparse attention methods. Notably, even for image generation on FLUX, PISA achieves a 1.2 times acceleration without compromising visual quality. Code is available at: https://github.com/xie-lab-ml/piecewise-sparse-attention.

1. Introduction

PISA addresses the quadratic attention bottleneck in Diffusion Transformers with training-free piecewise exact-or-approximate computation that preserves the full attention span. It achieves substantial speedups while maintaining high or state-of-the-art quality across image and video generation tasks.

  • Diffusion Transformers face quadratic self-attention costs as higher-resolution images and longer videos dramatically increase input-token sequence lengths.
  • Existing sparse attention methods can discard key-value context at high sparsity and alter pretrained attention distributions, requiring retraining and limiting efficiency on shorter sequences.
  • PISA preserves critical blocks through exact computation and covers non-critical blocks using block-wise Taylor approximation, enabling full attention span with sub-quadratic complexity.
  • 1.91× and 2.57× speedups are achieved on Wan2.1-14B and Hunyuan-Video, respectively, while preserving state-of-the-art quality.
  • PISA uses hybrid-order approximation and covariance-aware routing to improve accuracy and minimize approximation divergence at negligible cost.

2. Related Work

Prior efficient-attention methods reduce complexity through block selection, kernel feature mappings, or hybrid branches, but introduce limitations involving discarded context, approximation scope, or disrupted normalization. PISA instead addresses these shortcomings through its exact-or-approximate attention design.

  • Block Sparse Attention: Block-sparse attention limits computation to critical key-value blocks, but discarding unselected blocks increases output error and degrades performance at high sparsity.Existing methods use static priors or dynamic routing to select blocks.
  • Bidirectional Linear Attention: Bidirectional linear attention achieves linear complexity by replacing the exponential kernel with feature mappings, while Taylor-based methods typically expand around zero globally.Vision methods mainly modify feature mappings while retaining the canonical linear-attention formulation.
  • Native Hybrid Attention: Native hybrid attention combines sparse or sliding-window attention with linear attention, but additive branch outputs disrupt the intrinsic normalization of attention weights.SLA and NHA exemplify this additive strategy.

3. Methodology

PISA replaces keep-or-drop sparse attention with an exact-or-approximate design that computes critical blocks exactly and approximates unselected blocks within online softmax. A global first-order correction improves hardware efficiency while error analysis motivates covariance-aware block selection.

  • Key Insights: Uncritical blocks have pre-softmax scores centered near zero or negative values, making mean-centered Taylor approximation accurate.Figure 3 shows a symmetric bell-shaped block-score distribution and divergence between uncritical and important blocks.
  • PISA Formulation: PISA unifies exact sparse computation with block-wise approximation directly within online softmax, preserving the full attention span without training.Selected key-value blocks are computed exactly, while unselected blocks are approximated rather than discarded.
  • Practical Challenge: A block-wise first-order approximation is theoretically efficient but becomes memory-bound because each unselected block requires distinct matrix handling.The resulting low arithmetic intensity prevents the theoretical speedup from being realized in practice.
  • Global First-Order Correction: Global first-order correction shares a unified scale factor across all unselected blocks, eliminating block-specific weighted summation.The formulation combines an exact sparse term, block-wise zeroth-order approximation, and global first-order correction.
  • Error Analysis: Theorem 3.1 bounds replacement error by the product of tail probability mass and block-matrix variance, keeping error controlled when unselected-block weights are small.This establishes the approximation’s error dependence on both tail mass and block heterogeneity.
  • Covariance-Aware Block Selection: Covariance-aware Top-k selection preserves blocks with high semantic relevance or covariance heterogeneity, while delegating smoother blocks to Taylor expansion.The strategy uses block covariance norms as an importance-routing prior to reduce approximation error.

4. Experiments

Experiments show that PISA delivers consistent efficiency gains across sequence lengths and densities while preserving or improving generation quality in video and image tasks. Ablations further validate that approximation reduces attention error and that hybrid approximation best preserves fine-grained details.

  • Efficiency: At 12.5% density, PISA consistently outperforms FA3 and SpargeAttn, retaining a speed advantage over FA3 even at 4K sequences.SpargeAttn becomes slower than FA3 at shorter sequence lengths.
  • Efficiency: PISA outperforms FA2 even above 70% density, and surpasses FA3 below 50% density for sequences over 8K or below 70% density at 4K.These results demonstrate efficiency across varying densities and sequence lengths.
  • Video Generation: PISA achieves state-of-the-art VBench performance with substantial speedups, sometimes surpassing full attention on generation-quality metrics.Compared with SVG2 and SpargeAttn, PISA maintains visual quality consistent with full attention while avoiding blurring, flickering, and fine-detail loss.
  • Video Generation: PISA maintains high-fidelity video generation without warmup, unlike competing methods that deteriorate precipitously.This supports the advantage of piecewise computation covering the full attention span.
  • Image Generation: On SD3.5 and FLUX.1, PISA outperforms SparseAttn in FID, human-preference benchmarks, similarity, and speedup at similar or higher sparsity.On FLUX.1-dev, PISA also surpasses full attention on certain metrics.
  • Ablation and Validation: Hybrid approximation lowers normalized attention error beyond zeroth-order approximation and improves similarity toward the exact baseline while recovering fine-grained visual details.The approximation phase adds minimal runtime overhead, supporting the efficiency–fidelity trade-off.

5. Conclusion

PISA replaces sparse attention’s conventional “keep-or-drop” paradigm with unified exact-or-approximate execution, preserving the full attention span at sub-quadratic complexity while balancing efficiency and accuracy.

  • 5. Conclusion: PISA replaces the conventional “keep-or-drop” paradigm with unified “exact-or-approximate” execution.This design underpins the proposed Piecewise Sparse Attention method.
  • 5. Conclusion: PISA enables the full attention span with sub-quadratic complexity, achieving an optimal trade-off between efficiency and accuracy.

A. Expanded Related Work

Prior work reduces DiT attention cost through sparse attention, linear attention, Taylor-expansion approximations, and hybrid attention designs. These approaches vary in how they select key-value pairs, approximate softmax attention, and combine attention mechanisms.

  • Sparse Attention for DiTs: Sparse attention reduces complexity by selectively computing only critical key-value pairs in diffusion transformers.Video-generation methods include static masks based on spatio-temporal sparsity priors and dynamic selection strategies.
  • Bidirectional Linear Attention: Linear attention achieves linear complexity by decomposing exp(qk⊤) into ϕ(q)ϕ(k)⊤ and computing k⊤v first through associativity.Early formulations commonly used non-negative kernels such as ReLU or ELU(·) + 1 to preserve valid normalization.
  • Linear Attention with Taylor Expansion: LLM research uses Taylor expansion to linearize softmax attention, with Based and power attention adopting second-order expansions through exact polynomial feature maps.The cited feature map satisfies ϕp(q)ϕp(k) = (q⊤k)p for second-order terms.
  • Hybrid Attention: Hybrid attention combines softmax and linear attention across or within layers, typically by weighted output summation or head-wise allocation followed by projection.Recent work has also explored hybridization at a finer granularity.

B. Compare with Other Methods

PISA differs from prior methods through an exact-or-approximate paradigm, rigorous block-wise Taylor mixing, and local-mean expansion that enables training-free reuse of pre-trained weights. It unifies block-wise Taylor expansion and sparse attention in one framework.

  • Exact-or-approximate paradigm: PISA replaces pure sparse attention’s keep-or-drop paradigm with exact computation for critical blocks and efficient approximation for the remainder.This blends sparse exact computation with approximation instead of discarding most key-value pairs.
  • Hybrid attention mixing: PISA uses block-wise Taylor expansion to approximate softmax attention’s numerator and denominator within a canonical normalization framework.This provides a fine-grained, mathematically rigorous alternative to naively summing softmax and linear attention outputs.
  • Local expansion: PISA expands around each block’s local mean rather than globally, reducing approximation error and enabling direct training-free inheritance of pre-trained weights with first-order expansion.Prior global expansion causes prohibitive errors that prevent pre-trained-weight reuse.
  • Unified framework: PISA is presented as the first framework to natively integrate block-wise Taylor expansion with sparse attention.The method combines both components within a unified framework.

C. Error Analysis of the Hybrid-Order Approximation

The hybrid-order approximation replaces exact first-order contributions from unselected blocks with a global-mean approximation and bounds the resulting attention error. The bound separates block heterogeneity from unselected attention mass, while Taylor truncation introduces an additional residual error.

  • Hybrid-order approximation: The approximation replaces the exact first-order contribution of unselected blocks with a global statistic while retaining the same attention denominator.This isolates the error caused by global-mean replacement from denominator changes.
  • Residual bound: The resulting numerator error is controlled through a residual matrix measuring deviations between block-wise first-order contributions and their global mean.The proof uses triangle inequality and operator-norm submultiplicativity.
  • Error decomposition: The block size B reduces the bound because Jensen’s inequality makes each block-centroid exponential no larger than the average of its B per-row exponentials.This yields the block-size dependence in the denominator of the bound.
  • Error decomposition: Theorem C.2 separates the error bound into structural heterogeneity M and tail mass ρt = τt/Dt from unselected blocks.M measures block-to-global-statistic dissimilarity, whereas ρt measures the attention weight assigned to approximated blocks.
  • Limitations: The theorem covers only global-mean replacement, while Taylor-series truncation contributes an additional error governed by the second-order moment of block deviations.This residual is expected to be small on unselected blocks because their attention distribution is sparse and flat.

D. Covariance-Aware Block Selection

PISA selects blocks by combining attention scores with covariance-based approximation error, using a normalized routing factor whose mean term cancels during softmax normalization.

  • Covariance-Aware Block Selection: Theoretical error scales with exp(qk^T) · ||H||_2, motivating selection of blocks with both large attention scores and high approximation errors.The method rectifies attention scores using M_j := ||H_j − H̄||.
  • Covariance-Aware Block Selection: The block-wise routing factor is exp(q_t H̄k^T) · M_j/M̄, normalizing covariance deviation M_j by its mean over U.This removes the influence of M_j’s absolute magnitude.
  • Covariance-Aware Block Selection: Taking the routing factor’s logarithm integrates it into softmax, while the constant mean term cancels during normalization to yield Eq. (11).The cancellation follows because M̄ is constant.

E. Implementation Details for Reproducibility

The implementation uses explicitly documented warmup configurations while following official settings elsewhere to ensure fair, reproducible comparisons across video and image generation models.

  • E. Implementation Details for Reproducibility: All other sparse-attention parameters follow the corresponding official implementations for strictly fair comparison.The paper states that these parameters adhere to official implementations beyond the warmup configuration.
  • E. Implementation Details for Reproducibility: Remaining generation hyperparameters use official recommended configurations, including sampling steps and classifier-free guidance (CFG) scale.These settings are applied when integrating different attention methods into video and image generation models, without further enumeration.

F. More Generation Samples

Additional samples illustrate PISA’s image and video generation under high sparsity. Figure 10 covers FLUX.1-dev text-to-image generation, while Figure 11 covers Wan2.1-14B text-to-video generation.

  • Additional Samples: The paper provides additional image and video generation samples in Figures 10 and 11.The examples include prompts such as a photorealistic grizzly in a forest fire and a teddy bear playing drums in Times Square.
  • Image Generation: Figure 10 presents FLUX.1-dev text-to-image samples, comparing Dense Attention, PISA, and SpargeAttn.SpargeAttn uses 80% sparsity, while PISA uses 85% sparsity.
  • Video Generation: Figure 11 presents Wan2.1-14B text-to-video samples generated with PISA at 87.5% sparsity.The configuration uses 10 steps and 1 layer warmup.
Loading 2602.01077v2…