[논문 리뷰] Exploring Vision Transformers as Diffusion Learners

He Cao, Jianan Wang외 5명|arXiv (Cornell University)|2022. 12. 28.

Generative Adversarial Networks and Image Synthesis인용 수 7

한 줄 요약

이 논문은 비대칭 인코더-디코더 아키텍처를 가진 비전 트랜스포머 기반의 확산 백본인 ASCEND를 제안한다. 이는 이미지 생성 작업에서 성능을 크게 향상시킨다. U-Net 대신 강력한 Swin 트랜스포머 인코더와 경량 컨볼루션 디코더를 사용함으로써, CIFAR-10, CelebA, LSUN, CUB-Bird에서 경쟁적인 성능을 달성하며, 64×64 해상도를 초월하는 단일 단계 텍스트-이미지 확산 모델을 훈련한 최초의 사례로, 엔드 투 엔드 고해상도 생성 분야의 새로운 기준을 설정한다.

ABSTRACT

Score-based diffusion models have captured widespread attention and funded fast progress of recent vision generative tasks. In this paper, we focus on diffusion model backbone which has been much neglected before. We systematically explore vision Transformers as diffusion learners for various generative tasks. With our improvements the performance of vanilla ViT-based backbone (IU-ViT) is boosted to be on par with traditional U-Net-based methods. We further provide a hypothesis on the implication of disentangling the generative backbone as an encoder-decoder structure and show proof-of-concept experiments verifying the effectiveness of a stronger encoder for generative tasks with ASymmetriC ENcoder Decoder (ASCEND). Our improvements achieve competitive results on CIFAR-10, CelebA, LSUN, CUB Bird and large-resolution text-to-image tasks. To the best of our knowledge, we are the first to successfully train a single diffusion model on text-to-image task beyond 64x64 resolution. We hope this will motivate people to rethink the modeling choices and the training pipelines for diffusion-based generative models.

연구 동기 및 목표

• 비전 트랜스포머를 확산 모델 백본으로 체계적으로 평가하고, 특히 U-Net과의 비교를 통해 성능을 분석한다.
• 기본적인 ViT와 U-Net 간의 저해상도 이미지 생성에서의 성능 격차를 줄이기 위해 U-ViT 아키텍처를 개선함으로써 IU-ViT를 제안한다.
• 확장 가능하고 고해상도이며 다중 모odal 확산 생성에 적합한 비대칭 인코더-디코더 아키텍처(ASCEND)를 제안하고 검증한다.
• 다단계 훈련 파이프라인의 관행을 도전하기 위해, 비전 트랜스포머 백본을 사용한 엔드 투 엔드 단일 단계 고해상도 텍스트-이미지 생성을 입증한다.
• 확산 기반 생성 모델링에서 백본 설계와 훈련 파라다임을 재고할 것을 연구 공동체에 촉구한다.

제안 방법

• 패치 머징/확장 대신 스킵 커넥션과 잔차 다운/업샘플링을 향상시킨 IU-ViT를 제안한다.
• 스위프트 트랜스포머 인코더와 컨볼루션 디코더를 사용한 비대칭 인코더-디코더 백본인 ASCEND를 설계한다.
• 특징 계층과 기울기 흐름을 향상시키기 위해 패치 머징/확장 대신 잔차 다운/업샘플링을 사용한다.
• 텍스트 조건부 텍스트-이미지 생성 작업에서 텍스트를 위한 사전 훈련된 T5-XXL 인코더를 통합한다.
• LAION-Aesthetics와 같은 대규모 데이터셋에서 128×128 및 256×256 해상도 이미지 생성을 위해 확산 모델을 엔드 투 엔드로 훈련한다.
• CIFAR-10에서 아키텍처 선택 사항을 검증하기 위해 스킵 커넥션, 다운샘플링 방법, 인코더/디코더 구성 등을 포함한 아블레이션 스터디를 수행한다.

실험 결과

연구 질문

• RQ1 비전 트랜스포머는 확산 모델 백본으로 효과적으로 사용될 수 있으며, 저해상도 및 고해상도 생성에서 U-Net과 비교해 어떻게 성능을 내는가?
• RQ2 기본적인 ViT 기반 백본이 U-Net과 경쟁 가능하도록 하기 위해 어떤 아키텍처 개선이 필요한가?
• RQ3 비대칭 인코더-디코더 구조는 확산 기반 이미지 생성에서 성능과 확장성에 기여하는가?
• RQ4 비전 트랜스포머 백본을 사용한 단일 단계 엔드 투 엔드 확산 모델이 64×64 해상도를 초월하는 고해상도에서 고품질 이미지를 생성할 수 있는가?
• RQ5 인코더 아키텍처 선택(예: Swin 트랜스포머)이 고해상도 및 텍스트-이미지 생성에서 성능에 어떤 영향을 미치는가?

주요 결과

• ASCEND는 CIFAR-10에서 Fréchet Inception Distance(FID) 2.98을 기록하여 잘 튜닝된 U-Net 백본의 성능을 그대로 따라간다.
• ASCEND에서 패치 머징/확장을 잔차 다운/업샘플링으로 대체함으로써 FID가 12.81에서 2.98로 감소하여 아키텍처적 개선의 효과를 입증한다.
• ASCEND에서 스킵 커넥션을 줄이면 FID가 3.54점 상승하여 특징 재구성에서의 중요성을 확인한다.
• 590M 파라미터 모델을 사용해 고해상도 128×128 텍스트-이미지 샘플을 생성하며, 64×64 해상도를 초월하는 단일 단계 훈련에서 성공한 최초의 사례이다.
• ASCEND는 고해상도 생성에서 IU-ViT와 U-Net을 모두 능가하며, 후기 훈련 단계에서 더 빠른 수렴과 더 나은 샘플 품질을 보여준다.
• 아블레이션 스터디는 인코더에 SwinBlocks를 사용하고 디코더에 컨볼루션 구조를 사용할 경우 대칭 또는 다른 구성보다 열등한 결과를 낳는다.