CLEAR: Linearización tipo convolucional revierte transformadores de difusión pre-entrenados

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

December 20, 2024

Autores: Songhua Liu, Zhenxiong Tan, Xinchao Wang

cs.AI

Resumen

Los Transformadores de Difusión (DiT) se han convertido en una arquitectura líder en la generación de imágenes. Sin embargo, la complejidad cuadrática de los mecanismos de atención, responsables de modelar las relaciones entre tokens, resulta en una latencia significativa al generar imágenes de alta resolución. Para abordar este problema, nuestro objetivo en este documento es un mecanismo de atención lineal que reduce la complejidad de los DiTs pre-entrenados a lineal. Comenzamos nuestra exploración con un resumen exhaustivo de los mecanismos de atención eficientes existentes e identificamos cuatro factores clave cruciales para la exitosa linealización de los DiTs pre-entrenados: localidad, consistencia de formulación, mapas de atención de alto rango e integridad de características. Basándonos en estos conocimientos, presentamos una estrategia de atención local similar a una convolución denominada CLEAR, que limita las interacciones de características a una ventana local alrededor de cada token de consulta, logrando así una complejidad lineal. Nuestros experimentos indican que, al ajustar finamente la capa de atención en tan solo 10K muestras auto-generadas durante 10K iteraciones, podemos transferir efectivamente el conocimiento de un DiT pre-entrenado a un modelo estudiante con complejidad lineal, obteniendo resultados comparables al modelo docente. Al mismo tiempo, reduce los cálculos de atención en un 99.5% y acelera la generación en 6.3 veces para la generación de imágenes de resolución 8K. Además, investigamos propiedades favorables en las capas de atención destiladas, como la generalización de cero disparos entre varios modelos y complementos, y un mejor soporte para inferencia paralela multi-GPU. Los modelos y códigos están disponibles aquí: 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.