[논문 리뷰] Your ViT is Secretly a Hybrid Discriminative-Generative Diffusion Model

Xiulong Yang, Sheng-Min Shih외 3명|arXiv (Cornell University)|2022. 08. 16.

Generative Adversarial Networks and Image Synthesis인용 수 15

한 줄 요약

이 논문은 이미지 생성과 분류를 동시에 수행할 수 있는 새로운 모델인 GenViT와 HybViT를 제안한다. 이 모델들은 디퓨전 생성 모델링과 비전 트랜스포머(ViT)를 통합하여, 하나의 ViT 백본을 통해 이미지 생성과 분류 작업을 모두 수행한다. DDPM 프레임워크에 ViT를 통합함으로써, 기존의 하이브리드 모델보다 안정적인 MCMC를 사용하지 않는 훈련과 향상된 계산 효율성을 확보하면서도 두 작업에서 최신 기술 수준(SOTA)의 성능을 달성한다.

ABSTRACT

Diffusion Denoising Probability Models (DDPM) and Vision Transformer (ViT) have demonstrated significant progress in generative tasks and discriminative tasks, respectively, and thus far these models have largely been developed in their own domains. In this paper, we establish a direct connection between DDPM and ViT by integrating the ViT architecture into DDPM, and introduce a new generative model called Generative ViT (GenViT). The modeling flexibility of ViT enables us to further extend GenViT to hybrid discriminative-generative modeling, and introduce a Hybrid ViT (HybViT). Our work is among the first to explore a single ViT for image generation and classification jointly. We conduct a series of experiments to analyze the performance of proposed models and demonstrate their superiority over prior state-of-the-arts in both generative and discriminative tasks. Our code and pre-trained models can be found in https://github.com/sndnyang/Diffusion_ViT .

연구 동기 및 목표

• 단일 비전 트랜스포머(ViT)가 DDPM에서 전통적으로 사용되는 UNet 대신 생성 모델로 효과적으로 사용될 수 있는지 탐구한다.
• 단일 ViT가 이미지 분류와 생성을 동시에 수행할 수 있도록, 분류 모델과 생성 모델 간의 격차를 해소한다.
• JEM 및 JEM++와 같은 MCMC 기반 하이브리드 모델의 불안정성과 높은 계산 비용 문제를 해결한다.
• CNN 기반 또는 GAN 기반 대안 대비 ViT 기반 디퓨전 모델이 더 뛰어난 확장성, 안정성, 효율성을 제공함을 보여준다.
• 기존 기준 모델과 비교하여 제안된 모델의 강건성, 불확실성 캘리브레이션 및 OOD 탐지 성능을 분석한다.

제안 방법

• 표준 UNet 백본 대신 ViT를 사용하여 디퓨전 과정에서 노이즈 제거를 수행함으로써, 비전 트랜스포머 아키텍처를 직접 DDPM 프레임워크에 통합한다.
• 변분 하한(VLB) 목적함수를 사용하여 역디퓨전 과정에서 노이즈를 예측하도록 훈련함으로써, 순수 생성 모델인 GenViT를 설계한다.
• 공유된 ViT 특징 공간과 두 개의 독립된 헤드를 사용하여 이미지 재구성(디퓨전을 통한)과 분류를 동시에 최적화하는 하이브리드 모델인 HybViT를 개발한다.
• 생성 및 분류 작업 모두에 동일한 ViT 파라미터를 공유함으로써, 파라미터 효율성과 통계적 강도 공유를 달성한다.
• MCMC 샘플링이 필요 없도록 재구성 손실과 크로스엔트로피 손실을 통합한 단일 목적함수를 사용한다.
• 학습 가능한 선형 투영과 다중 헤드 자기주의를 사용하여 패치 기반 이미지 처리를 수행함으로써, 잠재 디퓨전 공간 내 장거리 의존성을 모델링한다.

실험 결과

연구 질문

• RQ1 단일 비전 트랜스포머가 표준 UNet 대신 디퓨전 기반 이미지 생성에 효과적으로 사용될 수 있는가?
• RQ2 공유된 ViT 백본에 생성 및 분류 헤드를 함께 훈련할 경우, 이미지 생성 및 분류 작업의 성능에 어떤 영향을 미치는가?
• RQ3 제안된 방법이 JEM 및 JEM++와 같은 MCMC 기반 하이브리드 모델보다 훈련 안정성, 수렴 속도 및 최종 성능 측면에서 뛰어나지 않는가?
• RQ4 패치 크기와 모델 차원이 제안된 모델에서 생성 품질과 분류 정확도 간의 상호 보완적 관계에 어떤 영향을 미치는가?
• RQ5 GenViT와 HybViT가 최신 기술 수준의 모델과 비교할 때, 가능도, 적대적 공격에 대한 강건성, 불확실성 캘리브레이션 및 OOD 탐지 성능 측면에서 어떻게 비교되는가?

주요 결과

• HybViT는 STL10에서 87.1%의 테스트 정확도, Inception Score 6.90, FID 125.5를 기록하여 기존 하이브리드 모델을 초월한다.
• 모델 차원이 768인 GenViT는 CIFAR-10에서 Inception Score 7.01과 FID 126.6을 기록하여 뛰어난 생성 성능를 입증한다.
• HybViT는 MCMC 기반 모델보다 훨씬 빠른 훈련 속도를 보이며, 단일 GPU로 CIFAR-10에서 31.2시간 만에 훈련을 완료했고, JEM(K=20)의 101.3시간 대비 빠른 성능을 보였다.
• 패치 크기가 6인 모델(HybViT-ps6)이 정확도와 생성 품질 간 최적의 균형을 이룩하여 81.7%의 정확도와 7.30의 IS를 기록했다.
• 기본 모델 대비 HybViT는 더 뛰어난 불확실성 캘리브레이션 및 OOD 탐지 성능를 보이며, 강건성 향상을 시사한다.
• 강력한 성능에도 불구하고 고해상도 이미지의 생성 품질은 패치 크기 6에서 정점에 도달하며, 이는 선형 투영층의 한계일 가능성이 높으며, 자기주의 메커니즘의 문제가 아님을 시사한다.