Technical Summary: PISA - Piecewise Sparse Attention for Efficient Diffusion Transformers

1. Problem Statement

Diffusion Transformers (DiTs) have become the standard for high-fidelity image and video generation. However, their scalability is bottlenecked by the quadratic complexity ($O(L^2)$) of the self-attention mechanism as input sequence lengths ($L$) increase for higher resolutions and longer videos.

Existing solutions face significant limitations:

• Sparse Attention ("Keep-or-Drop"): Methods that discard non-critical key-value (KV) blocks accelerate computation but suffer from performance degradation at high sparsity ratios due to the loss of context. They also struggle with shorter sequences where the overhead of selection outweighs benefits.
• Linear/Hybrid Attention: Approaches that alter the attention distribution (e.g., via feature mappings or additive combinations of sparse and linear branches) often disrupt the intrinsic normalization of pre-trained models. This necessitates expensive retraining or fine-tuning, preventing the direct reuse of pre-trained weights.
• Hard Truncation: Directly dropping blocks leads to sharp deviations from the original attention distribution, causing quality loss in generated content.

The core challenge is to design a mechanism that accelerates inference with sub-quadratic complexity while maintaining the accuracy of full attention and remaining compatible with pre-trained weights without retraining.

2. Methodology: Piecewise Sparse Attention (PISA)

PISA introduces a novel "exact-or-approximate" paradigm, moving beyond the traditional "keep-or-drop" strategy. It treats attention as a piecewise process that covers the full attention span with sub-quadratic complexity.

Key Insights

The authors analyzed pre-trained DiTs and discovered two critical properties:

1. Distributional Stability: The pre-softmax attention scores of non-critical (uncritical) blocks exhibit a symmetric, bell-shaped distribution centered around zero or negative values. This makes them highly amenable to accurate approximation via Taylor expansion, rather than requiring exact computation or discarding.
2. Normalization Consistency: To maintain training-free compatibility, approximations must be integrated internally into the softmax numerator and denominator, rather than simply adding outputs from different branches (which violates weighted sum rules).

Core Algorithm

PISA partitions the KV blocks into two sets:

1. Critical Blocks ($S_i$): Computed exactly to preserve high-salience information.
2. Non-Critical Blocks ($U_i$): Approximated efficiently using block-wise Taylor expansion.

The output $o_t$ is derived by normalizing the weighted value aggregation:
$$o_t = \frac{N_t}{D_t}$$

Where the denominator ($D_t$) and numerator ($N_t$) are composed of:

• Exact Sparse Terms: Computed for selected blocks $j \in S_i$.
• Block-wise Zeroth-Order Approximation: For unselected blocks, the contribution is approximated using the block centroid $\bar{k}_j$ and the sum of values.
• Global First-Order Correction: Instead of computing a distinct first-order term for every unselected block (which is memory-bound), PISA employs a hybrid-order approximation. It uses a global statistic $\bar{H}$ (the mean of block-wise first-order matrices across all blocks) to approximate the first-order gradient for the entire tail set.

The formulation for the numerator $N_t$ becomes:
$$N_t = \sum_{j \in S_i} \sum_{n} e^{q_t k_{j,n}^\top} v_{j,n} + \sum_{j \in U_i} \alpha_{t,j} \left(\sum_{n} v_{j,n}\right) + q_t \bar{H} \sum_{j \in U_i} \alpha_{t,j}$$
Where $\alpha_{t,j} = \exp(q_t \bar{k}_j^\top)$.

Covariance-Aware Routing

To minimize approximation error, PISA uses a Covariance-Aware Top-k Strategy. The selection score for a block incorporates not only the attention score but also the norm of the block's covariance matrix ($M_j = \|H_j - \bar{H}\|_2$). Blocks with high semantic relevance or high heterogeneity are selected for exact computation, while smoother blocks are delegated to the Taylor expansion.

Hardware Implementation

PISA is implemented via a fused kernel that interleaves exact computation and approximate scanning within the online softmax process (FlashAttention style). The pipeline involves:

1. Pre-computation: Calculating block centroids, value sums, and the global statistic $\bar{H}$.
2. Phase 1 (Exact): Loading selected blocks into SRAM for exact attention.
3. Phase 2 (Approximation): Scanning aggregated block vectors for unselected blocks using zeroth-order expansion.
4. Phase 3 (Correction): Injecting the global first-order correction term.

This design ensures the method is memory-efficient and avoids the bottlenecks associated with loading distinct matrices for every block.

3. Key Contributions

1. Novel Piecewise Sparse Attention: A training-free mechanism that enables full attention span coverage with sub-quadratic complexity by unifying exact sparse computation and block-wise approximation.
2. Hybrid-Order Approximation Scheme: A strategy that combines block-wise zeroth-order expansion with a global first-order correction. This significantly boosts accuracy with negligible computational cost compared to standard sparse attention.
3. Covariance-Aware Routing: A derivation from error analysis that minimizes approximation divergence by selecting blocks based on both attention scores and block covariance norms.
4. State-of-the-Art Performance: Demonstrated superior efficiency and quality across diverse tasks (video and image generation) without requiring retraining.

4. Experimental Results

The authors evaluated PISA on video generation models (Wan2.1-1.3B/14B, Hunyuan-Video) and image generation models (SD 3.5, FLUX.1).

• Video Generation:

  • Wan2.1-14B: Achieved a 1.91× speedup over dense attention (1468s $\to$ 818s) at 87.5% sparsity, while maintaining or slightly improving quality metrics (VBench, PSNR, SSIM) compared to dense attention.
  • Hunyuan-Video: Achieved a 2.57× speedup (1651s $\to$ 641s) with consistent quality preservation.
  • Comparison: Outperformed existing sparse methods like SpargeAttn and SVG2 in both speed and quality, particularly in preserving fine-grained details and temporal consistency.

• Image Generation:

  • FLUX.1-dev: Achieved a 1.2× speedup (8.32s $\to$ 6.87s) at 85% sparsity. Notably, PISA achieved higher scores on human preference metrics (ImageReward, HPSv2) and similarity metrics (PSNR, SSIM) compared to both dense attention and SpargeAttn.
  • SD 3.5: Demonstrated significant speedups (e.g., 1.4× on SD 3.5-M) while maintaining FID scores comparable to or better than dense attention.

• Kernel Efficiency:

  • PISA consistently outperformed FlashAttention-2/3 and SpargeAttn in latency across various sequence lengths (4K to 32K) and densities.
  • Unlike SpargeAttn, which degrades on shorter sequences, PISA maintained speed advantages even at 4K tokens.

5. Significance and Claims

The paper claims that PISA represents a paradigm shift from "keep-or-drop" to "exact-or-approximate" sparse attention. Its primary significance lies in:

• Bridging the Speed-Quality Gap: It effectively bridges the gap between inference speed and generation quality, achieving near-lossless approximation with negligible overhead.
• Training-Free Compatibility: By integrating approximations directly into the softmax normalization framework, PISA allows for the direct reuse of pre-trained weights, eliminating the need for expensive retraining or fine-tuning required by linear or hybrid attention methods.
• Scalability: It provides a unified solution for both high-resolution image and long-duration video generation, addressing the quadratic complexity bottleneck without sacrificing the fidelity of the generated content.

The authors emphasize that their method is the first to natively integrate block-wise Taylor expansion and sparse attention into a unified framework that preserves the intrinsic distribution of pre-trained models.