PixArt-Σ: Treinamento de Fraco para Forte do Transformer de Difusão para Geração de Imagens 4K a partir de Texto

PixArt-Σ: Weak-to-Strong Training of Diffusion Transformer for 4K Text-to-Image Generation

March 7, 2024

Autores: Junsong Chen, Chongjian Ge, Enze Xie, Yue Wu, Lewei Yao, Xiaozhe Ren, Zhongdao Wang, Ping Luo, Huchuan Lu, Zhenguo Li

cs.AI

Resumo

Neste artigo, apresentamos o PixArt-Σ, um modelo de Transformador de Difusão (DiT) capaz de gerar imagens diretamente em resolução 4K. O PixArt-Σ representa um avanço significativo em relação ao seu predecessor, o PixArt-α, oferecendo imagens com fidelidade notavelmente superior e melhor alinhamento com prompts de texto. Uma característica fundamental do PixArt-Σ é sua eficiência de treinamento. Aproveitando o pré-treinamento fundamental do PixArt-α, ele evolui de uma linha de base "mais fraca" para um modelo "mais forte" por meio da incorporação de dados de maior qualidade, um processo que denominamos "treinamento de fraco para forte". Os avanços no PixArt-Σ são duplos: (1) Dados de Treinamento de Alta Qualidade: o PixArt-Σ incorpora dados de imagem de qualidade superior, emparelhados com legendas de imagem mais precisas e detalhadas. (2) Compressão Eficiente de Tokens: propomos um novo módulo de atenção dentro da estrutura DiT que comprime tanto chaves quanto valores, melhorando significativamente a eficiência e facilitando a geração de imagens em ultra-alta resolução. Graças a essas melhorias, o PixArt-Σ alcança qualidade de imagem superior e capacidade de aderência a prompts do usuário com um tamanho de modelo significativamente menor (0,6 bilhões de parâmetros) em comparação com modelos de difusão texto-para-imagem existentes, como o SDXL (2,6 bilhões de parâmetros) e o SD Cascade (5,1 bilhões de parâmetros). Além disso, a capacidade do PixArt-Σ de gerar imagens em 4K suporta a criação de pôsteres e papéis de parede de alta resolução, impulsionando eficientemente a produção de conteúdo visual de alta qualidade em indústrias como cinema e jogos.

English

In this paper, we introduce PixArt-\Sigma, a Diffusion Transformer model~(DiT) capable of directly generating images at 4K resolution. PixArt-\Sigma represents a significant advancement over its predecessor, PixArt-\alpha, offering images of markedly higher fidelity and improved alignment with text prompts. A key feature of PixArt-\Sigma is its training efficiency. Leveraging the foundational pre-training of PixArt-\alpha, it evolves from the `weaker' baseline to a `stronger' model via incorporating higher quality data, a process we term "weak-to-strong training". The advancements in PixArt-\Sigma are twofold: (1) High-Quality Training Data: PixArt-\Sigma incorporates superior-quality image data, paired with more precise and detailed image captions. (2) Efficient Token Compression: we propose a novel attention module within the DiT framework that compresses both keys and values, significantly improving efficiency and facilitating ultra-high-resolution image generation. Thanks to these improvements, PixArt-\Sigma achieves superior image quality and user prompt adherence capabilities with significantly smaller model size (0.6B parameters) than existing text-to-image diffusion models, such as SDXL (2.6B parameters) and SD Cascade (5.1B parameters). Moreover, PixArt-\Sigma's capability to generate 4K images supports the creation of high-resolution posters and wallpapers, efficiently bolstering the production of high-quality visual content in industries such as film and gaming.