[논문 리뷰] DiffiT: Diffusion Vision Transformers for Image Generation

Ali Hatamizadeh, Jiaming Song외 3명|arXiv (Cornell University)|2023. 12. 04.

Generative Adversarial Networks and Image Synthesis인용 수 7

한 줄 요약

이 논문은 전통적인 U-Net 아키텍처를 대체하는 계층적 트랜스포머 인코더-디코더를 갖춘 비전 트랜스포머 기반의 확산 모델인 DiffiT를 소개한다. 시간에 따라 변화하는 자기주의 주의(TMSA) 메커니즘을 제안하여 다양한 노이즈 제거 단계에 동적으로 적응하도록 하여, 잠재 공간에서 ImageNet-256 기준으로 SOTA FID 점수 1.73, ImageNet-512 기준으로 2.67을 달성한다.

ABSTRACT

Diffusion models with their powerful expressivity and high sample quality have achieved State-Of-The-Art (SOTA) performance in the generative domain. The pioneering Vision Transformer (ViT) has also demonstrated strong modeling capabilities and scalability, especially for recognition tasks. In this paper, we study the effectiveness of ViTs in diffusion-based generative learning and propose a new model denoted as Diffusion Vision Transformers (DiffiT). Specifically, we propose a methodology for finegrained control of the denoising process and introduce the Time-dependant Multihead Self Attention (TMSA) mechanism. DiffiT is surprisingly effective in generating high-fidelity images with significantly better parameter efficiency. We also propose latent and image space DiffiT models and show SOTA performance on a variety of class-conditional and unconditional synthesis tasks at different resolutions. The Latent DiffiT model achieves a new SOTA FID score of 1.73 on ImageNet256 dataset while having 19.85%, 16.88% less parameters than other Transformer-based diffusion models such as MDT and DiT,respectively. Code: https://github.com/NVlabs/DiffiT

연구 동기 및 목표

• 확산 기반 이미지 생성에서 컨볼루션 기반 U-Net 아키텍처를 넘어서 비전 트랜스포머의 효과성을 조사한다.
• 기존 확산 모델에서 시간에 따른 세밀한 모델링이 부족한 문제를 해결한다. 즉, 자기주의 주의 메커니즘이 다양한 노이즈 제거 단계에 적응하지 못한다.
• 이미지 및 잠재 공간 생성을 위한 계층적 트랜스포머 기반 아키텍처를 도입함으로써 노이즈 제거 네트워크 설계 패턴을 통합한다.
• 다양한 벤치마크에서 FID 점수로 측정된 이미지 생성 품질 측면에서 SOTA 성능을 달성한다.

제안 방법

• 자기주의 주의의 쿼리, 키, 밸류 프로젝션에 시간 스텝 정보를 직접 통합하는 새로운 시간에 따라 변화하는 자기주의 주의(TMSA) 모듈을 제안하여, 노이즈 제거 단계에 따라 동적으로 적응할 수 있도록 한다.
• 특징 계층을 유지하고 다양한 스케일에서 효과적인 특징 정제를 가능하게 하는 계층적, U-형태의 트랜스포머 인코더-디코더 아키텍처(DiffiT)를 설계한다.
• 잠재 공간에 TMSA 모듈을 적용하여 고해상도 이미지 생성을 가능하게 하는 잠재 DiffiT를 도입한다. 이는 효율적이고 고해상도의 합성 성능을 제공한다.
• 시간 토큰을 학습 가능한 투영을 통해 자기주의 주의 메커니즘에 통합하여, 현재 노이즈 제거 단계에 따라 주의 행동을 조절할 수 있도록 한다.
• 상대적 위치 편향을 시간 토큰과 조합하여 공간적 인덕티브 바이어스를 유지하면서도 시간에 따른 적응성을 확보한다.
• 조건부 처리를 위해 적응형 레이어 정규화(AdaLN)를 사용하지 않고, 시간 의존성을 주로 주의 메커니즘에 직접 통합한다.

Figure 2 : Overview of the image-space DiffiT model. Downsample and Upsample denote convolutional downsampling and upsampling layers, respectively. Please see the supplementary materials for more information regarding the DiffiT architecture.

실험 결과

연구 질문

• RQ1 비전 트랜스포머가 확산 기반 이미지 생성에서 U-Net 아키텍처를 효과적으로 대체할 수 있을까? 이는 샘플 품질을 유지하거나 향상시킬 수 있는가?
• RQ2 확산 모델의 자기주의 주의 메커니즘을 어떻게 다양한 노이즈 제거 시간 스텝에 적응시킬 수 있을까? 이를 통해 변화하는 이미지 구조를 더 잘 포착할 수 있는가?
• RQ3 자신의 주의 메커니즘에 시간 의존성을 직접 통합할 경우, 단순한 위치 임베딩이나 정규화 레이어를 통한 방법과 비교해 어떤 영향을 미칠까?
• RQ4 계층적 설계를 가진 통합된 트랜스포머 기반 아키텍처가 이미지 및 잠재 공간 생성 모두에서 기존의 U-Net 및 DiT 스타일 모델을 초월할 수 있을까?
• RQ5 제안된 시간에 따라 변화하는 자기주의 주의 메커니즘이 CIFAR10, FFHQ-64, ImageNet-256/512와 같은 다양한 벤치마크에서 FID 점수에 얼마나 기여하는가?

주요 결과

• 제안된 시간에 따라 변화하는 자기주의 주의(TMSA) 모듈은 표준 자기주의 주의를 대체함으로써, VE 설정 기준으로 CIFAR10에서 FID 점수를 0.28 감소시키고, VP 설정 기준으로는 0.25 감소시킨다.
• 잠재 DiffiT는 ImageNet-256 데이터셋에서 기존 방법들인 DDPM++ 및 MDT를 모두 능가하는 새로운 SOTA FID 점수 1.73을 달성한다.
• 제거 실험 결과 TMSA가 다른 대안보다 훨씬 효과적임을 확인할 수 있었다. TMSA에서 시간 의존성을 제거하면 FID가 3.97로 상승하고, 시간 토큰을 MLP 레이어에만 적용할 경우 FID는 3.81로 나타났다.
• 질적 주의 맵 시각화 결과 TMSA가 노이즈 제거 단계에 따라 동적으로 적응하며 세밀한 물체의 세부 사항을 잘 포착하는 것으로 나타났다. TMSA가 없는 모델과는 대조적으로 이와 같은 특징을 보였다.
• 파wer-코사인 스케줄을 사용한 분류기 없는 가이던스는 ImageNet-256에서 FID를 1.73으로 향상시켰고, 최적의 가이던스 스케일은 4.6이었다. 이보다 높은 스케일은 성능 저하를 초래했다.
• 에코더와 디코더 블록을 모두 포함한 전체 DiffiT 아키텍처는 CIFAR10에서 FID 1.95를 기록하여, 고해상도 생성을 위한 완전한 U-형 아키텍처의 중요성을 입증했다.

Figure 3 : The DiffiT Transformer block applies linear projection to spatial and time-embedding tokens before combining them together to form query, key, and value vectors for each token. These vectors are then used to compute multi-head self-attention activations, followed by two linear layers.