CLEAR: Linearização semelhante a convolução reverte Transformadores de Difusão Pré-Treinados Acima

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

December 20, 2024

Autores: Songhua Liu, Zhenxiong Tan, Xinchao Wang

cs.AI

Resumo

Os Transformadores de Difusão (DiT) tornaram-se uma arquitetura líder na geração de imagens. No entanto, a complexidade quadrática dos mecanismos de atenção, responsáveis por modelar as relações entre tokens, resulta em latência significativa ao gerar imagens de alta resolução. Para abordar esse problema, este artigo propõe um mecanismo de atenção linear que reduz a complexidade dos DiTs pré-treinados para linear. Iniciamos nossa exploração com um resumo abrangente dos mecanismos de atenção eficientes existentes e identificamos quatro fatores-chave cruciais para a linearização bem-sucedida dos DiTs pré-treinados: localidade, consistência na formulação, mapas de atenção de alto escalão e integridade de características. Com base nessas percepções, introduzimos uma estratégia de atenção local semelhante a convolução denominada CLEAR, que limita as interações de características a uma janela local ao redor de cada token de consulta, alcançando assim complexidade linear. Nossos experimentos indicam que, ao ajustar a camada de atenção em meramente 10 mil amostras auto-geradas por 10 mil iterações, podemos transferir efetivamente o conhecimento de um DiT pré-treinado para um modelo aluno com complexidade linear, produzindo resultados comparáveis ao modelo professor. Simultaneamente, reduz as computações de atenção em 99,5% e acelera a geração em 6,3 vezes para imagens de resolução de 8K. Além disso, investigamos propriedades favoráveis nas camadas de atenção destiladas, como generalização zero-shot entre vários modelos e complementos, e suporte aprimorado para inferência paralela multi-GPU. Modelos e códigos estão disponíveis aqui: 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.