[논문 리뷰] All are Worth Words: A ViT Backbone for Diffusion Models

Fan Bao, Nie, Shen외 5명|arXiv (Cornell University)|2022. 09. 25.

Generative Adversarial Networks and Image Synthesis인용 수 7

한 줄 요약

이 논문은 시간, 조건, 노이즈 이미지 패치를 토큰으로 간주하고 얕은 층과 깊은 층 사이에 장거리 스킵 연결을 사용하는, 확산 모델을 위한 비전 트랜스포머(ViT)-기반 백본인 U-ViT를 제안한다. U-ViT는 이미지 품질을 높이기 위해 내부적으로 다운/업샘플링 연산을 필요로 하지 않으며, ImageNet 256×256의 클래스 조건부 생성에서 FID 점수 2.29, MS-COCO 텍스트-이미지 생성에서 FID 점수 5.48을 기록하여 최신 기술 수준(SOTA)을 달성한다. 이는 스킵 연결을 갖춘 ViT 백본이 CNN 기반 U-Net보다 우수한 성능을 낼 수 있음을 보여준다.

ABSTRACT

Vision transformers (ViT) have shown promise in various vision tasks while the U-Net based on a convolutional neural network (CNN) remains dominant in diffusion models. We design a simple and general ViT-based architecture (named U-ViT) for image generation with diffusion models. U-ViT is characterized by treating all inputs including the time, condition and noisy image patches as tokens and employing long skip connections between shallow and deep layers. We evaluate U-ViT in unconditional and class-conditional image generation, as well as text-to-image generation tasks, where U-ViT is comparable if not superior to a CNN-based U-Net of a similar size. In particular, latent diffusion models with U-ViT achieve record-breaking FID scores of 2.29 in class-conditional image generation on ImageNet 256x256, and 5.48 in text-to-image generation on MS-COCO, among methods without accessing large external datasets during the training of generative models. Our results suggest that, for diffusion-based image modeling, the long skip connection is crucial while the down-sampling and up-sampling operators in CNN-based U-Net are not always necessary. We believe that U-ViT can provide insights for future research on backbones in diffusion models and benefit generative modeling on large scale cross-modality datasets.

연구 동기 및 목표

• 비전 트랜스포머(ViT)가 확산 모델의 CNN 기반 U-Net 백본을 효과적으로 대체할 수 있는지 조사하기 위해.
• 확산 기반 이미지 생성에서 U-Net 아키텍처의 다운샘플링 및 업샘플링 연산이 필수적인지 평가하기 위해.
• ViT 기반 확산 모델 백본에서 장거리 스킵 연결의 영향을 규명하기 위해.
• 간단한 ViT 기반 아키텍처를 사용하여 무조건적, 클래스 조건부, 텍스트-이미지 생성 과제에서 최신 기술 수준의 성능을 달성하기 위해.

제안 방법

• U-ViT는 노이즈 이미지 패치, 타임스텝 임베딩, 조건 임베딩을 모두 토큰으로 간주하여 통합된 토큰 기반 입력 표현을 가능하게 한다.
• 확산 모델에서 픽셀 수준 복원에 중요한 저수준 특징을 유지하기 위해 얕은 층과 깊은 층 사이에 장거리 스킵 연결을 적용한다.
• 클래스 토큰과 위치 임베딩을 갖춘 표준 ViT 트랜스포머 인코더를 사용하고, 학습된 위치 임베딩을 갖춘 디코더 헤드를 이어받는다.
• 최종 출력 이전에 선택적 3×3 컨볼루션 블록을 추가하여 시각적 품질을 향상시킨다.
• 모델은 픽셀 스페이스 및 레이튼 스페이스 확산 모델에 모두 적용되며, 노이즈 예측 목표함수를 사용해 훈련된다.
• CIFAR-10, ImageNet 256×256, ImageNet 512×512, MS-COCO 등 다양한 벤치마크에서 FID 및 기타 지표를 사용해 평가된다.

실험 결과

연구 질문

• RQ1 CNN 기반 U-Net을 대체하는 비전 트랜스포머 백본이 성능 저하 없이 확산 모델에서 효과적으로 작동할 수 있는가?
• RQ2 U-Net 아키텍처의 다운샘플링 및 업샘플링 메커니즘이 확산 모델에서 고품질 이미지 생성에 필수적인가?
• RQ3 표준 ViT 아키텍처에 비해 ViT 기반 확산 모델 백본에서 장거리 스킵 연결이 얼마나 중요한가?
• RQ4 대규모 외부 데이터셋에 의존하지 않고도 ViT 기반 아키텍처가 텍스트-이미지 및 클래스 조건부 이미지 생성에서 최신 기술 수준의 성능을 달성할 수 있는가?

주요 결과

• U-ViT는 외부 대규모 데이터셋에 접근하지 못한 상태에서도 클래스 조건부 ImageNet 256×256 생성에서 최신 기술 수준의 FID 점수 2.29를 기록하였다.
• MS-COCO 텍스트-이미지 생성에서 U-ViT-S (Deep)는 새로운 최신 기술 수준의 FID 5.48을 달성하여 우수한 의미적 일치도와 이미지 품질을 보였다.
• U-ViT는 크기와 FLOPs가 유사한 CNN 기반 U-Net보다 무조건적 및 조건부 이미지 생성 과제에서 모두 뛰어난 성능을 보였다.
• 추이적 분석 결과, 장거리 스킵 연결이 성능에 매우 중요하며, U-Net의 다운/업샘플링 연산은 항상 고품질 생성에 필수적인 것은 아님을 확인하였다.
• 특히 텍스트-이미지 생성 과제에서 U-ViT는 U-Net보다 더 의미적으로 정확한 샘플을 생성하였으며, 프롬프트에서 언급된 물체들(예: 박스, 볼 등)을 더 잘 포착하였다.
• 모델은 다양한 샘플링 스텝과 인셉션 스코어, 정밀도, 재현율 등의 지표에서도 강력하고 고품질의 생성 성능을 유지하였다.