LinFusion: 1 GPU, 1 Minuto, Imagen de 16K

LinFusion: 1 GPU, 1 Minute, 16K Image

September 3, 2024

Autores: Songhua Liu, Weihao Yu, Zhenxiong Tan, Xinchao Wang

cs.AI

Resumen

Los modelos de difusión modernos, en particular aquellos que utilizan un UNet basado en Transformadores para el desruido, dependen en gran medida de operaciones de autoatención para gestionar relaciones espaciales complejas, logrando así un impresionante rendimiento en la generación. Sin embargo, este paradigma existente enfrenta desafíos significativos en la generación de contenido visual de alta resolución debido a su complejidad cuadrática en tiempo y memoria con respecto al número de tokens espaciales. Para abordar esta limitación, nos proponemos en este artículo un novedoso mecanismo de atención lineal como alternativa. Específicamente, comenzamos nuestra exploración a partir de modelos recientemente introducidos con complejidad lineal, como Mamba, Mamba2 y Atención Lineal con Compuertas, e identificamos dos características clave: normalización de atención e inferencia no causal, que mejoran el rendimiento en la generación visual de alta resolución. Basándonos en estos conocimientos, introducimos un paradigma generalizado de atención lineal, que sirve como una aproximación de rango bajo de un amplio espectro de mezcladores de tokens lineales populares. Para reducir el costo de entrenamiento y aprovechar mejor los modelos pre-entrenados, inicializamos nuestros modelos y destilamos el conocimiento del modelo pre-entrenado StableDiffusion (SD). Observamos que el modelo destilado, denominado LinFusion, logra un rendimiento igual o superior al SD original después de un entrenamiento modesto, al tiempo que reduce significativamente la complejidad en tiempo y memoria. Experimentos extensos en SD-v1.5, SD-v2.1 y SD-XL demuestran que LinFusion ofrece un rendimiento satisfactorio en la generación de resoluciones cruzadas sin entrenamiento, generando imágenes de alta resolución como 16K. Además, es altamente compatible con componentes pre-entrenados de SD, como ControlNet e IP-Adapter, sin necesidad de esfuerzos de adaptación. El código está disponible en https://github.com/Huage001/LinFusion.

English

Modern diffusion models, particularly those utilizing a Transformer-based UNet for denoising, rely heavily on self-attention operations to manage complex spatial relationships, thus achieving impressive generation performance. However, this existing paradigm faces significant challenges in generating high-resolution visual content due to its quadratic time and memory complexity with respect to the number of spatial tokens. To address this limitation, we aim at a novel linear attention mechanism as an alternative in this paper. Specifically, we begin our exploration from recently introduced models with linear complexity, e.g., Mamba, Mamba2, and Gated Linear Attention, and identify two key features-attention normalization and non-causal inference-that enhance high-resolution visual generation performance. Building on these insights, we introduce a generalized linear attention paradigm, which serves as a low-rank approximation of a wide spectrum of popular linear token mixers. To save the training cost and better leverage pre-trained models, we initialize our models and distill the knowledge from pre-trained StableDiffusion (SD). We find that the distilled model, termed LinFusion, achieves performance on par with or superior to the original SD after only modest training, while significantly reducing time and memory complexity. Extensive experiments on SD-v1.5, SD-v2.1, and SD-XL demonstrate that LinFusion delivers satisfactory zero-shot cross-resolution generation performance, generating high-resolution images like 16K resolution. Moreover, it is highly compatible with pre-trained SD components, such as ControlNet and IP-Adapter, requiring no adaptation efforts. Codes are available at https://github.com/Huage001/LinFusion.