LinFusion: 1 графический процессор, 1 минута, изображение 16K

LinFusion: 1 GPU, 1 Minute, 16K Image

September 3, 2024

Авторы: Songhua Liu, Weihao Yu, Zhenxiong Tan, Xinchao Wang

cs.AI

Аннотация

Современные модели диффузии, особенно те, которые используют Transformer-основанный UNet для подавления шума, сильно полагаются на операции самовнимания для управления сложными пространственными отношениями, что позволяет достичь впечатляющей производительности генерации. Однако существующая парадигма сталкивается с существенными вызовами при создании визуального контента высокого разрешения из-за квадратичной сложности по времени и памяти относительно количества пространственных токенов. Для преодоления этого ограничения мы предлагаем в данной статье новый линейный механизм внимания в качестве альтернативы. Конкретно, мы начинаем наше исследование с недавно представленных моделей с линейной сложностью, например, Mamba, Mamba2 и Gated Linear Attention, и выделяем две ключевые особенности - нормализацию внимания и некаузальное вывод, которые улучшают производительность генерации визуального контента высокого разрешения. Основываясь на этих идеях, мы представляем обобщенную парадигму линейного внимания, которая служит низкоранговым приближением широкого спектра популярных линейных смесителей токенов. Для снижения затрат на обучение и более эффективного использования предварительно обученных моделей мы инициализируем наши модели и дистиллируем знания из предварительно обученной StableDiffusion (SD). Мы обнаружили, что дистиллированная модель, названная LinFusion, достигает производительности на уровне или превосходящей оригинальную SD уже после небольшого обучения, при этом значительно снижая временную и памятьную сложность. Обширные эксперименты на SD-v1.5, SD-v2.1 и SD-XL показывают, что LinFusion обеспечивает удовлетворительную производительность по генерации изображений разного разрешения без обучения на таких разрешениях, как 16K. Более того, он хорошо совместим с предварительно обученными компонентами SD, такими как ControlNet и IP-Adapter, не требуя усилий по адаптации. Коды доступны по адресу 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.