논문 <Scalable Diffusion Models with Transformers>의 내용을 정리한 글입니다.
이 논문은 Facebook Research(현 Meta AI)에서 공개한 논문이며, 저자 중 William Peebles는 OpenAI에서 Research Scientist로 일하고 있으며, 영상 생성 AI Sora의 개발을 공동으로 리드하고 있다고 한다.
DiT 아키텍처에 관한 관심은 OpenAI의 Sora 공개 이후 뜨거워진 것 같으며, 최근에 공개된 이미지 생성 모델인 Stable Diffusion 3, PixArt 계열의 모델에도 적용되어 더욱 각광받고 있다고 생각한다.
개괄적인 논문 요약
개요
이 연구는 이미지 생성을 위한 새로운 접근 방식으로 트랜스포머 기반 확산 모델을 제안한다. 기존의 U-Net 구조를 대체하는 트랜스포머 아키텍처를 사용하여, 확장성과 성능 면에서 우수한 결과를 달성하려 한다.
1. 소개
최근 딥러닝은 다양한 분야에서 획기적인 발전을 이루었으며, 특히 트랜스포머 모델은 그 가능성을 널리 입증하고 있다. 본 논문에서는 이미지 생성 분야에서 확산 모델의 기본 구조를 트랜스포머로 전환하는 새로운 시도를 소개한다. 이러한 전환을 통해 모델의 확장성과 성능을 크게 향상할 수 있다.
2. 관련 연구
2.1 트랜스포머의 발전
트랜스포머는 주로 자연어 처리 분야(NLP)에서 시작되었으나, 시간이 지남에 따라 시각 인식과 같은 다른 영역에서도 중요한 역할을 하게 되었다. 트랜스포머의 주요 강점은 대규모 데이터셋에 대한 우수한 학습 능력과 더불어 높은 확장성이다.
2.2 확산 모델(Diffusion Model)의 진화
확산 모델은 이미지 생성을 위한 강력한 방법으로 자리 잡았다. 기존 연구들은 주로 컨볼루셔널 U-Net 아키텍처(Latent Diffusion Model, LDM)에 기반하여 발전되어 왔지만, 본 연구는 트랜스포머를 사용하여 이러한 모델들을 뛰어넘고자 한다.
3. 확산 트랜스포머
3.1 기본 개념
확산 프로세스는 실제 이미지 데이터에 점진적으로 노이즈를 추가하고, 이를 역전시키는 과정을 모델링한다. 본 연구에서는 이러한 과정을 트랜스포머 아키텍처를 통해 수행한다.
3.2 모델 아키텍처
트랜스포머를 사용한 확산 모델의 핵심은 이미지를 패치로 분할하고 각 패치를 독립적인 토큰으로 처리하는 것이다. 이 토큰들은 트랜스포머의 인코더와 디코더를 통해 처리되며, 최종적으로 노이즈가 제거된 이미지를 생성한다.
4. 실험 설정
DiT 모델은 다양한 구성과 패치 크기를 가지고 ImageNet 데이터셋에서 훈련되었다. 모델 성능은 FID(Fréchet Inception Distance)를 통해 평가되었다.
5. 실험 결과
다양한 구성의 DiT 모델들이 훈련을 거치며 FID 점수가 개선되었음을 확인할 수 있었다. FID(Fréchet Inception Distance)는 생성된 영상이나 이미지의 품질을 평가하는데 자주 사용된다. 좌측의 그래프에서는 모델의 복잡도가 높아질수록, FID는 낮아져서 생성 품질이 좋아진다는 것을 설명한다. 우측의 그림은 기존의 U-Net 기반의 Diffusion Model과의 연산량을 비교하여 효율적이라는 것을 주장하는 그래프이다. 이는 트랜스포머 기반 구조가 기존 확산 모델을 뛰어넘는 성능을 달성할 수 있음을 의미한다.
결론
본 논문에서 제안한 트랜스포머 기반 확산 모델은 높은 확장성과 우수한 성능을 제공한다. 앞으로 이 모델을 텍스트-이미지 생성 모델 등 다양한 분야에 적용할 수 있는 가능성을 탐구할 예정이다.
논문 핵심 정리 & 공부
이미지 생성 모델링에서 표준 트랜스포머 아키텍처 도입
이 논문에서는 Diffusion 기반의 생성 모델링에서, 기존 U-Net 기반의 LDM 외에 미래로 나아갈 방향을 탐색해 보는 데에 중요한 역할을 했다고 생각한다. 그리고 U-Net의 inductive bias(고층부의 레이어에서는 디테일한 속성에 집중, 저층부에 가까울수록 더욱 coarse 한 특성에 집중, 논문에 자세히 설명은 되어있지 않음)가 생성 퀄리티에 핵심적인 부분은 아님을 보여주었다. 논문에서는 U-Net이 트랜스포머로 교체가 가능함을 보여주었고, 이를 통해 표준화된 transformer의 아키텍처를 도입할 수 있었다고 한다. 이는 이후에 영상 생성(시간 축에 대한 attend), 여러 도메인(cross-domain), 여러 모달리티(multi-modal)를 활용한 연구에도 영향을 주었다고 생각한다. 아키텍처 설계에 있어서 Vision Transformer의 모범 예시를 잘 따르도록 하였다고 하며, 그 덕에 논문 이름에 Scalable이라는 말을 추가한 것 같다.
그리고 세부적으로 Network complexity vs Sample quality 측면에서 scaling behavior를 비교해보았다고 한다. 그리고 VAE의 latent space에서 학습된 LDM과 비교하였다. 결론적으로는 Network complexity(Gflops로 측정된 값)이 높아질수록 Sample Quality(생성 퀄리티, FID로 측정)가 좋아진다(FID가 낮을수록 좋은 퀄리티를 뜻함)고 한다.
DiT 아키텍처
위의 아키텍처 그림을 보면, LDM처럼 latent 공간에서 ViT 아키텍처를 차용한 듯 보인다. 그리고, Cross Attention과 Multi-Head Self-Attention 또한 실험해 보았으나 adaLN-Zero 아키텍처를 최종적으로 선택하였다 (Adaptive Instance Normalization (AdaIN)을 사용한 Style GAN과 일부 유사하다고 볼 수 있음, 두 가지의 차이는 Adaptive 하게 normalization을 하지만, instance 차원에서 하는지, layer 차원에서 하는지이다).
디퓨전 모델의 기본 구조
이 논문의 "Diffusion formulation" 부분에서는 확산 모델, 특히 Gaussian diffusion models(가우시안 확산 모델)에 대한 기본 개념과 수학적 접근을 설명하고 있다. 여기에서 설명하는 주요 개념들을 간단하게 정리해보았다.
- 전방 노이즈 과정(Forward Noising Process):
- 가우시안 확산 모델은 실제 데이터 $x_0$에 점진적으로 노이즈를 적용하는 전방 과정 $q(x_t|x_0)$을 가정한다. 이 과정은 $x_0$에 가우시안 노이즈를 추가하여 새로운 상태 $x_t$를 생성한다. 수학적으로는 $q(x_t|x_0) = \mathcal{N}(x_t; \sqrt{\alpha_t}x_0, (1-\alpha_t)I)$로 표현되며, 여기서 $\alpha_t$는 하이퍼파라미터이다.
- 재매개변수화 트릭(Reparameterization Trick)을 사용하여 $x_t$를 샘플링할 수 있다. 이는 $x_t = \sqrt{\alpha_t}x_0 + \sqrt{1-\alpha_t}\epsilon_t$, $\epsilon_t \sim \mathcal{N}(0, I)$과 같이 표현된다.
- 역 과정(Reverse Process):
- 확산 모델은 전방 과정에 의해 손상된 데이터를 복구하는 역 과정을 학습한다. 이를 수학적으로는 $p_\theta(x_{t-1}|x_t) = \mathcal{N}(\mu_\theta(x_t), \Sigma_\theta(x_t))$로 표현하며, 여기서 $\mu_\theta$와 $\Sigma_\theta$는 신경망을 통해 예측된다.
- 역 과정 모델은 변분 하한(Variational Lower Bound)을 사용하여 훈련된다. 이는 로그-우도(log-likelihood) $\mathcal{L}(\theta) = -p(x_0|x_1) + \sum_t D_{KL}(q^*(x_{t-1}|x_t, x_0)||p_\theta(x_{t-1}|x_t))$를 최대화하는 형태로, 여기서 $D_{KL}$은 두 확률 분포 간의 Kullback-Leibler 발산을 나타낸다.
- 노이즈 예측 네트워크(Noise Prediction Network):
- 확산 모델은 $\mu_\theta$를 노이즈 예측 네트워크 $\theta$로 재매개변수화하여 간단한 평균 제곱 오차(mean-squared error)를 사용하여 훈련할 수 있다. 이를 통해 $\epsilon_\theta(x_t)$와 실제 샘플링된 가우시안 노이즈 $\epsilon_t$ 간의 차이를 최소화한다.
- 학습과 샘플링:
- 모델이 훈련되면, 최대 시간 $t_{\text{max}}$에서 $x_{t_{\text{max}}} \sim \mathcal{N}(0, I)$을 초기화하여 새로운 이미지를 샘플링할 수 있다. 이후 $x_{t-1} \sim p_\theta(x_{t-1}|x_t)$를 사용하여 역방향으로 샘플링을 진행한다.
Classifier-free Guidance의 개념
논문의 섹션 3에서는 조건부 확산 모델(conditional diffusion models)에서의 향상된 샘플링 절차를 설명한다. 이 접근 방식은 특정 클래스 라벨 $c$와 같은 추가적인 정보를 입력으로 사용하며, 이는 역 과정 $p_\theta(x_{t-1}|x_t, c)$에 영향을 미치게 된다.
- 조건부 입력의 사용:
- 조건부 확산 모델은 클래스 라벨과 같은 추가 정보를 입력으로 사용한다. 이러한 정보는 역 과정 $p_\theta(x_{t-1} | x_t, c)$에서 모델의 예측 $\epsilon_\theta$와 분산 $\Sigma_\theta$가 조건부 입력 $c$에 따라 조정되도록 한다.
- Classifier-free guidance의 역할:
- Diffusion 모델의 output을 score function으로 생각하는 경우, 이 접근법(DDPM 샘플링 절차)은 $p(c|x)$가 높은 $x$를 샘플링 하도록 유도된다. $p(c|x)$는 $x$가 주어졌을 때 클래스 $c$에 속할 확률을 의미한다.
- Bayes Rule에 따라 $\log p(c|x) \propto \log p(x|c) - \log p(x)$ 이므로, $\nabla_x \log p(c|x) \propto \nabla_x \log p(x|c) - \nabla_x \log p(x)$로 계산된다. 이는 클래스 라벨 $c$가 주어졌을 때 $x$의 생성을 유도하는 방향으로 작용한다.
- 샘플링 절차:
- $\hat{\epsilon}_\theta(x_t, c) = \epsilon_\theta(x_t, \phi) + s \cdot (\epsilon_\theta(x_t, c) - \epsilon_\theta(x_t, \phi))$와 같은 식을 사용하여 샘플링을 유도한다. 여기서 $\phi$는 라벨이 없는 경우를 대비한 학습된 "null" 임베딩이고, $s$는 유도의 강도를 나타내는 스케일 파라미터이다.
- $s > 1$은 유도(가이던스)의 강도를 나타내며, $s = 1$일 때는 표준 샘플링과 동일하다.
- 평가:
- $c = \phi$ 가 포함된 디퓨전 모델의 성능을 평가할 수 있다. 훈련 중에는 무작위로 $c$를 드롭아웃시키고 "null" 임베딩($\phi$)으로 대체하여 이를 학습한다.
Classifier-free Guidance의 중요성
- 이 방법은 일반적인 샘플링 기법보다 훨씬 개선된 샘플을 생성할 수 있으며, DiT(확산 트랜스포머) 모델에도 동일하게 적용되어 효과를 입증했다. 이는 특히 생성 모델이 특정 클래스의 특성을 더 잘 반영하도록 할 때 유용하다.
Latent Diffusion Model(LDM)
디퓨전 모델을 고해상도의 이미지 pixel 공간에서 학습하는 것은 비효율적이고, 비용이 많이 든다. 따라서, LDM에서는 2가지 stage로 이를 극복하였다.
- Auto Encoder:
- 이미지를 더욱 작은 representational space로 보내는 AutoEncoder를 학습하고, latent space에서 학습을 진행한다. 즉, 고해상도 이미지를 압축할 수 있는 Auto Encoder($E$)를 도입하였다.
- Latent Space에서 모델 학습:
- 이미지 $x$를 바로 만들어내는 모델을 학습하지 않고, Latent $z$를 모델링한다. $z = E(x)$ 이며, 이때 $E$는 frozen. 이후, 이 $z$를 디코딩하는 디코더 ($D$)를 사용하여, 디퓨전 모델이 만들어낸 latent $z$로부터 이미지 $x$를 생성한다 ($x= D(z)$).
Paper에서 언급은 명확히 하지 않았지만, U-Net 아키텍처를 추가로 채택하여, 더욱 저차원에서도 학습을 진행한다.
기술적인 특징 요약
이 논문의 기술적인 설명들을 종합해 보면, off-the-shelf의 Convolutional VAE를 활용하는 Transformer-based DDPM이라고 요약할 수 있다.
Diffusion Transformer Design Space
DiT를 설계할 때, 표준적인 ViT의 scaling property들을 유지하고자 했다고 한다. 따라서, ViT의 모범 사례들을 잘 유지하고 있다고 하며, 이 챕터에서는 patchify, DiT Block design, Model size, Transformer Decoder에 대해서 자세히 설명한다.
- patchify: 패치 사이즈 2, 4, 8로 모델을 학습하였으며, 전체적인 모델의 연산량 Gflops에 큰 영향을 끼치지만, 이후 모델의 파라미터 개수에는 영향을 끼치지 않는다.
- DiT Block Design: Diffusion 모델은 noised image inputs을 주요 입력으로 받지만, 추가적인 조건부 입력을 받는다. 예를 들자면, noise timesteps $t$, class labels $c$, natural language 등이 있다. 따라서, 이러한 추가적인 입력들을 고려하여, ViT 아키텍처를 수정하였다고 한다.
- In-context conditioning: ViT의 CLS 토큰과 유사한 방식으로, 추가적인 조건부 입력 $t$, $c$ 토큰을 input sequence 마지막에 붙여서 contextualizing을 시켜보았다고 한다. 이후, 마지막 block에서 이 추가적인 토큰을 제거한다. 이 방식은 Gflops에는 거의 영향을 미치지 않았다고 한다. Figure 3의 마지막 block design
DiT Block with In-Context Conditioning에 대한 설명이다. - Cross-attentionblock: MSA(Multi-Head Self Attention) 블록 뒤에 Cross Attention 블록을 추가하였다고 한다. 이는 전체 모델의 연산량(Gflops)의 15%를 차지할 정도로 큰 영향을 끼친다. 이는 Figure 3의 두 번째 DiT 블록
DiT Block with Cross-Attention에 대한 설명으로 파악하였다. - Adaptive layer norm (adaLN) block: GAN에서부터 자주 사용되는 Adaptive Layer Norm 블록을 적용하였으며, 가장 계산 효율적이었다고 한다.
- adaLN-Zero block: 큰 규모로 지도학습을 진행할 때, 마지막 batch norm의 scale factor를 0으로 초기화(zero-initializing)하는 것이 학습을 더욱 용이하게 한다는 것을 밝힌 선행 연구를 참고하였다고 한다. 그래서 adaLN DiT 블럭을 변형하여, scale factor $\gamma$, $\beta$ 외에도 dimension-wise scaling factor $\alpha$를 추가로 도입하여 이를 적용하였다고 한다(자세한 내용은 안나옴..). MLP 레이어가 모든 $\alpha$에 대하여 zero-vector를 출력하도록 초기화하였다고 한다. 이 방식이 전체 DiT 블럭이 identity function이 되도록 만들어 준다고 한다. 기존 adaLN에 비해 연산량의 차이는 없다고 한다.
- In-context conditioning: ViT의 CLS 토큰과 유사한 방식으로, 추가적인 조건부 입력 $t$, $c$ 토큰을 input sequence 마지막에 붙여서 contextualizing을 시켜보았다고 한다. 이후, 마지막 block에서 이 추가적인 토큰을 제거한다. 이 방식은 Gflops에는 거의 영향을 미치지 않았다고 한다. Figure 3의 마지막 block design
- Model Size: DiT block의 개수($N$)와 각 block의 hidden dim size $d$ 그리고 Attention Head 개수를 변경해가면서 테스트를 했다고 한다. 모델 크기가 커지는 순서로 DiT-S, DiT-B, DiT-L and DiT-XL라고 명명하고 있다.
- Transformer Decoder: 마지막 DiT 블럭 이후, 이미지 토큰을 output noise prediction과 output diagonal covariance prediction으로 디코딩을 해야했다고 한다. 이 모델에서는 Linear Decoder를 학습하여, 사용하였다고 함(linear projection or fully connected ?).
실험 세팅 및 결과
실험 세팅
- Training:
- class-conditional latent DiT models를 이미지넷의 256x256, 512x512 데이터셋으로 학습
- 마지막 linear layer를 0으로 초기화, 나머지 weight init은 vit의 표준을 따른다.
- AdamW 사용.
- Diffusion: Stable diffusion 모델에 쓰인 VAE 사용, ADM의 하이퍼파라미터 세팅을 그대로 이용.
- Evaluation metrics: FID, Inception Score(IS), sFID, Precision/Recall(class 기반 생성이기 때문) 등을 사용
실험 결과
- DiT Block 비교: adaLN-zero이 Gflops가 가장 적으면서도 FID가 가장 낮게(좋게) 나옴. 이후, 이 블럭 아키텍처로 고정 후 뒤의 실험 진행.
- Scaling model size and patch size: 모델 사이즈를 키우고, 패치 사이즈를 작게 만드는 것이 성능을 향상시키는 것을 관찰함.
- DiT Gflops are critical to improving performance: 말 그대로, 모델 연산량이 전체 성능 향상에 결정적이었다고 한다.
- Larger DiT models are more compute-efficien: 어느 정도의 Gflops 이상으로 가면, 큰 모델이 성능 대비 연산량 또한 효율적이었다고 한다.
- Scaling Model vs. Sampling Compute: 여기서는 모델의 capacity or flexibility가 작을 경우 sampling step을 더욱 높여도, 모델 자체의 성능을 극복할 수 없음을 보여주었다.
