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

LinFusion: 1 GPU, 1 Minute, 16K Image

September 3, 2024

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

cs.AI

Resumo

Os modelos de difusão modernos, especialmente aqueles que utilizam um UNet baseado em Transformer para remoção de ruído, dependem fortemente de operações de autoatenção para gerenciar relações espaciais complexas, alcançando assim um desempenho impressionante na geração. No entanto, esse paradigma existente enfrenta desafios significativos na geração de conteúdo visual de alta resolução devido à sua complexidade quadrática de tempo e memória em relação ao número de tokens espaciais. Para lidar com essa limitação, propomos um novo mecanismo de atenção linear como alternativa neste artigo. Especificamente, iniciamos nossa exploração a partir de modelos recentemente introduzidos com complexidade linear, como Mamba, Mamba2 e Gated Linear Attention, e identificamos dois recursos-chave - normalização de atenção e inferência não causal - que aprimoram o desempenho na geração visual de alta resolução. Com base nessas percepções, introduzimos um paradigma de atenção linear generalizada, que serve como uma aproximação de baixa ordem de uma ampla gama de misturadores de tokens lineares populares. Para reduzir o custo de treinamento e aproveitar melhor os modelos pré-treinados, inicializamos nossos modelos e destilamos o conhecimento do StableDiffusion pré-treinado (SD). Descobrimos que o modelo destilado, denominado LinFusion, alcança desempenho equivalente ou superior ao SD original após apenas um treinamento modesto, enquanto reduz significativamente a complexidade de tempo e memória. Experimentos extensivos no SD-v1.5, SD-v2.1 e SD-XL demonstram que o LinFusion oferece um desempenho satisfatório na geração de alta resolução entre diferentes resoluções, gerando imagens de alta resolução, como resolução 16K. Além disso, é altamente compatível com componentes SD pré-treinados, como ControlNet e IP-Adapter, não exigindo esforços de adaptação. Os códigos estão disponíveis em 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.