Chiara: Linearizzazione simile alle reti neurali convoluzionali per rivisitare i trasformatori a diffusione pre-addestrati.

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

December 20, 2024

Autori: Songhua Liu, Zhenxiong Tan, Xinchao Wang

cs.AI

Abstract

I Transformer di Diffusione (DiT) sono diventati un'architettura di spicco nella generazione di immagini. Tuttavia, la complessità quadratica dei meccanismi di attenzione, responsabili della modellazione delle relazioni tra i token, comporta significative latenze nella generazione di immagini ad alta risoluzione. Per affrontare questo problema, in questo articolo ci proponiamo di implementare un meccanismo di attenzione lineare che riduca la complessità dei DiT pre-addestrati a lineare. Iniziamo la nostra esplorazione con un ampio riassunto dei meccanismi di attenzione efficienti esistenti e identifichiamo quattro fattori chiave cruciali per il successo della linearizzazione dei DiT pre-addestrati: località, coerenza della formulazione, mappe di attenzione ad alto rango e integrità delle caratteristiche. Sulla base di queste intuizioni, introduciamo una strategia di attenzione locale simile a una convoluzione chiamata CLEAR, che limita le interazioni tra le caratteristiche a una finestra locale intorno a ciascun token di query, ottenendo così una complessità lineare. I nostri esperimenti indicano che, mediante il raffinamento del livello di attenzione su soli 10K campioni auto-generati per 10K iterazioni, possiamo trasferire efficacemente le conoscenze da un DiT pre-addestrato a un modello studente con complessità lineare, ottenendo risultati paragonabili al modello docente. Allo stesso tempo, riduce i calcoli di attenzione del 99,5% e accelera la generazione di immagini con risoluzione 8K di 6,3 volte. Inoltre, esaminiamo le proprietà favorevoli nei livelli di attenzione distillati, come la generalizzazione senza addestramento incrociato tra vari modelli e plugin, e un miglior supporto per l'inferenza parallela multi-GPU. Modelli e codici sono disponibili qui: 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.