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

오늘 소개할 논문은 "Your ViT is Secretly a Hybrid Discriminative-Generative Diffusion Model"이라는 논문이다.

https://arxiv.org/abs/2208.07791

Your ViT is Secretly a Hybrid Discriminative-Generative Diffusion Model

 

 

 

최근 Diffusion nodel 스터디를 진행하면서 스터디 내용을 블로그로 정리한다는게 늦어지고 늦어져 첫 논문 리뷰 포스터로 기본 Diffusion model이 아닌 응용 논문으로 시작하게 되었다. 그리고 사실 이 논문은 읽으면서 Diffusion 논문보다는 Hybrid model 혹은 ViT 논문이라는 생각을 강하게 받은 논문이었다. (지극히 나의 주관적인 느낌이다.)

 

Introduction

이 논문을 읽으면서 Hybrid model이나 JEM++과 같은 (내 기준) 다소 생소한 용어들을 많이 접했는데, 논문을 다 읽고 저자에 대해 검색해니 조금이나마 그 이유를 알 것만 같았다. 아래는 저자 Xiulong Yang에 대한 논문 기록이다.

논문 저자 Xiulong Yang에 대한 검색 결과

저자의 논문 기록을 보면 이번에 소개할 논문 이전에 Classification과 Generation을 동시에 하는 Hybrid model에 대해 연구해 왔으며, JEM을 발전시킨 JEM++이라는 논문도 작성한 것을 볼  수가 있다. 즉, Hybrid model을 계속 연구해 왔고, 본인이 연구한 JEM++이라는 모델을 이번 논문에 등장시켰던 것이다. 앞서 이번 논문이 Diffusion model 논문보다는 Hybrid model 논문에 더 가깝다고 생각했던 이유 중 하나이다. 아마 Hybrid model에 대한 또 다른 방법론으로 ViT를 사용하고자 하였으며, 단일 ViT로는 힘들어 Diffusion model의 방법론을 가져온 것이 아닌가라는 생각이 든다. 그리고 후술 하겠지만 Hybrid model의 성능이 단일 모델들의 성능을 능가하지 못하는데, 아마 이러한 이유 때문에 이 저자는 더 이상 Hybrid model에 대한 연구를 그만둔 것이 아닌가라는 생각이 들었다.

 

Method

사실 ViT와 Diffusion model에 대해 안다면 method로 설명할 것이 거의 없다. 기존 Diffusion model에서 각 timestep 사이의 Unet 구조를 ViT로 치환한 것이 이번 논문의 핵심이다. 또한, ViT에서는 cls token을 이용한 classification이 가능한데, HybViT에서는 Unet 구조를 ViT로 바꾸면서 ViT의 classification 능력을 같이 활용하였다.

GenViT와 HybViT

위의 사진은 논문에 나와 있는 GenViT와 HybViT에 대한 그림이다. 기존의 Diffusion model 방식과 비교했을 때, 정말로 Unet 구조를 ViT 구조로 바꾼 것 외에는 크게 달라진 것이 없다. 또한 사용된 ViT 식도 기존 Transformer 식과 거의 달라진 것이 없다.

GenViT 수식과 기존 Transformer의 Encoder

첫번째 줄에서, 이미지에서 패치 단위로 쪼개어 positional embedding까지 하는 것도 기존의 ViT와 동일하며,

3, 4번째 줄에서 Norm 후, MSA(Multi-Head Attention)을 한 것과, Norm 후, MLP를 태운 것도 동일하다. 또한, 중간중간 skip connection을 해준 것과 Transformer를 여러 번(L번) 쌓는 것도 동일하다.

다른 점은 timestep에 대한 정보를 주어야 하기 때문에 time embedding을 하는 A layer가 추가되었으며, 각 패치에 time embedding을 섞어주는 layer가 각 norm layer 전에 들어간다는 것이다. 또한 Diffusion model에서는 똑같은 크기의 이미지가 나와야 하기 때문에 마지막에 reconstruction layer를 추가하여 패치를 이미지 크기로 decoding 하는 부분이 추가되었다.

 

Loss function

Loss 함수에 대한 수식 및 그림

논문에서 기존의 Hybrid model은 학습을 위해 classification과 generation의 성능을 모두 높이기 위해 22번과 같은 식이 나와야 한다고 말하였다. 22 번식의 우항에서 앞에 항은 classification에 대한 확률이고, 뒤의 항은 generation에 대한 확률이다. HybViT에서는 앞의 항은 classification에서 많이 쓰이는 cross entropy loss를, 뒤의 항은 DDPM에서 등장하는 simple loss를 사용한다고 하였다. GenViT와 HybViT의 차이점은 cross entropy loss를 사용하느냐 사용하지 않느냐의 차이이다. 또한, 이 논문에서는 알파 값으로 Loss의 가중치를 결정하는데, 알파 값이 100일 때, classification의 성능도 어느 정도 유지하면서 gneration의 성능이 괜찮게 나온다고 주장하고 있다.

 

Experiments

실험 결과이다. 신기하게도 어느 것 하나 매우 잘났다고 하는 부분이 없다. 매우 정직하다.

CIFAR10에 대한 기존 모델들과의 성능 비교

위의 결과를 보면 classification에 대한 정확도도, IS와 FID score도 어떠한 SOTA 모델보다 나은 게 없다. 그러나 HybViT는 Hybrid model로 classification과 generation을 동시에 하는 것에 중점을 두고 있다. 따라서 개별 task를 해결하는 ViT나 StyleGAN2-ADA에 비해 각각의 성능은 떨어지지만, 종합적으로 준수한 성능을 보이고 있으며, Hybrid model 중에서는 최고의 성능을 내고 있다고 주장하고 있다. 그러나 아쉽게도 각각의 task에 대해서 높은 성능을 내지 못하고 있으며, 아마 이러한 이유 때문에 더 이상 변별력이 없다고 판단하여 Hybrid model에 대한 연구를 그만둔 것이 아닌가 싶다.

 

Conclusion을 살펴보면 이 논문에서는 다양한 실험을 진행하였다고도 주장하고 있다. 그러나 실제로 다양한 실험을 진행하지만 정말로 다양한 실험을 진행할 뿐, HybViT가 낫다거나 유의미한 내용은 없는 것 같다.

 

Code review

코드는 아래 github 페이지에 공유되어 있으며, 공유된 code 위주로 code review를 진행하겠다.

https://github.com/sndnyang/Diffusion_ViT?tab=readme-ov-file

GitHub - sndnyang/Diffusion_ViT: PyTorch Implementation of "Your ViT is Secretly a Hybrid Discriminative-Generative Diffusion Mo

 

우선 Diffusion 부분은 거의 동일하게 작성되어 있으며, denoise 부분이 Unet에서 ViT로 바뀌었다는 점을 기억하면서 보면 이해가 편할 것이다.

 

main.py

우선 main 함수를 살펴보면 처음 model을 create_model이란 함수를 통해 만들고, 이를 GaussianDiffusion class의 인자로 넣어 최종 diffusion_model을 만든다. (create_model이랑 GaussianDiffusion이 뭐지? 라고 하지 말고 바로 아래를 봐라)

GaussianDiffusion class

여기서 GaussianDiffusion이란 기존의 DDPM에 있는 함수로, 첫 인자로 denoise layer에 사용될 model을 인자로 받는다. 기존의 DDPM의 경우 이 부분에 Unet이 들어가겠지만, 이번 코드에서는 이 부분이 create_model로 만든 ViT model이 들어가게 될 것이다.

create_model 함수

바로 create_model 함수를 살펴보면, 다양한 Transformer에 대해 구현되어 있는 것을 볼 수 있는데, difvit라는 부분이 논문에서 설명한 model이라고 할 수 있다. difvit 조건문 마지막에 DifViT class로 ViT 모델을 생성하는 것을 볼 수 있다.

 

DifViT.py

DifViT 코드를 살펴보면 논문에 나와 있는 수식이 그대로 구현되어 있는 것을 볼 수 있다. 수식의 각 부분들에 해당하는 코드를 같은 색으로 하이라이트를 해 놨다.

forward 함수의 경우 인자로 t와 recon, feat를 추가로 받을 수 있게 되어 있는데, t의 경우 diffusion 단계, 즉 sampling이나 학습 중 sample loss를 구할 때 주어지며, classification 단계에서는 쓰이지 않는다. 또한 diffusion 단계에서는 recond=True로 하여 input과 동일한 이미지 크기의 output을 반환하며, classification 단계에서는 feat=True로 하여 cls token만을 취한다.

 

vit.py

vit.py에 있는 Transformer부분을 살펴보면 위와 같다. 이번에도 수식의 각 부분들에 해당하는 코드를 같은 색으로 하이라이트 해 놨다. attn과 ff은 기존 Transformer와 똑같이 작성되어 있다.

 

loss.py

마지막으로 loss function에 대한 코드이다. 앞서 cross entropy loss와 simple loss를 사용하였다고 하였는데, 그 부분에 대한 식이다. 여기서 살펴봐야 할 점은 simple loss에 해당하는 주황색 부분의 경우 model이 그대로 사용되었고, cross entropy loss에 해당하는 파란색 부분에서는 model.denoise_fn이 사용되었다. 즉, classification 단계에서는 diffusion의 전체 과정이 쓰이는 것이 아니라 단순히 한 번의 ViT만으로 loss가 계산된다는 것을 의미한다. 이때, model의 인자로 t가 들어가지 않으며, feat=True로 사용된다.

또한 loss 수식에서 보면 두 loss의 가중치를 알파라는 하나의 미지수로 관리하는 것처럼 보이나, 실제 코드에서는 px과 pyx를 인자로 받아 관리되며, 앞서 Loss Function 단원에서 설명했던 대로 알파값이 100일 때 최적의 성능을 보였다고 되어 있어, main.py 코드에서도 pyx=1, px=100이 default로 설정되어 있는 것을 볼 수 있다.