Klare: Conv-ähnliche Linearisierung verbessert vortrainierte Diffusions-Transformer Hoch

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

December 20, 2024

Autoren: Songhua Liu, Zhenxiong Tan, Xinchao Wang

cs.AI

Zusammenfassung

Diffusion Transformers (DiT) sind zu einer führenden Architektur in der Bildgenerierung geworden. Die quadratische Komplexität der Aufmerksamkeitsmechanismen, die für die Modellierung der tokenweisen Beziehungen verantwortlich sind, führt jedoch zu erheblicher Latenz bei der Generierung von hochauflösenden Bildern. Um dieses Problem zu lösen, zielen wir in diesem Papier auf einen linearen Aufmerksamkeitsmechanismus ab, der die Komplexität der vortrainierten DiTs auf lineare Weise reduziert. Wir beginnen unsere Untersuchung mit einer umfassenden Zusammenfassung bestehender effizienter Aufmerksamkeitsmechanismen und identifizieren vier entscheidende Faktoren, die für eine erfolgreiche Linearisierung vortrainierter DiTs entscheidend sind: Lokalität, Formulierungskonsistenz, hochrangige Aufmerksamkeitskarten und Merkmalsintegrität. Basierend auf diesen Erkenntnissen stellen wir eine konvolutionsähnliche lokale Aufmerksamkeitsstrategie namens CLEAR vor, die die Merkmalsinteraktionen auf ein lokales Fenster um jeden Abfrage-Token beschränkt und so eine lineare Komplexität erreicht. Unsere Experimente zeigen, dass wir durch Feinabstimmung der Aufmerksamkeitsschicht an lediglich 10.000 selbstgenerierten Beispielen für 10.000 Iterationen effektiv Wissen von einem vortrainierten DiT auf ein Schülermodell mit linearer Komplexität übertragen können, wodurch Ergebnisse erzielt werden, die mit dem Lehrermodell vergleichbar sind. Gleichzeitig reduziert es die Aufmerksamkeitsberechnungen um 99,5 % und beschleunigt die Generierung um das 6,3-fache für die Erzeugung von 8K-Auflösungsbildern. Darüber hinaus untersuchen wir vorteilhafte Eigenschaften in den destillierten Aufmerksamkeitsschichten, wie beispielsweise Zero-Shot-Verallgemeinerung über verschiedene Modelle und Plugins sowie verbesserte Unterstützung für Multi-GPU-Parallelinferenz. Modelle und Codes sind hier verfügbar: 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.