APT: Accelerating Diffusion Transformers via Attention Probability-Guided Pruning and Quantization
APT accelerates Diffusion Transformers via attention probability-guided pruning and quantization, achieving up to 8.16× speedup and 14.98× energy efficiency improvement.
Key Findings
Methodology
APT combines software-hardware co-design with Attention Probability-guided Adaptive Dual Thresholding (APDT) and Timestep-Aware FlashAttention (TAFA). APDT dynamically prunes and quantizes elements based on attention probabilities, while TAFA predicts attention probabilities across timesteps to maintain memory efficiency. The hardware architecture supports irregular sparsity and dual-precision execution with a tile-based dataflow.
Key Results
- APT achieves up to 8.16× speedup and 14.98× energy efficiency improvement on PixArt-α, Stable Diffusion 3, and FLUX, outperforming EXION by 3.01× speedup and 2.04× energy efficiency.
- APDT reduces computational burden by 65% while maintaining a quantization signal-to-noise ratio (QSNR) of 28.5 dB, comparable to 12-bit precision.
- TAFA achieves >97% similarity in execution masks compared to full Softmax while reducing memory overhead by up to 103.4×.
Significance
APT addresses the computational bottleneck of self-attention in high-resolution Diffusion Transformers, enabling efficient deployment of generative AI models. Its techniques can generalize to other Transformer-based architectures requiring sparsity and precision control.
Technical Contribution
APT introduces APDT for joint pruning and quantization, and TAFA to integrate with FlashAttention. Its hardware design supports irregular sparsity and dual-precision execution, significantly enhancing efficiency for high-resolution generative tasks.
Novelty
APT is the first to leverage attention probabilities for joint pruning and quantization. TAFA innovatively addresses FlashAttention's memory limitations, and the hardware design uniquely supports tile-based execution with dynamic sparsity.
Limitations
- APT's performance relies on temporal similarity in attention probabilities, limiting its applicability to models with low timestep correlation.
- The hardware design's complexity may increase development and deployment costs.
- Currently validated only on generative tasks, its effectiveness in other Transformer applications remains untested.
Future Work
Future work could explore applying APT to other Transformer models, such as those in natural language processing. Optimizing the hardware design for higher resolutions and more complex sparsity patterns is another promising direction.
AI Executive Summary
Diffusion Transformers (DiTs) excel in high-resolution generative tasks but face severe computational bottlenecks due to the O(N^4) complexity of self-attention. APT addresses this challenge through Attention Probability-guided Adaptive Dual Thresholding (APDT) and Timestep-Aware FlashAttention (TAFA), significantly reducing computational overhead. Its hardware architecture supports dynamic sparsity and dual-precision execution with a tile-based dataflow.
Experiments on PixArt-α, Stable Diffusion 3, and FLUX demonstrate up to 8.16× speedup and 14.98× energy efficiency improvement over NVIDIA A100, with 3.01× speedup and 2.04× energy efficiency over EXION. APDT achieves a 65% reduction in computation while maintaining a QSNR of 28.5 dB, comparable to 12-bit precision. TAFA generates execution masks with >97% similarity to full Softmax while reducing memory requirements by 103.4×.
APT provides an efficient solution for high-resolution generative tasks and opens new directions for hardware design in sparsity and precision control. However, its reliance on temporal similarity and hardware complexity pose challenges, warranting further exploration in broader applications and design optimizations.
Deep Analysis
Background
Diffusion models have become central to generative AI, producing high-quality images and videos through iterative denoising. Diffusion Transformers (DiTs) have emerged as the standard for high-resolution tasks due to their scalability and output quality. However, their self-attention operations scale with O(N^4), creating significant computational challenges.
Core Problem
The primary bottleneck in DiTs at high resolutions is the computational cost of self-attention. Existing sparsity and quantization techniques fail to balance efficiency and accuracy, particularly in high-resolution scenarios.
Innovation
APT introduces: 1) Attention Probability-guided Adaptive Dual Thresholding (APDT) for dynamic pruning and quantization; 2) Timestep-Aware FlashAttention (TAFA) to predict attention probabilities using temporal similarity; 3) a hardware architecture supporting dynamic sparsity and dual-precision execution.
Methodology
- �� APDT dynamically adjusts pruning and quantization thresholds based on the standard deviation of attention probabilities.
- �� TAFA reuses normalization statistics from previous timesteps to predict attention probabilities, avoiding full Softmax computation.
- �� The hardware architecture employs a tile-based dataflow, supporting irregular sparsity and dual-precision execution.
Experiments
Experiments were conducted on PixArt-α, Stable Diffusion 3, and FLUX, comparing APT against NVIDIA A100 and EXION. Metrics included speedup, energy efficiency, and QSNR. Ablation studies evaluated the impact of pruning and quantization.
Results
APT achieved up to 8.16× speedup and 14.98× energy efficiency improvement, with a QSNR of 28.5 dB. TAFA's execution masks showed >97% similarity to full Softmax, significantly reducing memory requirements.
Applications
APT is ideal for high-resolution image and video generation, particularly in resource-constrained environments. Its techniques can extend to other Transformer-based models requiring sparsity and precision control.
Limitations & Outlook
APT's reliance on temporal similarity limits its applicability to models with low timestep correlation. The hardware design's complexity may increase development costs.
Plain Language Accessible to non-experts
Imagine APT as a smart delivery system. Traditional systems deliver every package (compute every element), but APT analyzes each package's importance (attention probability) to decide which to deliver first, which to simplify (quantize), and which to skip (prune). It also uses yesterday's delivery data (temporal similarity) to optimize today's routes (TAFA). This drastically improves efficiency while maintaining quality.
ELI14 Explained like you're 14
Think of APT like playing a huge puzzle game where you can only place one piece at a time. APT is like a super-smart helper that tells you which pieces are most important to place first, which can be simplified, and which can wait. Plus, it remembers what you did in earlier steps to make the next ones faster. Cool, right?
Glossary
Attention Probability
Indicates the importance of each element in self-attention, guiding pruning and quantization.
Used in APT to decide which elements require high precision.
Pruning
Removes less important elements to reduce computation.
APT prunes elements with low attention probabilities.
Quantization
Reduces precision of data to lower computational cost.
APT uses 6-bit and 12-bit quantization for efficiency.
FlashAttention
A memory-efficient self-attention algorithm using tiled computation.
APT extends FlashAttention with TAFA for compatibility.
Tile-based Dataflow
Breaks computation into smaller tiles for higher hardware efficiency.
Core to APT's hardware design.
Open Questions Unanswered questions from this research
- 1 APT's performance in non-generative tasks remains unexplored, requiring further validation.
- 2 Optimizing hardware for more complex sparsity patterns is an open challenge.
- 3 Models with low temporal similarity may limit APT's effectiveness.
Applications
Immediate Applications
High-Resolution Image Generation
Accelerates models like PixArt-α and Stable Diffusion for art creation and advertising.
Video Generation and Editing
Speeds up high-resolution video generation for filmmaking and VR applications.
Long-term Vision
General Transformer Acceleration
Extends APT's methods to NLP and scientific computing, enhancing Transformer efficiency.
Abstract
Recent advances in generative AI have significantly increased the demand for high-resolution image and video generation, positioning diffusion models as a core technology. Among them, Diffusion Transformers (DiTs) have emerged as the state-of-the-art (SOTA) models due to their scalability and output quality. However, self-attention in DiTs incurs significant computational overhead, leading to excessively long latency as the complexity grows with the fourth power of the output resolution. While prior works have attempted to mitigate this cost using sparsity and quantization techniques, they fall short of effectively reducing the computational cost in high-resolution DiTs. In this paper, we present APT, a software-hardware co-designed accelerator for high-resolution DiTs. APT leverages attention probabilities as a unified importance metric to jointly optimize computation through fine-grained pruning and adaptive precision scaling. At the algorithm level, we propose Attention Probability-guided Adaptive Dual Thresholding (APDT), which dynamically performs element selection and precision assignment using dual thresholds. To ensure compatibility with memory-efficient FlashAttention, we introduce Timestep-Aware FlashAttention (TAFA), which predicts attention probabilities across timesteps by exploiting temporal similarity. At the architecture level, we co-design a specialized accelerator that efficiently supports irregular sparsity and dual-precision execution, featuring dynamic mask management, address translation, dual-precision compute units, and a tile-based dataflow. Finally, we evaluate APT on SOTA DiT models, including PixArt-$α$, Stable Diffusion 3, and FLUX. APT achieves up to 8.16$\times$ speedup and 14.98$\times$ higher energy efficiency over NVIDIA A100, and up to 3.01$\times$ speedup and 2.04$\times$ higher energy efficiency over EXION, a SOTA diffusion model accelerator.