PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2403.04692v2cs.CV

TL;DR

High-quality text-to-image generation is computationally expensive, while efficiently integrating improved data and algorithms into an existing model remains a challenge. PixArt-Σ addresses this by fine-tuning PixArt-α with higher-quality data and key-value token compression, achieving direct 4K generation with a compact model and efficient training. Its results include strong image quality and prompt adherence, although specific scenes, text, hands, complex prompts, faces, and safety remain imperfect.

  • Problem

    Integrating continuously improving datasets and algorithms into existing text-to-image models efficiently, under limited resources, remains an open challenge.

  • Method

    PixArt-Σ uses weak-to-strong training from PixArt-α, combining higher-quality images and captions with key-value token compression for efficient high-resolution generation.

  • Results

    PixArt-Σ directly generates high-quality 4K images with 0.6B parameters, using only 9% of PixArt-α’s GPU days for fine-tuning while closely following textual prompts.

  • Takeaways & Limitations

    PixArt-Σ provides an efficient Diffusion Transformer route to high-quality 4K text-to-image generation and competitive visual results against open-source and commercial systems.

  • Takeaways & Limitations

    The model still struggles with specific scenes and objects, especially text and hands, and may fail to fully align complex prompts or generate flawless faces.

Abstract

from arXiv · show

In this paper, we introduce PixArt-Σ, a Diffusion Transformer model~(DiT) capable of directly generating images at 4K resolution. PixArt-Σrepresents a significant advancement over its predecessor, PixArt-α, offering images of markedly higher fidelity and improved alignment with text prompts. A key feature of PixArt-Σis its training efficiency. Leveraging the foundational pre-training of PixArt-α, 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-Σare twofold: (1) High-Quality Training Data: PixArt-Σ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-Σ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-Σ'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.

1 Introduction

PixArt-Σ develops PixArt-α into a stronger 4K text-to-image model through weak-to-strong training, combining higher-quality data, token compression, and fine-tuning strategies. It targets efficient improvement under limited resources while enhancing image quality and prompt alignment.

  • Motivation: Training top-tier text-to-image models requires substantial computational resources, creating barriers for researchers with limited resources.Training SD1.5 from scratch requires about 6000 A100 GPU days.
  • Weak-to-Strong Training: Weak-to-strong training builds PixArt-Σ on PixArt-α’s pre-trained foundation and evolves the weaker baseline into a stronger model through efficient training.The strategy includes replacing the VAE, scaling from low to high resolution, and adding KV compression.
  • Higher-Quality Training Data: The training dataset contains 33M high-resolution images, including 2.3M images around 4K, paired with dense captions from Share-Captioner.All images exceed 1K resolution, and the text encoder’s token length is extended to improve textual-visual alignment.
  • Efficient Token Compression: Key-value token compression uses stride-2 group convolutions for local aggregation, reducing high-resolution training and inference time by ∼34%.A specialized initialization scheme adapts compression smoothly from a pre-trained model without KV compression.
  • Results: PixArt-Σ directly generates high-quality 4K images with 0.6B parameters and uses only 9% of PixArt-α’s GPU days for fine-tuning.The model also demonstrates fine-grained textual prompt alignment and aesthetic quality comparable to leading text-to-image products.

2 Related Work

Related work establishes the computational difficulty of scaling diffusion Transformers and prior approaches for efficient attention and high-resolution generation. PixArt-Σ extends this line by directly generating 4K images with a Transformer architecture.

  • Diffusion Transformers: PixArt-α achieved high-quality 1024px image generation through efficient text-to-image training schemes within the Diffusion Transformer framework.Related models include DiT, UViT, GenTron, and later Transformer-based diffusion systems.
  • Comparisons: Qualitative comparisons report improved portrait realism and semantic analysis over PixArt-α, plus stronger instruction following than SDXL.The figure highlights prompt keywords in blue.
  • PixArt-Σ: PixArt-Σ is presented as the first work to use a Transformer architecture for direct 4K ultra-high-resolution image generation.The approach addresses the computational complexity caused by long token sequences.
  • High-Resolution Generation: Prior high-resolution methods commonly add super-resolution networks or multiple diffusion stages, which can introduce cumulative errors.Imagen, GigaGAN, Stable Diffusion, and Stable Cascade exemplify these combined-model solutions.
  • Efficient Transformer Architecture: Transformer self-attention has quadratic computational complexity in token count, motivating sparse, local, and low-rank attention alternatives.PVT v2 condenses key and value spaces with convolution, while Swin Transformer restricts computation to windows.

3 Framework

PixArt-Σ improves training data quality and DiT efficiency to support direct high-resolution generation. Its framework combines denser captions, high-resolution images, KV token compression, and weak-to-strong fine-tuning.

  • Higher Aesthetic and higher Resolution: As image resolution increases, model fidelity and semantic alignment improve according to FID and CLIP Score.The paper uses this observation to motivate training directly on higher-resolution data.
  • Efficient DiT Design: KV compression merges spatially adjacent key and value tokens while retaining all query tokens, reducing attention computation for high-resolution generation.The module uses convolutional aggregation and compresses keys and values rather than queries.
  • Efficient DiT Design: The Conv Avg Init convolution starts as an average operator, enabling coarse initial outputs during fine-tuning while adding only 0.018% parameters.This initialization uses group convolution with weights initialized to 1/R^2.
  • Weak-to-Strong Training Strategy: Weak-to-strong training combines VAE replacement, low-to-high-resolution scaling, and progression from uncompressed to KV-compressed models.The strategy is illustrated as accelerating convergence during these transitions.

4 Experiment

PixArt-Σ is evaluated with qualitative comparisons, human and AI preference studies, and high-resolution generation tests. It directly generates 4K images, follows detailed prompts, and compares favorably with open-source and commercial systems.

  • Image Quality Assessment: PixArt-Σ improves portrait realism and semantic analysis over SDXL and PixArt-α while adhering more closely to user instructions than SDXL.These findings come from qualitative comparisons with open-source models.
  • Image Quality Assessment: PixArt-Σ produces photo-realistic images and adheres closely to user instructions, performing competitively with contemporary closed-source products.The comparison is reported in Fig. 4 against current commercial T2I products.
  • High-resolution Generation: PixArt-Σ directly generates 4K images without post-processing and follows complex, detailed, and long user prompts.The paper states that satisfactory results do not require prompt engineering.
  • High-resolution Generation: The paper argues that direct 4K generation avoids cascade-related cumulative errors and learns text alignment from 4K image distributions.This contrasts with post-processing and super-resolution pipelines described in the discussion.
  • Performance Comparisons: Human and GPT-4 Vision preference studies show PixArt-Σ surpassing PixArt-α and matching or exceeding Stable Cascade on image quality and instruction following.Human evaluators ranked quality and prompt matching, while GPT-4 Vision voted on image quality and image-text alignment.

5 Ablation Studies

The ablation studies evaluate KV compression using FID, CMMD, CLIP-Score, and speed across compression positions, operators, ratios, and resolutions. Deep-layer compression with Conv 2×2 is identified as the strongest design choice among tested variants.

  • Experimental setup: Unless specified, KV-compression ablations use 512px generation, with detailed experimental settings provided in the appendix.FID and CLIP-Score are used to compare generated-image distribution and prompt alignment.
  • Compression positions: KV compression in deep Transformer layers achieves superior performance to compression in shallow or intermediate layers.The authors attribute this to deep layers encoding higher-level semantics, which are less affected by compression than shallow texture features.
  • Compression operators: Conv 2×2 outperforms random discarding and average pooling when merging 2×2 tokens into one token.The result supports using a learnable kernel to reduce redundant features.
  • Compression ratios on different resolutions: The experiments vary compression ratios across resolutions and report speed using seconds per iteration for batch size 384.The study evaluates compression ratio, position, operator, resolution, and speed.

6 Conclusion

PixArt-Σ builds on PixArt-α through weak-to-strong training, combining higher-quality data and efficient token compression to generate high-quality 4K images with strong prompt adherence.

  • PixArt-Σ directly generates high-quality images at 4K resolution.
  • Weak-to-strong training upgrades PixArt-α by incorporating higher-quality data and efficient token compression.
  • PixArt-Σ produces high-fidelity images while adhering closely to textual prompts.

A.1 Extension: Inference Acceleration (PixArt + DMD).

PixArt integrates DMD for one-step inference, with the denoising timestep T investigated experimentally and performance compared against LCM.

  • DMD accelerates PixArt inference through distribution matching distillation and a one-step generator.The generated image is represented as x0 = Gθ(T, text).
  • The denoising timestep T was initially set to 999, matching the training denoising timesteps.
  • Figure 10 reports Base model + DMD performance over T.
  • Table 4 compares PixArt + DMD with PixArt + LCM at 512x512 resolution and batch size 1.

A.2 Training Details of PixArt-Σ

PixArt-Σ training uses staged adaptation and evaluates efficiency gains from KV token compression, while Figure 11 compares DMD and LCM performance.

  • A.2 Training Details of PixArt-Σ: PixArt-Σ training stages record image resolution, training samples, steps, batch size, learning rate, and GPU-day compute.
  • A.2 Training Details of PixArt-Σ: KV token compression reduces 512px-to-1024px fine-tuning time from 50 to 20 V100 GPU days.
  • A.2 Training Details of PixArt-Σ: Figure 11 provides a visual comparison of Base model + DMD and Base model + LCM.
  • A.2 Training Details of PixArt-Σ: KV token compression reduces 2K training time from 20 to 14 A800 GPU days and 4K training time from 25 to 20 A800 GPU days.

A.3 Detailed Settings of the Ablation Studies

The ablation studies evaluate FID and CLIP scores using curated prompts and a high-quality evaluation dataset, with testing-image count and resolution specified as important settings.

  • Ablations evaluate FID on 30,000 curated prompts against images from the High-Quality Evaluation Dataset.
  • CLIP scores compare model-generated images with the original 30,000 prompts.
  • GPU-day counts exclude offline VAE-feature and T5-text-feature extraction because both features are prepared before training.
  • Testing-image count affects FID, with 10,000-image datasets typically yielding higher scores.
  • Compression-position and operator experiments use images generated at 512px resolution.

A.4 Samples of High-Quality Evaluation Dataset

The paper presents a 30,000-pair high-quality evaluation dataset and visual comparisons showing PixArt-Σ’s image quality, prompt alignment, and competitiveness with commercial T2I products. Additional samples demonstrate fine-grained detail, varied aspect ratios, and direct 4K generation, while the authors acknowledge remaining generation limitations.

  • High-Quality Evaluation Dataset: 30,000 high-quality text-image pairs form the proposed evaluation dataset for assessing visually appealing images.The dataset emphasizes superior aesthetic quality and precise text-image alignment.
  • Comparisons with T2I Products: PixArt-Σ produces high-quality, photo-realistic images with rich details comparable with Firefly 2, Imagen 2, Dalle 3, and Midjourney 6.The comparison is presented through visual results in Figures 14 and 15.
  • Generated Samples: The authors report that PixArt-Σ samples closely match provided textual prompts and exhibit high fidelity.These qualities are illustrated in Figures 16, 17, and 18.
  • Limitations: PixArt-Σ still struggles with specific scenes and objects, especially text and hands, and may fail on complex prompts or produce flawed faces.The paper also notes that sensitive content may be generated.
  • Generated Samples: PixArt-Σ generates high-quality images with fine-grained details and diverse aspect ratios.The samples illustrate variation in both visual detail and image proportions.
  • 4K Generation: PixArt-Σ directly generates 4K HD images at 3840×2560 while preserving fine-grained details.Figures 17 and 18 provide high-resolution sample outputs.
Loading 2403.04692v2…