PixArt-α: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis

[Project Page]

[Github]

[arXiv](Current version v2)

 

 

Abstract

낮은 훈련 비용으로 고해상도 이미지 생성이 가능한 transformer 기반 T2I 확산 모델 PixArt-α 제안

 

 

 

Introduction

이미지 생성 품질을 유지하면서 훈련의 계산 요구를 크게 줄이는 3가지 핵심 디자인:

1. Training strategy decomposition : 이미지 생성 작업을 픽셀 분포 학습 - text-image 정렬 학습 - 이미지의 미적 품질 향상으로 분해하는 훈련 패러다임
2. Efficient T2I Transformer : DiT를 기반으로 cross-attention 추가, class-condition branch 간소화, reparameterization 기술 소개
3. High-informative data : 기존 text-image data에서 결함이 발견되었으며 LLaVA를 활용하여 SAM에서 캡션을 생성하는 auto-labeling pipeline 제안.

 

훈련 데이터는 Imagen의 0.2% 미만, 훈련 시간은 RAPHAEL의 2% 미만, 훈련 비용은 1%이며 기존 SOTA T2I 모델보다 우수한 생성 품질과 정렬을 제공한다.

 

 

 

Method

Motivation

기존 T2I 모델의 느린 훈련의 이유는 훈련 파이프라인과 데이터에 있다.

• 픽셀 분포 이해, 정렬, 미적 품질 향상을 동시에 훈련하므로 비효율적
• 낮은 품질의 데이터가 포함됨

 

Training strategy decomposition

Stage1: Pixel dependency learning

Class-condition model인 DiT의 사전 훈련된 가중치 사용.

 

Stage2: Text-image alignment learning

후술 할 개념 밀도가 높은 정확한 text-image pair dataset을 구축하여 효과적으로 정렬

 

Stage3: High-resolution and aesthetic image generation

고품질의 심미적 데이터를 활용하여 모델을 fine-tuning. 이전 단계의 강력한 사전 지식 덕분에 매우 빠르게 수렴함.

 

Efficient T2I transformer

Cross-Attention layer

Cross-attention layer를 추가하고 출력 projection을 0으로 초기화.

 

AdaLN-single

DiT에서는 모든 transformer block에서 각각 다른 MLP(c + t) 가중치를 통해 컨디셔닝 하였는데, 이 때문에 adaLN 모듈이 피라미터의 상당한 비율(27%)을 차지하였다.

 

하지만 PixArt-α에서는 class-condition이 필요 없으므로 제거하고 모든 transformer block에 동일한 MLP(t) 가중치를 사용하는 대신 각 블록에서 학습 가능한 임베딩 E_i를 추가하여 적응적으로 조절한다.

 

Re-parameterization

사전 훈련된 가중치와의 호환성을 유지하기 위해 각 블록에서 c가 없는 DiT와 동일한 S를 생성하도록 초기화된다.

 

Dataset construction

LAION dataset의 캡션은 다음 그림과 같이 설명 부족, 정렬 불량, 희귀한 어휘 사용 등의 문제가 있다.

 

“Describe this image and its style in a very detailed manner”라는 프롬프트와 함께 LLaVA를 사용하여 고품질 캡션 생성.

하지만 LAION dataset은 다양성 문제도 있어서 대신에 SAM dataset을 채택하였다.

 

 

 

Experiment

 

Project Page

PIXART-α:Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis