LinFusion: 1 GPU, 1 minuut, 16K afbeelding

LinFusion: 1 GPU, 1 Minute, 16K Image

September 3, 2024

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

cs.AI

Samenvatting

Moderne diffusiemodellen, met name die gebruikmaken van een Transformer-gebaseerde UNet voor denoising, zijn sterk afhankelijk van self-attention-operaties om complexe ruimtelijke relaties te beheren, waardoor ze indrukwekkende generatieprestaties bereiken. Dit bestaande paradigma wordt echter geconfronteerd met aanzienlijke uitdagingen bij het genereren van visuele inhoud met hoge resolutie vanwege de kwadratische tijd- en geheugencomplexiteit ten opzichte van het aantal ruimtelijke tokens. Om deze beperking aan te pakken, streven we in dit artikel naar een nieuw lineair aandachtmechanisme als alternatief. Specifiek beginnen we onze verkenning met recent geïntroduceerde modellen met lineaire complexiteit, zoals Mamba, Mamba2 en Gated Linear Attention, en identificeren we twee belangrijke kenmerken - aandachtnormalisatie en niet-causale inferentie - die de prestaties van visuele generatie met hoge resolutie verbeteren. Op basis van deze inzichten introduceren we een gegeneraliseerd lineair aandachtsparadigma, dat dient als een lage-rangbenadering van een breed spectrum van populaire lineaire tokenmixers. Om de trainingskosten te besparen en beter gebruik te maken van vooraf getrainde modellen, initialiseren we onze modellen en distilleren we de kennis van vooraf getrainde StableDiffusion (SD). We ontdekken dat het gedistilleerde model, genaamd LinFusion, prestaties bereikt die vergelijkbaar zijn met of superieur aan de originele SD na slechts bescheiden training, terwijl de tijd- en geheugencomplexiteit aanzienlijk wordt verminderd. Uitgebreide experimenten op SD-v1.5, SD-v2.1 en SD-XL tonen aan dat LinFusion bevredigende zero-shot cross-resolutie generatieprestaties levert, waarbij afbeeldingen met hoge resolutie zoals 16K-resolutie worden gegenereerd. Bovendien is het zeer compatibel met vooraf getrainde SD-componenten, zoals ControlNet en IP-Adapter, zonder dat aanpassingsinspanningen nodig zijn. Codes zijn beschikbaar op 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.