CLEAR: Conv-Like Linearization Revs Pre-Trained Diffusion Transformers Up

TL;DR

This paper introduces CLEAR, a convolution-like local attention mechanism that linearizes pre-trained diffusion transformers, significantly reducing computational complexity while maintaining high-quality image generation.

Contribution

The paper proposes a novel linear attention method called CLEAR, enabling efficient fine-tuning of pre-trained diffusion transformers with minimal data and computational resources.

Findings

Reduces attention computation by 99.5%

Accelerates image generation by 6.3 times for 8K images

Maintains comparable quality to original models after fine-tuning

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.