Scalable Diffusion Models with Transformers

Scalable Diffusion Models with Transformers

 

William Peebles, Saining Xie

 

Abstract

우리는 트랜스포머 아키텍처를 기반으로 한 새로운 종류의 디퓨전 모델을 탐구합니다.

우리는 일반적으로 사용되는 U-Net 백본을 잠재 패치에서 작동하는 트랜스포머로 대체하여 이미지의 잠재 디퓨전 모델을 학습합니다.

우리는 Gflops로 측정한 순방향 패스 복잡성 렌즈를 통해 Diffusion Transformers (DiT)의 확장성을 분석합니다.

우리는 더 높은 Gflops를 가진 DiT가 —트랜스포머 깊이/폭을 늘리거나 입력 토큰 수를 늘림으로써— 일관되게 낮은 FID를 가지고 있다는 것을 발견했습니다.

우수한 확장성 특성을 보유하고 있을 뿐만 아니라, 가장 큰 DiT-XL/2 모델은 클래스 조건부 ImageNet 512x512 및 256x256 벤치마크에서 이전의 모든 디퓨전 모델을 능가하여 후자에서 2.27의 SOTA FID를 달성했습니다.

 

 

1. Introduction

머신러닝은 트랜스포머를 기반으로 한 르네상스를 경험하고 있습니다.

지난 5년 동안 자연어 처리를 위한 신경 아키텍처 [8, 42], 비전 [10] 및 기타 여러 영역이 트랜스포머 [60]에 의해 대부분 포함되었습니다.

하지만 많은 이미지-수준 생성 모델 클래스는 여전히 이러한 추세를 고수하고 있습니다—트랜스포머는 자기 회귀 모델[3,6,43,47]에서 널리 사용되는 반면, 다른 생성 모델링 프레임워크에서는 채택률이 감소했습니다.

예를 들어, 디퓨전 모델은 최근 이미지-수준 생성 모델 [9,46]의 발전에 앞장서고 있지만; 모두 백본의 사실상 선택으로 컨볼루션 U-Net 아키텍처를 채택하고 있습니다.

 

Ho et al. [19]의 중요한 연구는 디퓨전 모델을 위한 U-Net 백본을 처음 도입했습니다.

픽셀 레벨의 autoregressive 모델과 Conditional GAN [23]에서 처음 성공을 거둔 U-Net은 몇 가지 변경 사항을 통해 PixelCNN+ [52, 58]에서 상속되었습니다.

이 모델은 주로 ResNet [15] 블록으로 구성된 컨볼루션입니다.

표준 U-Net [49]과 달리 트랜스포머의 필수 구성 요소인 추가 공간 셀프 어텐션 블록은 낮은 해상도에서 간섭을 받습니다.

Dhariwal 및 Nichol [9]은 컨벌루션 레이어에 대한 조건부 정보와 채널 수를 주입하기 위해 adaptive normalization layers [40]를 사용하는 등 U-Net을 위한 몇 가지 아키텍처 선택 사항을 완화했습니다.

그러나 Ho et al.의 U-Net의 하이레벨 설계는 대부분 그대로 유지되었습니다.

 

이 연구를 통해 디퓨전 모델에서 아키텍처 선택의 중요성을 설명하고 향후 생성 모델링 연구를 위한 경험적 베이스라인을 제공하는 것을 목표로 합니다.

우리는 U-Net 귀납적 편향이 디퓨전 모델의 성능에 결정적이지 않으며 트랜스포머와 같은 표준 설계로 쉽게 대체할 수 있음을 보여줍니다.

결과적으로 디퓨전 모델은 최근 아키텍처 통합 추세의 이점을 누릴 수 있는 좋은 위치에 있습니다—예를 들어, 다른 도메인에서 모범 사례 및 학습 레시피를 상속하고 확장성, 견고성 및 효율성과 같은 유리한 속성을 유지함으로써.

표준화된 아키텍처는 또한 도메인 간 연구를 위한 새로운 가능성을 열어줄 것입니다.

 

본 논문에서는 트랜스포머를 기반으로 한 새로운 종류의 디퓨전 모델에 초점을 맞춥니다.

우리는 이를 Diffusion Transformers 또는 줄여서 DiT라고 부릅니다.

DiT는 기존 컨볼루션 네트워크 (예: ResNet [15])보다 시각 인식을 위해 더 효과적으로 확장되는 것으로 나타난 Vision Transformers (ViT) [10]의 모범 사례를 준수합니다.

 

보다 구체적으로, 우리는 네트워크 복잡성 vs. 샘플 품질과 관련하여 트랜스포머의 확장 동작을 연구합니다.

우리는 디퓨전 모델이 VAE의 잠재 공간 내에서 학습되는 Latent Diffusion Models (LDM) [48] 프레임워크 하에서 DiT 설계 공간을 구성하고 벤치마킹함으로써 U-Net 백본을 트랜스포머로 성공적으로 대체할 수 있음을 보여줍니다.

또한 DiT가 디퓨전 모델을 위한 확장 가능한 아키텍처임을 보여줍니다: 네트워크 복잡성 (Gflops로 측정)과 샘플 품질 (FID로 측정) 사이에는 강한 상관관계가 있습니다.

DiT를 확장하고 대용량 백본 (118.6 Gflops)으로 LDM을 학습하기만 하면 클래스 조건부 256x256 ImageNet 생성 벤치마크에서 2.27 FID의 SOTA 결과를 얻을 수 있습니다.

 

 

2. Related Work

Transformers.

Transformers [60]는 언어, 비전 [10], 강화 학습 [5, 25] 및 메타 학습 [39] 전반에 걸쳐 도메인별 아키텍처를 대체했습니다.

이들은 증가하는 모델 크기, 언어 도메인 [26]의 컴퓨팅 및 데이터, 일반 autoregressive 모델 [17] 및 ViT [63]에서 놀라운 확장 특성을 보여주었습니다.

언어 외에도 트랜스포머는 픽셀 [6, 7, 38]을 자동 회귀적으로 예측하도록 학습되었습니다.

또한 autoregressive 모델 [11,47]과 마스킹 생성 모델 [4, 14]로 개별 codebooks [59]에 대해 학습되었으며, 전자는 최대 20B 매개변수 [62]까지 뛰어난 확장 동작을 보였습니다.

마지막으로, 트랜스포머는 DDPM에서 비공간 데이터를 합성하기 위해 탐색되었으며, 예를 들어 DALL-E 2 [41,46]에서 CLIP 이미지 임베딩을 생성하기 위해 탐색되었습니다.

본 논문에서는 이미지 디퓨전 모델의 백본으로 사용될 때 트랜스포머의 확장 특성을 연구합니다.

 

Denoising diffusion probabilistic models (DDPMs).

Diffusion [19, 54] 및 score-based generative models [22, 56]은 이미지 [35,46,48,50]의 생성 모델로서 특히 성공적이었으며, 많은 경우 이전에 SOTA였던 generative adversarial networks (GAN)[12]를 능가하는 성능을 보였습니다.

지난 2년 동안 DDPM의 개선은 주로 샘플링 기술 [19, 27, 55]의 개선, 특히 classifier-free guidance [21], 픽셀 대신 노이즈를 예측하도록 디퓨전  모델을 재구성하고 [19] 저해상도 base 디퓨전 모델을 업샘플러 [9, 20]와 병렬로 학습하는 계단식 DDPM 파이프라인 사용에 의해 주도되었습니다.

위에 나열된 모든 디퓨전 모델에 대해 컨볼루션 U-Nets[49]는 백본 아키텍처의 사실상 선택입니다.

동시 작업 [24]은 DDPM에 대한 어텐션을 기반으로 하는 새롭고 효율적인 아키텍처를 도입했으며, 순수 트랜스포머를 탐구합니다.

 

Architecture complexity.

이미지 생성 문헌에서 아키텍처 복잡성을 평가할 때는 매개 변수 수를 사용하는 것이 상당히 일반적인 관행입니다.

일반적으로 매개 변수 수는 성능에 큰 영향을 미치는 이미지 해상도를 고려하지 않기 때문에 이미지 모델의 복잡성에 대한 좋지 않은 프록시가 될 수 있습니다 [44, 45].

대신, 본 논문에서 모델 복잡성 분석의 대부분은 이론적 Gflops의 렌즈를 통해 이루어집니다.

이는 Gflops가 복잡성을 측정하는 데 널리 사용되는 아키텍처 설계 문헌과 일치합니다.

실제로 골든 복잡도 메트릭은 특정 애플리케이션 시나리오에 따라 달라지는 경우가 많기 때문에 여전히 논쟁의 여지가 있습니다.

Nichol 및 Dhariwal의 중요한 연구인 디퓨전 모델 개선 [9, 36]은 우리와 가장 관련이 있습니다—그곳에서 그들은 U-Net 아키텍처 클래스의 확장성과 Gflop 속성을 분석했습니다.

본 논문에서는 트랜스포머 클래스에 초점을 맞춥니다.

 

 

3. Diffusion Transformers

3.1. Preliminaries

Diffusion formulation.

아키텍처를 소개하기 전에 디퓨전 모델 (DDPM)을 이해하는 데 필요한 몇 가지 기본 개념을 간략하게 검토합니다 [19, 54].

가우시안 디퓨전 모델은 상수 α_t가 하이퍼파라미터인 실제 데이터 x_0: q(x_t|x_0) = N(x_t; √α_t x_0; (1 - α_t)I)에 점진적으로 노이즈를 적용하는 순방향 노이즈 프로세스를 가정합니다.

재파라미터화 트릭을 적용하여 x_t = √α_t x_0 + √(1 - α_t) ε_t를 샘플링할 수 있으며, 여기서 ε_t ~ N(0, I)입니다.

 

디퓨전 모델은 순방향 프로세스 손상을 반전시키는 역 프로세스를 학습하도록 학습됩니다: p_θ(x_(t-1))|x_t) = N(μ_θ(x_t), ∑_θ(x_t)), 여기서 신경망은 p_θ의 통계를 예측하는 데 사용됩니다.

역 프로세스 모델은 학습과 무관한 추가 항을 제외하고 L(θ) = -p(x_0|x_1) + ∑_t D_KL(q*(x_(t-1))|x_t, x_0)|p_θ(x_(t-1)|x_t)로 감소하는 log likelihood의 변동 하한 [30]으로 학습됩니다.
q*와 p_θ은 모두 가우시안이므로 D_KL은 두 분포의 평균과 공분산으로 평가할 수 있습니다.

μ_θ을 노이즈 예측 네트워크 ε_θ로 재파라미터화함으로써 예측 노이즈 ε_θ(x_t)와 샘플링된 가우스 노이즈 ε_t 사이의 간단한 평균 제곱 오차를 사용하여 모델을 학습시킬 수 있습니다: L_simple(θ) = |ε_θ(x_t) - ε_t|.

그러나 학습된 역 프로세스 공분산 ∑_θ로 디퓨전 모델을 학습하려면 전체 D_KL 항을 최적화해야 합니다.
우리는 Nichol 및 Dhariwal’s의 접근 방식을 따릅니다[36]: L_simple로 ε_θ을 학습하고 전체 L로 ∑_θ을 학습합니다.
p_θ이 학습되면 재파라미터화 트릭을 통해 x_t_max ~ N(0, I)을 초기화하고 x_(t-1) ~ p_θ(x_(t-1))을 샘플링하여 새로운 이미지를 샘플링할 수 있습니다.

 

Classifier-free guidance.

조건부 디퓨전 모델은 클래스 레이블 c와 같은 추가 정보를 입력으로 받습니다.
이 경우 역 과정은 p_θ(x_(t-1))|x_t; c)가 되며, 여기서 ε_θ와 ∑_θ는 c를 조건으로 합니다.
이 설정에서는 classifier-free guidance를 사용하여 샘플링 절차가 log p(c|x)가 높도록 x를 찾도록 유도할 수 있습니다 [21].
Bayes Rule에 의해, log p(c|x) ∝ log p(x|c) - log p(x), and 그러므로 ∇_x log p(c|x) ∝ ∇_x log p(x|c) - ∇_x log p(x) 이다.
디퓨전 모델의 출력을 score 함수로 해석하여 DDPM 샘플링 절차를 높은 p(x|c)를 가진 샘플 x로 유도할 수 있습니다: ^ε_θ(x_t, c) = ε_θ(x_t, ∅) + s · ∇_x log p(x|c) ∝ ε_θ(x_t, ∅) + s · (ε_θ(x_t, c)-ε_θ(x_t, ∅)), 여기서 s > 1은 가이던스의 척도를 나타냅니다(s = 1은 표준 샘플링을 복구한다는 점에 유의하세요).

c = ∅로 디퓨전 모델을 평가하는 것은 학습 중에 c를 랜덤으로 드롭하고 학습된 "null" 임베딩 ∅로 대체하여 수행됩니다.

classifier-free guidance은 일반 샘플링 기술 [21, 35, 46]에 비해 크게 개선된 샘플을 생성하는 것으로 널리 알려져 있으며, 이러한 추세는 DiT 모델에서도 유지됩니다.

 

Latent diffusion models.

고해상도 픽셀 공간에서 직접 디퓨전 모델을 학습하는 것은 계산적으로 금지될 수 있습니다.

Latent diffusion models (LDM) [48]은 2단계 접근 방식으로 이 문제를 해결합니다: (1) 학습된 인코더 E로 이미지를 더 작은 공간 표현으로 압축하는 오토인코더를 학습하고, (2) 이미지 x의 디퓨전 모델 대신 표현 z = E(x)의 디퓨전 모델을 학습합니다(E는 frozen).

그런 다음 디퓨전 모델에서 표현 z를 샘플링한 다음 학습된 디코더 x = D(z)가 있는 이미지로 디코딩하여 새로운 이미지를 생성할 수 있습니다.

그림 2. Diffusion Transformers (DiT)를 사용한 ImageNet 생성. 버블 영역은 디퓨전 모델의 flops을 나타냅니다. 왼쪽: 400K 학습 반복에서 우리의 DiT 모델의 FID-50K (낮을수록 좋습니다). 모델 flops이 증가함에 따라 FID 성능이 꾸준히 향상됩니다. 오른족: 저희 최고의 모델인 DiT-XL/2는 계산 효율성이 뛰어나며 ADM 및 LDM과 같은 이전의 모든 U-Net 기반 디퓨 모델을 능가합니다.

그림 2에서 볼 수 있듯이 LDM은 ADM과 같은 픽셀 공간 디퓨전 모델의 Gflops 중 일부를 사용하면서 좋은 성능을 달성합니다.

계산 효율성에 관심이 있기 때문에 아키텍처 탐색을 위한 매력적인 시작점이 됩니다.

본 논문에서는 잠재 공간에 DiT를 적용하지만, 픽셀 공간에도 수정 없이 적용할 수 있습니다.

따라서 이미지 생성 파이프라인은 하이브리드 기반 접근 방식이며 기성 컨볼루션 VAE와 트랜스포머 기반 DDPM을 사용합니다.

 

3.2. Diffusion Transformer Design Space

우리는 디퓨전 모델을 위한 새로운 아키텍처인 Diffusion Transformers (DiTs)를 소개합니다.
표준 트랜스포머 아키텍처에 최대한 충실하여 확장 특성을 유지하는 것을 목표로 합니다.
우리는 이미지의 DDPM (특히 이미지의 공간 표현)을 학습하는 데 중점을 두고 있기 때문에, DiT는 패치 시퀀스에서 작동하는 Vision Transformer (ViT) 아키텍처를 기반으로 합니다 [10].
DiT는 ViT의 모범 사례를 많이 보유하고 있습니다.
그림 3은 전체 DiT 아키텍처의 개요를 보여줍니다.
이 섹션에서는 DiT 클래스의 설계 공간 구성 요소뿐만 아니라 DiT의 순방향 패스에 대해 설명합니다.

그림 3. Diffusion Transformer (DiT) 아키텍처. 왼쪽: 조건부 잠재 DiT 모델을 학습합니다. 입력 잠재는 패치로 분해되어 여러 DiT 블록에 의해 처리됩니다. 오른쪽: 우리의 DiT 블록에 대한 세부 사항입니다. 우리는 적응형 레이어 노름, 크로스 어텐션 및 추가 입력 확인을 통한 컨디셔닝을 포함하는 표준 트랜스포머 블록의 변형을 실험합니다. 적응형 레이어 노름이 가장 잘 작동합니다.

Patchify.

DiT에 대한 입력은 공간 표현 z (256 x 256 x 3 이미지의 경우 z는 32 x 32 x 4 모양입니다) 입니다.
DiT의 첫 번째 레이어는 "patchify"로, 각 패치를 입력에 선형적으로 삽입하여 공간 입력을 차원 d마다 T개의 토큰 시퀀스로 변환합니다.
패치파이 후, 모든 입력 토큰에 표준 ViT 주파수 기반 위치 임베딩 (사인-코사인 버전)을 적용합니다.
패치파이에 의해 생성된 토큰 T의 수는 패치 크기 하이퍼파라미터 p에 의해 결정됩니다.
그림 4에서 볼 수 있듯이, p를 절반으로 줄이면 T가 4배가 되고, 따라서 전체 트랜스포머 G플롭이 최소 4배가 됩니다.
비록 Gflops에 큰 영향을 미치지만, p를 변경해도 다운스트림 매개변수 수에는 의미 있는 영향을 미치지 않는다는 점에 유의하세요.

 

우리는 DiT 설계 공간에 p = 2, 4, 8을 추가합니다.

그림 4. DiT의 입력 사양. 패치 크기 p x p가 주어졌을 때, 형태 I x I x C의 공간 표현(VAE로부터의 노이즈 잠재)은 숨겨진 차원 d를 가진 길이 T = (I/p)^2의 시퀀스로 "patchified"됩니다. 패치 크기 p가 작을수록 시퀀스 길이가 길어지고 따라서 Gflops이 더 많아집니다.

DiT block design.

패치파이 후 입력 토큰은 일련의 트랜스포머 블록에 의해 처리됩니다.
노이즈가 있는 이미지 입력 외에도 디퓨전 모델은 때때로 노이즈 타임스텝 t, 클래스 라벨 c, 자연어 등과 같은 추가 조건 정보를 처리합니다.
조건부 입력을 다르게 처리하는 네 가지 변형 트랜스포머 블록을 살펴봅니다.
디자인은 표준 ViT 블록 디자인에 작지만 중요한 수정 사항을 도입합니다.
모든 블록의 디자인은 그림 3에 나와 있습니다.

 

– In-context conditioning. 

입력 시퀀스에 t와 c의 벡터 임베딩을 두 개의 추가 토큰으로 추가하기만 하면 이미지 토큰과 다르지 않게 처리할 수 있습니다.
이는 ViT의 cls 토큰과 유사하며, 표준 ViT 블록을 수정 없이 사용할 수 있게 해줍니다.
마지막 블록이 끝나면 시퀀스에서 조건화 토큰을 제거합니다.
이 접근 방식은 모델에 무시할 수 있는 새로운 Gflops를 도입합니다.

 

– Cross-attention block.

우리는 이미지 토큰 시퀀스와 별도로 t와 c의 임베딩을 길이 2 시퀀스로 연결합니다.
트랜스포머 블록은 Vaswani et al. [60]의 원래 설계와 유사하게 멀티 헤드 셀프 어텐션 블록 뒤에 추가적인 멀티 헤드 크로스 어텐션 레이어를 포함하도록 수정되었으며, 클래스 라벨을 조건화하기 위해 LDM에서 사용하는 것과도 유사합니다.
크로스 어텐션은 모델에 가장 많은 Gflops를 추가하며, 대략 15%의 오버헤드를 발생시킵니다.

 

– Adaptive layer norm (adaLN) block.

GAN [2, 28]과 U-Net 백본 [9]을 사용하는 디퓨전 모델에서 적응형 정규화 레이어 [40]가 널리 사용된 후, 우리는 트랜스포머 블록에서 표준 레이어 노름 레이어를 adaptive layer norm (adaLN)으로 대체하는 방법을 탐구합니다.
차원별 스케일과 시프트 매개변수 γ와 β를 직접 학습하는 대신, 우리는 이를 t와 c의 임베딩 벡터의 합에서 회귀합니다.
우리가 탐구하는 세 가지 블록 디자인 중에서 adaLN은 가장 적은 Gflops를 추가하여 가장 계산 효율적입니다.
또한 모든 토큰에 동일한 함수를 적용하도록 제한된 유일한 조건화 메커니즘입니다.

 

– adaLN-Zero block.

ResNets에 대한 이전 연구에서는 각 잔여 블록을 항등 함수로 초기화하는 것이 유익하다는 것을 발견했습니다.
예를 들어, Goyal et al.은 각 블록에서 최종 배치 노름 척도 인자 γ를 0으로 초기화하면 지도 학습 환경에서 대규모 학습이 가속화된다는 것을 발견했습니다 [13].
디퓨전 U-Net 모델은 유사한 초기화 전략을 사용하여 잔여 연결 전에 각 블록의 최종 컨볼루션 레이어를 0으로 초기화합니다.
우리는 동일한 역할을 하는 adaLN DiT 블록의 수정을 탐구합니다.
γ 및 β 회귀 외에도 DiT 블록 내의 잔여 연결 직전에 적용되는 차원별 스케일링 매개변수 α도 회귀합니다.
MLP를 초기화하여 모든 α에 대해 영벡터를 출력합니다; 이렇게 하면 전체 DiT 블록이 항등 함수로 초기화됩니다.
vanilla adaLN 블록과 마찬가지로 adaLNZero는 모델에 무시할 수 있는 Gflops를 추가합니다.

 

우리는 DiT 설계 공간에 in-context, 크로스 어텐션, 적응형 레이어 노름 및 adaLN-Zero 블록을 포함합니다.

 

Model size.

우리는 히든 차원 크기 d에서 작동하는 N개의 DiT 블록 시퀀스를 적용합니다.
ViT 이후, 우리는 N, d, 그리고 어텐션 헤드를 공동으로 스케일링하는 표준 트랜스포머 구성을 사용합니다 [10, 63].

구체적으로, 우리는 네 가지 구성을 사용합니다: DiT-S, DiT-B, DiT-L 및 DiT-XL.

이들은 0.3에서 118.6 Gflops까지 다양한 모델 크기와 flop 할당을 포함하여 스케일링 성능을 측정할 수 있습니다.
표 1은 구성에 대한 세부 정보를 제공합니다.

 

우리는 DiT 설계 공간에 B, S, L, XL 구성을 추가합니다.

표 1. DiT 모델의 세부 사항. Small (S), Base (B) 및 Large (L) 변형에 대한 ViT [10] 모델 구성을 따르며, 가장 큰 모델로 XLarge (XL) 구성도 소개합니다.

Transformer decoder.

최종 DiT 블록 후에는 이미지 토큰 시퀀스를 출력 노이즈 예측과 출력 대각선 공분산 예측으로 디코딩해야 합니다.
이 두 출력 모두 원래 공간 입력과 동일한 모양을 가지고 있습니다.
이를 위해 표준 선형 디코더를 사용합니다; 최종 계층 노름 (adaLN을 사용하는 경우 적응형)을 적용하고 각 토큰을 p x p x 2C 텐서로 선형 디코딩합니다, 여기서 C는 DiT에 대한 공간 입력의 채널 수입니다.
마지막으로, 디코딩된 토큰을 원래의 공간 레이아웃으로 재배열하여 예측된 노이즈와 공분산을 얻습니다.

 

우리가 탐구하는 완전한 DiT 설계 공간은 패치 크기, 트랜스포머 블록 아키텍처 및 모델 크기입니다.

 

 

 

4. Experimental Setup

우리는 DiT 설계 공간을 탐구하고 모델 클래스의 스케일링 특성을 연구합니다.
우리 모델은 구성과 잠재 패치 크기 p에 따라 이름이 붙여집니다; 예를 들어, DiT-XL/2는 XLarge 구성을 의미하고 p = 2를 의미합니다.

 

Training.

우리는 경쟁이 치열한 생성 모델링 벤치마크인 ImageNet 데이터셋 [31]에서 256 x 256 및 512 x 512 이미지 해상도로 클래스 조건부 잠재 DiT 모델을 학습합니다.
최종 선형 레이어를 0으로 초기화하고, 그 외에는 ViT의 표준 가중치 초기화 기법을 사용합니다.
우리는 AdamW [29,33]로 모든 모델을 학습시킵니다.
우리는 1 x 10^-4의 일정한 학습 속도, 가중치 감소 없음, 배치 크기 256을 사용합니다.
우리가 사용하는 유일한 데이터 증강은 수평 플립입니다.
이전의 많은 ViTs [57, 61] 연구와 달리, 우리는 DiTs를 고성능으로 학습시키기 위해 학습률 워밍업이나 정규화가 필요하지 않다는 것을 발견했습니다.
이러한 기술이 없어도 모든 모델 구성에서 학습이 매우 안정적이었으며, 트랜스포머를 학습할 때 흔히 볼 수 있는 loss 급증 현상은 관찰되지 않았습니다.
생성 모델링 문헌의 일반적인 관행에 따라, 우리는 0.99999의 감소와 함께 학습 중 DiT 가중치의 exponential moving average (EMA)을 유지합니다.
보고된 모든 결과는 EMA 모델을 사용합니다.

우리는 모든 DiT 모델 크기와 패치 크기에 걸쳐 동일한 학습 하이퍼파라미터를 사용합니다.
우리의 학습 하이퍼파라미터는 거의 전적으로 ADM에서 유지됩니다.
학습 속도, 감쇠/워밍업 일정, Adam β_1/β_2 또는 가중치 감소를 조정하지 않았습니다.

 

Diffusion.

우리는 Stable Diffusion [48]의 기성품 사전 학습된 variational autoencoder (VAE) 모델 [30]을 사용합니다.
VAE 인코더의 다운샘플 팩터는 8입니다—256 x 256 x 3 모양의 RGB 이미지 x가 주어졌을 때, z = E(x)는 32 x 32 x 4 모양입니다.
이 섹션의 모든 실험에서 디퓨전 모델은 이 Z 공간에서 작동합니다.
디퓨전 모델에서 새로운 잠재 변수를 샘플링한 후, VAE 디코더 x = D(z)를 사용하여 픽셀로 디코딩합니다.
우리는 ADM [9]의 디퓨전 하이퍼파라미터를 유지합니다; 구체적으로, 우리는 1 x 10^-4에서 2 x 10^-2까지의 t_max = 1000 선형 분산 스케줄, 공분산 ∑_θ의 ADM 매개변수화 및 입력 시간 단계와 레이블을 임베딩하는 방법을 사용합니다.

 

Evaluation metrics.

우리는 이미지 생성 모델을 평가하기 위한 표준 지표인 Fr'echet Inception Distance (FID) [18]을 사용하여 스케일링 성능을 측정합니다.

 

우리는 이전 연구들과 비교할 때 관례를 따르고, 250개의 DDPM 샘플링 단계를 사용하여 FID-50K를 보고합니다.
FID는 작은 구현 세부 사항에 민감한 것으로 알려져 있습니다 [37]; 정확한 비교를 위해 이 논문에서 보고된 모든 값은 샘플을 내보내고 ADM의 TensorFlow 평가 모음을 사용하여 얻습니다 [9].
이 섹션에 보고된 FID 수치는 달리 명시된 경우를 제외하고는 classifier-free guidance를 사용하지 않습니다.
추가로 Inception Score [51], sFID[34] 및 Precision/Recall [32]을 2차 지표로 보고합니다.

 

Compute.

우리는 모든 모델을 JAX [1]에 구현하고 TPU-v3 포드를 사용하여 학습시킵니다.
가장 계산 집약적인 모델인 DiT-XL/2는 전역 배치 크기가 256인 TPU v3-256 포드에서 초당 약 5.7회의 반복 학습을 수행합니다.

 

 

 

5. Experiments

DiT block design.

우리는 각각 다른 블록 설계를 사용하여 가장 높은 Gflop DiT-XL/2 모델 네 가지를 학습시킵니다—in-context (119.4 Gflops), cross-attention (137.6 Gflops), adaptive layer norm (adaLN, 118.6 Gflops) 또는 adaLN-zero (118.6 Gflops).

우리는 학습 과정에서 FID를 측정합니다.
그림 5는 결과를 보여줍니다.

adaLN-Zero 블록은 크로스-어텐션 및 in-context 조건화보다 낮은 FID를 제공하면서도 가장 계산 효율적입니다.
400K의 학습 반복에서 adaLN-Zero 모델로 달성한 FID는 in-context 모델의 거의 절반에 달하며, 이는 컨디셔닝 메커니즘이 모델 품질에 중요한 영향을 미친다는 것을 보여줍니다.
초기화도 중요합니다—각 DiT 블록을 항등 함수로 초기화하는 adaLNZero는 vanilla adaLN을 크게 능가합니다.

나머지 논문에서는 모든 모델이 adaLN-Zero DiT 블록을 사용할 것입니다.

그림 5. 다양한 컨디셔닝 전략 비교하기. adaLNZero는 모든 학습 단계에서 크로스 어텐션 및 in-context 조건화를 능가합니다.

Scaling model size and patch size.

우리는 모델 구성 (S, B, L, XL)과 패치 크기 (8, 4, 2)를 포괄하여 12개의 DiT 모델을 학습시킵니다.
DiT-L과 DiT-XL은 다른 구성보다 상대적인 Gflops 측면에서 서로 훨씬 가깝다는 점에 유의하세요.
그림 2 (왼쪽)는 각 모델의 Gflops과 400K의 학습 반복 시 FID에 대한 개요를 제공합니다.
모든 경우에 모델 크기를 늘리고 패치 크기를 줄이면 디퓨전 모델이 상당히 개선된다는 것을 발견했습니다.

그림 6. DiT 모델을 확장하면 모든 학습 단계에서 FID가 향상됩니다. 우리는 12개의 DiT 모델에 대해 학습 반복을 통해 FID-50K를 보여줍니다. 위 행: FID 홀딩 패치 크기를 일정하게 비교합니다. 하단 행: FID 보유 모델 크기를 일정하게 비교합니다. 트랜스포머 백본을 확장하면 모든 모델 크기와 패치 크기에서 더 나은 생성 모델을 얻을 수 있습니다.

그림 6 (위)은 모델 크기가 증가하고 패치 크기가 일정하게 유지됨에 따라 FID가 어떻게 변화하는지 보여줍니다.
네 가지 구성 모두에서 트랜스포머를 더 깊고 넓게 만들면 모든 학습 단계에서 FID가 크게 향상됩니다.
마찬가지로 그림 6 (아래쪽)은 패치 크기가 감소하고 모델 크기가 일정하게 유지되는 FID를 보여줍니다.
우리는 DiT로 처리되는 토큰의 수를 간단히 조정하고 매개변수를 대략적으로 고정함으로써 학습 전반에 걸쳐 상당한 FID 개선을 다시 한 번 관찰했습니다.

그림 8. 트랜스포머 Gflops은 FID와 강한 상관관계가 있습니다. 우리는 각 DiT 모델의 Gflops과 각 모델의 FID-50K를 400K 학습 단계 후에 플롯합니다.

DiT Glops are critical to improving performance.

그림 6의 결과는 매개변수 수가 DiT 모델의 품질을 고유하게 결정하지 않는다는 것을 시사합니다.
모델 크기가 일정하게 유지되고 패치 크기가 줄어들면 트랜스포머의 총 매개변수는 사실상 변하지 않고 (사실 총 매개변수는 약간 감소하지만) Gflops만 증가합니다.
이 결과는 모델 Gflops를 확장하는 것이 실제로 성능 향상의 열쇠임을 나타냅니다.
이를 더 자세히 조사하기 위해 그림 8의 모델 Gflops에 대해 400K 학습 단계에서 FID-50K를 플롯합니다.
결과는 서로 다른 DiT 구성이 총 Gflops가 유사할 때 유사한 FID 값을 얻는다는 것을 보여줍니다 (예: DiT-S/2 및 DiT-B/4).
우리는 모델 Gflops와 FID-50K 사이에 강한 음의 상관관계가 있음을 발견했으며, 이는 추가적인 모델 계산이 개선된 DiT 모델의 중요한 요소임을 시사합니다.
그림 12 (부록)에서 우리는 이러한 추세가 Inception Score와 같은 다른 지표에도 적용된다는 것을 발견했습니다.

그림 12. 여러 생성 모델링 지표에서 DiT 스케일링 동작. 왼쪽: 우리는 FID, sFID, Inception Score, Precision 및 Recall에 대한 총 학습 계산의 함수로서 모델 성능을 도표화합니다. 오른쪽: 우리는 트랜스포머 Gflops에 대해 모든 12개의 DiT 변형에 대해 400K 학습 단계에서 모델 성능을 그래프로 나타내어, 지표 간에 강한 상관관계를 발견했습니다. 모든 값은 ft-MSE VAE 디코더를 사용하여 계산되었습니다.

그림 9. 더 큰 DiT 모델은 대용량 컴퓨팅을 더 효율적으로 사용합니다. 우리는 FID를 전체 학습 계산의 함수로 표시합니다.

Larger DiT models are more compute-efficient.

그림 9에서 우리는 모든 DiT 모델에 대한 총 학습 계산의 함수로 FID를 표시합니다.

우리는 학습 계산을 모델 Gflops · 배치 크기 · 학습 단계 · 3으로 추정하며, 여기서 3의 배수는 역방향 패스가 순방향 패스보다 두 배 더 무거운 것과 대략적으로 유사합니다.
우리는 작은 DiT 모델이 더 오래 학습되더라도 더 적은 단계로 학습된 큰 DiT 모델에 비해 결국 계산 효율성이 떨어진다는 것을 발견했습니다.
마찬가지로, 패치 크기를 제외하고 동일한 모델은 Gflops 학습을 제어할 때에도 성능 프로필이 다르다는 것을 발견했습니다.
예를 들어, XL/4는 약 10^10 Gflops 후 XL/2보다 성능이 뛰어납니다.

그림 7. 트랜스포머 순방향 패스 Gflops을 증가시키면 샘플 품질이 향상됩니다. 우리는 동일한 입력 잠재 노이즈와 클래스 레이블을 사용하여 400K 학습 단계를 거친 후, 모든 12개의 DiT 모델에서 샘플을 채취했습니다. 모델에서 Gflops을 증가시키면—트랜스포머 깊이/폭을 증가시키거나 입력 토큰의 수를 증가시킴으로써— 시각적 충실도가 크게 향상됩니다.

Visualizing scaling.

우리는 그림 7에서 스케일링이 샘플 품질에 미치는 영향을 시각화합니다.
400K 학습 단계에서 동일한 시작 노이즈 x_t_max, 샘플링 노이즈 및 클래스 레이블을 사용하여 12개의 DiT 모델 각각에서 이미지를 샘플링합니다.
이를 통해 스케일링이 DiT 샘플 품질에 미치는 영향을 시각적으로 해석할 수 있습니다.
실제로 모델 크기와 토큰 수를 모두 확장하면 시각적 품질이 눈에 띄게 향상됩니다.

그림 1. 트랜스포머 백본을 사용한 디퓨전 모델은 SOTA 이미지 품질을 달성합니다. 우리는 ImageNet에서 각각 512x512 및 256x256 해상도로 학습된 두 개의 클래스 조건부 DiT-XL/2 모델에서 선택된 샘플을 보여줍니다.

표 2. ImageNet 256x256에서 클래스 조건부 이미지 생성 벤치마킹. DiT-XL/2는 SOTA FID를 달성합니다.

5.1. State-of-the-Art Diffusion Models

256x256 ImageNet.

스케일링 분석에 이어, 우리는 7M 단계에 대해 최고 Gflop 모델인 DiT-XL/2를 계속해서 학습하고 있습니다.
우리는 그림 1에서 모델의 샘플을 보여주며, SOTA 클래스 조건부 생성 모델과 비교합니다.
우리는 표 2에 결과를 보고합니다.
classifier-free guidance를 사용할 때, DiT-XL/2는 이전의 모든 디퓨전 모델을 능가하여 LDM이 달성한 최고 FID-50K인 3.60을 2.27로 줄였습니다.
그림 2 (오른쪽)는 DiT-XL/2 (118.6 Gflops)가 LDM-4 (103.6 Gflops)와 같은 잠재 공간 U-Net 모델에 비해 계산 효율적이며, ADM (1120 Gflops)이나 ADM-U (742 Gflops)와 같은 픽셀 공간 U-Net 모델보다 훨씬 더 효율적임을 보여줍니다.

 

우리의 방법은 이전 SOTA StyleGAN-XL [53]을 포함하여 모든 이전 생성 모델 중 가장 낮은 FID를 달성합니다.
마지막으로, 우리는 DiT-XL/2가 LDM-4 및 LDM-8에 비해 모든 테스트된 classifier-free guidance 척도에서 더 높은 리콜 값을 달성한다는 것을 관찰했습니다.
2.35M 단계 (ADM과 유사)로만 학습했을 때, XL/2는 여전히 FID 2.55로 모든 이전 디퓨전 모델을 능가합니다.

표 3. ImageNet 512 x 512에서 클래스 조건부 이미지 생성 벤치마킹. 이전 연구 [9]는 512 x 512 해상도에서 1000개의 실제 샘플을 사용하여 정밀도와 재현율을 측정했으며, 일관성을 위해 동일하게 수행했습니다.

512x512 ImageNet.

우리는 ImageNet에서 새로운 DiT-XL/2 모델을 256 x 256 모델과 동일한 하이퍼파라미터를 사용하여 3M 반복에 대해 512 x 512 해상도로 학습시킵니다.
패치 크기가 2인 이 XL/2 모델은 64 x 64 x 4 입력 잠재값 (524.6 Gflops)을 패치한 후 총 1024개의 토큰을 처리합니다.
표 3은 최신 방법들과의 비교를 보여줍니다.
XL/2는 이 해상도에서 이전의 모든 디퓨전 모델을 다시 능가하여 ADM이 달성한 이전 최고 FID인 3.85를 3.04로 개선했습니다.

토큰 수가 증가하더라도 XL/2는 여전히 계산 효율성이 높습니다.
예를 들어, ADM은 1983 Gflops을 사용하고 ADM-U는 2813 Gflops을 사용합니다; XL/2는 524.6 Gflops을 사용합니다.
그림 1과 부록에 있는 고해상도 XL/2 모델의 샘플을 보여줍니다.

그림 10. 스케일링 샘플링 계산은 모델 계산의 부족을 보상하지 않습니다. 400K의 반복을 위해 학습된 각 DiT 모델에 대해 [16, 32, 64, 128, 256, 1000] 샘플링 단계를 사용하여 FID-10K를 계산합니다. 각 단계마다 FID와 각 이미지를 샘플링하는 데 사용되는 Gflops을 표시합니다. 소형 모델은 대형 모델보다 더 많은 테스트 시간 Gflops로 샘플링하더라도 대형 모델과의 성능 격차를 좁힐 수 없습니다.

5.2. Scaling Model vs. Sampling Compute

디퓨전 모델은 이미지를 생성할 때 샘플링 단계의 수를 늘려 학습 후 추가 계산을 사용할 수 있다는 점에서 독특합니다.
모델 Gflops이 샘플 품질에 미치는 영향을 고려할 때, 이 섹션에서는 더 많은 샘플링 컴퓨팅을 사용하여 더 작은 모델의 DiT가 더 큰 모델의 DiT보다 더 나은 성능을 발휘할 수 있는지 연구합니다.
우리는 이미지당 [16, 32, 64, 128, 256, 1000] 샘플링 단계를 사용하여 400K 학습 단계를 거친 후 12개의 DiT 모델 모두에 대한 FID를 계산합니다.
주요 결과는 그림 10에 나와 있습니다.

1000개의 샘플링 단계를 사용하는 DiT-L/2와 128개의 단계를 사용하는 DiT-XL/2를 비교해 보겠습니다.
이 경우 L/2는 각 이미지를 샘플링하는 데 80.7 Tflops을 사용하고, XL/2는 각 이미지를 샘플링하는 데 5배 적은 연산량—15.2 Tflops—을 사용합니다.
그럼에도 불구하고 XL/2는 더 나은 FID-10K (23.7 vs 25.9)를 가지고 있습니다.
일반적으로, 스케일링 샘플링 계산은 모델 계산의 부족을 보완할 수 없습니다.

 

 

6. Conclusion

디퓨전 모델을 위한 간단한 트랜스포머 기반 백본인 Diffusion Transformers (DiTs)를 소개합니다, 이 백본은 이전 U-Net 모델을 능가하며 트랜스포머 모델 클래스의 뛰어난 확장 특성을 계승합니다.
이 논문의 유망한 스케일링 결과를 고려할 때, 향후 연구에서는 DiT를 더 큰 모델과 토큰 수로 계속 확장해야 합니다.
DiT는 DALL·E 2 및 Stable Diffusion과 같은 text-to-image 모델을 위한 drop-in 백본으로도 탐구될 수 있습니다.