LinFusion : 1 GPU, 1 minute, image de 16K

LinFusion: 1 GPU, 1 Minute, 16K Image

September 3, 2024

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

cs.AI

Résumé

Les modèles de diffusion modernes, en particulier ceux utilisant un UNet basé sur un Transformer pour le débruitage, s'appuient fortement sur les opérations d'auto-attention pour gérer des relations spatiales complexes, permettant ainsi d'obtenir des performances de génération impressionnantes. Cependant, ce paradigme existant est confronté à des défis importants dans la génération de contenu visuel haute résolution en raison de sa complexité temporelle et de mémoire quadratique par rapport au nombre de jetons spatiaux. Pour remédier à cette limitation, nous proposons un nouveau mécanisme d'attention linéaire comme alternative dans cet article. Plus précisément, nous commençons notre exploration à partir de modèles récemment introduits avec une complexité linéaire, tels que Mamba, Mamba2 et Gated Linear Attention, et identifions deux caractéristiques clés - la normalisation de l'attention et l'inférence non causale - qui améliorent les performances de génération visuelle haute résolution. En nous appuyant sur ces observations, nous introduisons un paradigme d'attention linéaire généralisé, qui sert d'approximation de rang faible à un large éventail de mélangeurs de jetons linéaires populaires. Pour réduire le coût d'entraînement et mieux exploiter les modèles pré-entraînés, nous initialisons nos modèles et distillons les connaissances du StableDiffusion pré-entraîné (SD). Nous constatons que le modèle distillé, appelé LinFusion, atteint des performances comparables ou supérieures à celles du SD original après seulement un entraînement modeste, tout en réduisant considérablement la complexité temporelle et de mémoire. Des expériences approfondies sur SD-v1.5, SD-v2.1 et SD-XL montrent que LinFusion offre des performances satisfaisantes en génération inter-résolution sans apprentissage, générant des images haute résolution telles que la résolution 16K. De plus, il est hautement compatible avec les composants SD pré-entraînés, tels que ControlNet et IP-Adapter, ne nécessitant aucun effort d'adaptation. Les codes sont disponibles sur 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.