CLAR: Linearisation semblable à une convolution révèle les transformateurs de diffusion pré-entraînés

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

December 20, 2024

Auteurs: Songhua Liu, Zhenxiong Tan, Xinchao Wang

cs.AI

Résumé

Les Transformateurs à Diffusion (DiT) sont devenus une architecture de premier plan dans la génération d'images. Cependant, la complexité quadratique des mécanismes d'attention, qui sont responsables de la modélisation des relations entre jetons, entraîne une latence significative lors de la génération d'images haute résolution. Pour résoudre ce problème, nous visons un mécanisme d'attention linéaire dans cet article qui réduit la complexité des DiTs pré-entraînés à linéaire. Nous commençons notre exploration par un résumé complet des mécanismes d'attention efficaces existants et identifions quatre facteurs clés cruciaux pour la linéarisation réussie des DiTs pré-entraînés : la localité, la cohérence de la formulation, les cartes d'attention de haut rang et l'intégrité des caractéristiques. Sur la base de ces informations, nous introduisons une stratégie d'attention locale de type convolution appelée CLEAR, qui limite les interactions entre caractéristiques à une fenêtre locale autour de chaque jeton de requête, et atteint ainsi une complexité linéaire. Nos expériences indiquent que, en affinant la couche d'attention sur seulement 10 000 échantillons auto-générés pendant 10 000 itérations, nous pouvons transférer efficacement les connaissances d'un DiT pré-entraîné à un modèle étudiant avec une complexité linéaire, donnant des résultats comparables au modèle enseignant. En même temps, cela réduit les calculs d'attention de 99,5 % et accélère la génération de 6,3 fois pour la génération d'images de résolution 8K. De plus, nous étudions les propriétés favorables dans les couches d'attention distillées, telles que la généralisation sans apprentissage croisé entre différents modèles et plugins, et un meilleur support pour l'inférence parallèle multi-GPU. Les modèles et les codes sont disponibles ici : https://github.com/Huage001/CLEAR.

English

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.