CLEAR: Conv-Like Linearization Revs Pre-Trained Diffusion Transformers Up (2412.16112v1)

Abstract: Diffusion Transformers (DiT) have become a leading architecture in image generation. However, the quadratic complexity of attention mechanisms, which are responsible for modeling token-wise relationships, results in significant latency when generating high-resolution images. To address this issue, we aim at a linear attention mechanism in this paper that reduces the complexity of pre-trained DiTs to linear. We begin our exploration with a comprehensive summary of existing efficient attention mechanisms and identify four key factors crucial for successful linearization of pre-trained DiTs: locality, formulation consistency, high-rank attention maps, and feature integrity. Based on these insights, we introduce a convolution-like local attention strategy termed CLEAR, which limits feature interactions to a local window around each query token, and thus achieves linear complexity. Our experiments indicate that, by fine-tuning the attention layer on merely 10K self-generated samples for 10K iterations, we can effectively transfer knowledge from a pre-trained DiT to a student model with linear complexity, yielding results comparable to the teacher model. Simultaneously, it reduces attention computations by 99.5% and accelerates generation by 6.3 times for generating 8K-resolution images. Furthermore, we investigate favorable properties in the distilled attention layers, such as zero-shot generalization cross various models and plugins, and improved support for multi-GPU parallel inference. Models and codes are available here: https://github.com/Huage001/CLEAR.

• The paper presents CLEAR, which replaces global attention with convolution-like local attention to achieve linear complexity in Diffusion Transformers.
• It employs knowledge distillation to reduce attention computations by 99.5% while accelerating 8K image generation by 6.3 times.
• The method supports multi-GPU parallelization, ensuring scalability and competitive performance in high-resolution image synthesis.

Overview of CLEAR: Efficient Linearization in Diffusion Transformers

The paper "CLEAR: Conv-Like Linearization Revs Pre-Trained Diffusion Transformers Up" addresses the computational challenges involved in generating high-resolution images using Diffusion Transformers (DiTs). While DiTs have emerged as a prominent architecture for text-to-image generation, their reliance on attention mechanisms leads to quadratic computational complexity, particularly problematic for high-resolution images. This paper introduces an approach called CLEAR (Convolution-Like Linearization) that aims to reduce the complexity of DiTs from quadratic to linear.

Key Contributions and Methodology

The authors begin by examining existing efficient attention mechanisms, identifying critical factors for successfully linearizing pre-trained DiTs, namely locality, formulation consistency, high-rank attention maps, and feature integrity. These insights form the basis for their proposed convolution-like local attention strategy, CLEAR. CLEAR limits feature interactions to a local window around each query token, achieving linear complexity.

Key aspects of their methodology include:

• Convolution-Like Local Attention: CLEAR employs local attention windows similar to convolutional operations, restricting interactions to predefined circular areas around query tokens. This strategy achieves linear complexity in terms of image resolution, allowing for efficient high-resolution image generation.
• Knowledge Distillation: By distilling knowledge from pre-trained DiTs to models implementing CLEAR, the paper demonstrates the ability to significantly reduce attention computations by 99.5% while accelerating image generation by 6.3 times for 8K-resolution images.
• Multi-GPU Parallelization: The CLEAR strategy supports multi-GPU parallel inference, offering efficient image generation by minimizing communication overhead between GPUs.

Results and Implications

The experimental results confirm that CLEAR can perform comparably to the original DiT designs while vastly reducing computational demands. Notable achievements include maintaining performance with a mere 10,000 fine-tuning iterations on only 10,000 self-generated samples, illustrating both efficiency and practical viability.

1. Quantitative Results: The application of CLEAR yields a substantial reduction in the computational burden associated with attention layers, and accelerates high-resolution image generation significantly. In benchmarks against models with full quadratic attention, CLEAR maintains competitive FID scores and image quality metrics.
2. Generalization: The CLEAR strategy demonstrated zero-shot generalization across different DiT models and successfully integrated additional features like ControlNet, highlighting its compatibility and adaptability.
3. Multi-GPU Scalability: The ability of CLEAR to facilitate multi-GPU parallel processing without significant adaptation further extends its applicability to scenarios requiring extensive computational resources, making linear attention methods practical for widespread use in computationally intensive settings.

Future Directions

CLEAR represents a significant step towards scalable and efficient high-resolution image synthesis using DiTs. Moving forward, extensive optimizations for hardware implementation of sparse attention mechanisms could further bridge the gap between theoretical and practical efficiency improvements noted in the paper. Additionally, exploring further optimizations and potentially integrating global attention elements could address specific challenges like overall scene coherence when required.

The advancements outlined in this paper hold substantial promise for the future of efficient large-scale AI applications, particularly in contexts where high-resolution imagery plays a central role. With improvements in both performance and computation, CLEAR sets the stage for broader adoption of advanced AI techniques in diverse use cases ranging from creative content generation to scientific visualization.