Vector Quantized Diffusion Model for Text-to-Image Synthesis, CVPR 2022 논문 리뷰

이번에 리뷰 할 논문은 VQ-diffusion입니다.

CVPR 2022년도에 발표되었고 기존 DALL-E 같이 기존 VQ-VAE를 기반으로 auto-regressive하게 이미지를 생성하는 모델들의 문제를 tackling하는 논문으로

같은 년도인 2022년에 DALL-E 2가 나와서 묻힌 감이 없지 않아 있지만 그 컨셉에 대해서 한번 쯤 보고 넘어가면 좋을 거 같아 한번 정리를 합니다.

 

간략하게 보실 분들은 2-4. Discrete diffusion 파트와 3. Mask diffusion 파트 그리고 8. Summary 파트 위주로 보시면 될것 같습니다.

 

* 해당 리뷰의 내용은 출처를 남겨주시면 자유롭게 인용하셔도 됩니다.

Index

1. Preview

2. Preliminary

    2-1. VQ-VAE

    2-2. DALL-E

    2-3. Limitations

    2-4. Diffusion Model for Discrete state space

        2-4-1. Markov transition matrix

        2-4-2. Discrete diffusion

        2-4-3. Limitations

3. Mask diffusion (Contribution)

    3-1. Forward process

    3-2. Reverse process

4. Loss function

5.  Architecture

6. Quantitative results

7. Qualitative results

8. Summary

 

1. Preview

Fig1. model architecture

들어가기 앞서 VQ-diffusion의 전체적인 구성을 보면 VQ-VAE를 학습 시키는 step1 pretrain part와

step1에서 학습된 codebook vector로 이루어진 latent space를 만드는 step2 VQ-diffusion part로 이루어져 있습니다.

 

조금 더 구체적으로 말하자면

입력과 출력을 같은 image로 하여 train 된 VQ-VAE를 통해

image를 restore할 수 있을 만큼 잘 나타내는 code book vector들을 학습해 놓고

 

Text encoder를 통해 encoding된 text condition과 attention score가 높은 방향으로 reverse diffusion process를 거쳐

code book vector 집합을 만들고 이를 디코딩 함으로서 우리가 원하는 condition에 맞는 이미지를 생성하는 것이

VQ-diffusion의 기본 구조입니다.

 

2. Preliminary

사실 본 논문 같은 경우 Discrete diffusion의 후속 연구에 가까운 논문이기 때문에 VQ-VAE와 DALL-E, Discrete diffusion에 대해서 설명하게 되면 이 논문에 대해 거의 설명을 마쳤다고 할 수 있을 것입니다.

 

하여, Preliminary 파트에서는 본 논문의 기반이 되는 VQ-VAE에 대한 설명과

VQ-VAE 기반의 Image generation 방식인 DALL-E에 대한 소개 그리고

기존의 Auto-regressive한 방법들이 가지고 있는 문제점들과 그 해결책으로 제시하는 Discrete diffusion에 대해 서술하고

 

마지막으로 3번 Mask diffusion 에서 본 논문에서 제시하는 Discrete diffusion을 개선한 Mask-and-replace diffusion strategy를 서술하는 식으로 리뷰를 진행하도록 하겠습니다.

 

2-1. VQ-VAE

Fig2. architecture of VQ-VAE

Text-to-image generation task에서 특유의 query key 연산을 통해 text와 image 사이의 연관성을 파악할 수 있는 transformer는 매우 좋은 performance를 내고 있습니다.

하지만 transformer의 경우 raw pixel을 직접 embedding해서 사용하기에는 일반적인 dot product attention 연산의 구조상 computational cost가 상당히 들어가기 마련이죠

이런 문제를 해결하기 위한 방법중 하나로 image generation에서 VQ-VAE를 사용하여 이미지를 discrete한 image token들로 나타내 sequence length를 줄이고 이를 transformer를 사용하는 방식이 사용 되고 있고 DALL-E를 비롯한 본 논문에서도 이 방식을 채택하고 있습니다.

 

VQ-VAE의 구조에 대해 간략한 설명을 드리면 기본적인 encoder와 decoder로 이루어진 autoencoder 구조를 사용하고 있고 입력 이미지를 넣어 출력으로 입력이미지를 복원하는 task를 수행합니다.

다만 차이가 있다면 디코딩을 진행할 때 인코딩 결과로 나온 latent space를 그대로 사용하는 것이 아니라

미리 사전에 정해놓은 K개 만큼의 learnable한 codebook vector로 encoder의 latent space를 대체하여 그 결과를 가지고 디코딩을 한다는 것이 차이입니다.

대체 방식의 경우 상기 Fig2 처럼 encoder의 output feature map의 x,y에 해당하는 위치의 vector와 K개의 code book vector사이의 distance를 재서 가장 가까운 index의 code book vector를 찾아로 교체를 하게 되고 이를 0,0부터 h,w까지 수행하여 최종적으로 encoder의 output feature map h,w,d를 전부 code book vector로 교체하게 됩니다.

 

이와 같은 latent space의 교체연산의 경우 미분이 불가능하다는 문제점 때문에 교체 연산을 기준으로 output부터 디코더를 포함하는 code book vector까지 backpropagation을 하고 해당 loss를 복사하여 인코더 파트에 전달해주는 stop gradient라는 기법이 사용되는데 더 궁금하신 분의 경우 자세한 내용은 VQ-VAE 논문을 참고하시면 좋을 것 같습니다.

 

최종적으로 VQ-VAE를 사용하여 입력이미지를 잘 복원할 수 있도록 train에 성공한다는 것은 K개의 discrete한 codebook vector만으로도 이미지를 충분히 잘 표현할 수 있다는것이 될것이고 이미지를 train된 VQ-VAE를 사용하여 code vector의 인덱스로 표현하게 되면 이미지를 표현하는데 드는 sequence의 길이를 획기적으로 줄일 수 있습니다.

 

예를 들면 (h,w,d)의 shape를 가진 인코더의 feature map을 code book vector의 인덱스로 표현시 하단 Fig3 과 같이 (h,w)로 표현할 수 있고 이 quantize 된 feature map을 transformer에 사용함으로서 computational cost를 절약하게 됩니다.

Fig3. Quantized image

2-2. DALL-E

Fig4. DALL-E

기존의 distance를 기준으로 인코더의 feature map을 codebook vector로 교체하는 VQ-VAE와는 다르게

DALL-E 같은 경우 GPT 처럼 Decoder only transformer 구조를 사용하여

입력 문장 토큰들을 바탕으로 auto-regressive하게 latent space의 0,0 부터 h,w까지에 들어갈 적절한 code book vector의 index를 예측을 하고 이렇게 만들어진 latent space를 디코딩하는 방식으로 학습이 이루어집니다.

 

Fig4를 기준으로 설명을 하면 an armchair in the shape of an avocado라는 문장이 들어갔을 때

latent space의 0,0 위치에 0번부터 K-1번까지 codebook vector의 index 중에 어떤 codebook vector의 index가 나오는 것이 적절한지 확률을 예측하여 이중 가장 확률이 높은 codebook vector의 index를 선택하게 되고 예시에서는 12번 index의 codebook vector가 선택된 상황입니다. 이후 다음 위치인 0,1 위치에 올 codebook vector를 예측할 때는 앞선 문장과 함께 12번 인덱스에 대한 정보까지 넣어 auto-regressive하게 예측한 결과 21번 codebook vector가 선택 되었고 0,2 위치는 7번 0,3위치에는 2번 codebook vector가 선택 된 것입니다. 해당 과정을  좌상단(0,0)부터 우하단(h,w)까지 반복하면 적절한 codebook vector로 이루어진 latent space를 얻을 수 있는데 이를 디코딩하여 이미지를 생성하게 됩니다.

 

 

2-3. Limitations

Image generation task 에서 이러한 auto-regressive 방식은 몇가지 문제점을 가지고 있고 본 논문에서는 이를 지적하고 있습니다.

Fig5. Limitations of auto-regressive method

첫번째로 이미지의 경우 상하좌우 공간적인 특성이 중요한데 앞서 설명한 auto-regressive하게 좌상단 부터 우하단까지 codebook vector의 인덱스를 예측 하게 되는 경우 이러한 점이 무시되고 이전 위치까지의 정보만을 가지고 예측을 해야한다는 점이 있습니다.

두번째로는 일단 한번 codebook vector를 예측하는데 error(실수)가 발생하게 된다면 그 error가 계속 다음 codebook vector를 예측하는데 사용이 되고 이는 error의 축적을 야기한다는 것입니다.(이 부분에 있어 사실 개인적으로 저는 그렇게 크게 문제라고 생각하진 않습니다. error가 누적이야 되겠지만 transformer 구조상 연속해서 크게 잘못된 예측을 하지 않는 이상 attention mechanism의 softmax를 취하는 과정중에 outlier가 어느정도 무시가 될테니까요)

 

본 논문에서는 discrete diffusion model을 사용하여 동시에 latent space의 모든 좌표에 해당하는 codebook vector index를 생성하여 이와 같은 문제를 해결하고자 합니다.

 

2-4. Diffusion Model for Discrete state space

Fig6. Normal diffusion process

여기서 의문점이 하나가 들 수 있습니다. 일반적으로 저희가 아는 diffusion model같은 경우는

보통 원본 이미지 x_0부터 시작하여 각 스탭별로 gaussian distribution을 따른다고 가정한 noise를 더해나가 gaussian noise 그자체인 x_T를 만드는 diffusion process와

gaussian noise x_T로부터 시작해 이전 T-1스탭에서 추가되었을 guassian noise를 예측하고 이를 x_T로부터 빼주어 x_(T-1) 스탭의 이미지를 만들고 이를 반복해서 x_0 상태의 이미지를 생성하는 reverse diffusion process로 구성되어 있습니다.

 

문제는 continuous한 일반적인 이미지와는 다르게 VQ-VAE나 DALL-E와 같이 latent space를 discrete한 codebook vector의 index를 사용하여 표현하는 경우 이를 적용하기가 어렵다는데에 있습니다.

 

따라서 본 논문에서 저자들은 discrete state space에서도 적용가능한 diffusion model을 사용하여 이를 극복하고자 하였습니다.

2-4-1. Markov transition matrix

Fig7. markov transition matrix

discrete state space상에서 diffusion을 적용하기 위한 방식을 알아보기 위해 우선 Markov transition matrix에 대해 알아보도록 하겠습니다.

 

Markov transition matrix Q_t는 각 상태로 전이할 확률을 모아놓은 matrix로 각 여기서 상태란 codebook vector를 의미합니다. row가 의미하는 것은 시작상태 column이 의미하는 것은 전이될 상태로 각 대각행렬이 의미하는 바는 자기 자신으로 전이될 확률입니다.

 

예를 들면 K=3(codebook vector의 개수)의 code book category로 이루어진 code book이 있다고 할때 (x, y) 좌표계를 기준으로 Q_t의 (0,0)은 0번 index(category)의 codebook vector가 0번으로 전이할 확률을 의미하고 (0,1)은 0번 codebook vector가 1번으로 전이할 확률, (0,2)은 0번 codebook vector가 2번으로 전이할 확률을 의미합니다.

 

discrete state space상에서 diffusion은 이와같은 markov transition matrix와 categorical 변수로 바꾼 현재 codebook vector의 index간의 dot product를 통해 이루어지게 됩니다.

 

FIg7을 기준으로 1번 index(0~k-1 까지 인덱스 기준)의 codebook vector를 one hot encoding을 해주어 categorical하게 바꾼 뒤 markov transition matrix Q_t를 matrix multiplication 해주는 경우 0번 인덱스로 전이할 확률 0.1, 1번 인덱스로 전이할 확률 0.8, 2번 인덱스로 전이할 확률 0.1이라는 결과가 나오게 됩니다. 하단의 예를 통해 조금 더 자세히 설명하도록 하죠

 

2-4-2. Discrete diffusion

Fig8. Example of the discrete state diffusion process

Fig8을 기준으로 좌측 (2,2)의 codebook vector index로 이루어진 latent space가 있다고 가정하면 우선 이를 one-hot encoding을 해줍니다. 그리고 앞서 설명한 markov transition matrix와 dot product를 해주면 다음 step에서 다른 codebook vector로 전이할 확률이 나오게 되고 이를 적용하여 다시 (2,2)의 latent space를 만들게 됩니다. 이런식으로 discrete diffusion을 사용하게 되면 앞서 설명한 바와 같이 이전 step의 latent space 전체를 고려하여 dot product 한번에 다음 step의 latent space를 구할 수 있게 되어 2-3파트에서 서술한 auto-regressive한 transformer의  첫번째 단점을 해결할 수 있다는 것이 저자의 주장입니다.

 

 

Constraints

Fig9. Constraints of markov transition matrix for discrete diffusion process

물론 이러한 markov transition matrix는 아무거나 사용할 수 있진 않고 몇가지 조건이 있습니다.

당연한 얘기지만 markov transition matrix Q_t의 row는 다른 index의 codebook vector로 전이할 확률들임으로 row를 합하였을 때 1이어야 합니다.

 

또한 일반적인 diffusion에서 forward process의 결과가 gaussian distribution에 수렴하고 그렇기 때문에 reverse process를 진행할 때 gaussian distribution에서 시작할 수 있듯

discrete diffusion에서 markov transition matrix도 Q_1부터 Q_t까지 곱하였을 때 우리가 알고 있는 stationary distribution에 수렴해야 합니다.

 

이 조건을 만족하는 matrix로는 uniform distribution을 갖는 row, column 합이 1인 doubly stochastic matrix가 있고 discrete diffusion process에선 이를 사용하게 됩니다.

 

Fig10. conversion of markov transition matrix (uniform distribution, doubly stochastic)

실제로 Fig 10의 좌측에 해당하는 matrix Q를 만들어 t를 일정 크기 이상 줘보시면 matrix의 각 element들이 uniform하게 수렴함을 확인할 수 있습니다. 하단은 해당 코드입니다.

# https://cv-ai.tistory.com
import numpy as np
from matplotlib import pyplot as plt
import seaborn as sns
def make_Q_t(K=4,beta_t =0.5):
    Q_t = np.zeros([K,K])
    Q_i,Q_j = Q_t.shape

    for i in range(Q_i):
        for j in range(Q_j):
            if i == j:
                Q_t[i,j] = 1-((K-1)/K)*beta_t
            elif i !=j:
                Q_t[i,j] =(1/K)*beta_t
    return Q_t

t = 15
Q_1 = make_Q_t()

print('#'*10+'Q_1'+'#'*10)
print(Q_1)
Q_t = Q_1
for i in range(t):
    Q_t_next = Q_t_next.dot(make_Q_t())
print('#' * 10 + 'Q_t' + '#' * 10)
print(Q_t_next)

fig, ax = plt.subplots(1,1)
img = sns.heatmap(Q_t,annot=True)
plt.show()
plt.cla()
img2 = sns.heatmap(Q_t_next,annot=True)
plt.show()
plt.cla()

 

2-4-3. Limitations 

지금까지의 나온 내용에 따르면 discrete diffusion을 사용하게 되면 2-3 파트에서 저자가 제기한 auto-regressive한 transformer구조를 사용했을 때 벌어지는 문제를 해결할 수 있지만 여전히 몇가지 문제가 남아있습니다.

Fig11. Limitations of the discrete diffusion

첫번째로 forward process중에 운이 없게 image token이 연관성이 없는 다른 token으로 한번 대체된 경우 토큰이 가진 의미를 잃어버리게 될 수 있다는 문제점과

두번째로 reverse process중에 모델이 어느 부분을 fix하고 어떤 부분을 바꿔야할지 파악하는데 많은 resource가 든다는 문제점입니다.

 

두번째 문제에 대해 추가적인 설명을 해보면 Fig11의 하단의 예시 그림을 보시면 현재 I want to go school이라는 단어가 있고 이를 reverse process를 통해 original sentence인 I want to go home으로 돌려야 한다고 생각해 봅시다.

school만을 home으로 바꾸면 되는 과정이지만 model의 입장에서는 어느 부분을 fix하고 어느 부분을 바꿔야할지 train을 통해서 파악해야합니다.

즉 school이라는 단어를 home으로 바꾸는 것을 학습하는 것 보다 어떤 부분을 fix하고 바꿔야할지 학습하는데 더 많은 리소스가 들게 되고 이는 상당히 비효율적이다는 것이 해당 논문의 주장입니다.

 

결론적으로 이 두가지 문제점을 해결하기 위한 방법으로 본 논문의 저자는 Mask-and-replace diffusion strategy를 제시하게 되는데 컨셉은 상당히 간단합니다. 말그대로 forward process를 통해 일정 확률로 token (여기서는 code vector)를 masking하여 discrete diffusion을 진행하는 것으로 masking을 제외한 기본적인 룰은 앞서 설명한 discrete diffusion을 따릅니다.

 

이런 방식의 masking을 통해서 저자는 masking을 통해 token이 다른 의미를 가진 token으로 바뀌게 되는 문제가 발생하지 않고 masking을 통해 어느 부분이 바뀌었는지를 모델에 쉽게 알려줌으로써 상기 문제들을 해결했다고 주장을 합니다.

 

3. Mask diffusion

해당 파트에서는 Mask diffusion에 대해 forward와 reverse process위주로 간단하게 설명을 하도록 하겠습니다. 사실 앞서 discrete diffusion에서 대부분의 내용을 다 언급하여 몇줄 되지는 않을 것 같네요 ㅎㅎ

3-1. forward process

Fig12. Forward process of mask diffusion

Mask diffusion의 forward process의 경우 앞서 설명한 discrete diffusion matrix에서 단지 감마 값 즉 masking될 확률로 row 한줄이 추가가 됩니다. 그 외에는 앞에 파트에서 서술한 일반적인 discrete diffusion process와 동일한 과정을 거칩니다. Fig12를 참조하시면 쉽게 이해가 되실 겁니다.

3-2. reverse process

Fig13. Reverse process of mask diffusion

Mask diffusion의 reverse process또한 discrete diffusion process와 동일하게 이루어지는데 한가지 특이한 것이 있다면 Fig.13에 서술한 것과 같이 x_t에서 x_t-1 스탭을 바로 예측 하는 것이 아닌 x_t에서 x_0를 예측하고 거기서 x_t1 만드는 방식으로 reverse process가 이루어지게 됩니다.

4. Loss function

Fig14. loss function of VQ-diffusion

loss function또한 간단합니다. 일반적인 diffusion process에 사용되는 loss에 Auxiliary loss term이 추가 된 형태입니다.

각각의 loss term에 대해 설명을 간단하게 하고 넘어가도록 하겠습니다.

Fig.15 Reconstruction term

x_0에 대한 reconstruction을 위한 reconstruction term입니다.

Fig16. Denoising process term

Reverse process를 위한 term입니다. 부가 설명을 조금 하자면

우리가 diffusion model이나 어떤 다른 모델을 사용하지 않고 단순 계산으로 x_t가 주어졌을 때 x_t-1을 아는것은 불가능 합니다.  그러나 x_0가 조건으로 주어진다면 일반적인 diffusion process를 기준으로 gaussian noise를 더해서 x_t-1을 만들 수 있습니다.

하지만 x_0를 생성하기 위해 diffusion model을 사용하는 건데 x_0를 조건으로 준다는거는 사실 주객전도에 해당하는 일일 겁니다.

그래서 이를 neural network를 활용하여 해결하는 텀이 해당 텀입니다.

p theta 즉 우리의 neural network가 x_t와 y (text guide)를 받아서 x_t-1을 예측하도록하고 실제 그 둘사이의 KL divergence를 줄이는 방향으로 학습을 하게 해주는 텀이 해당 텀입니다.

 

Fig17.  Regularization term

diffusion process가 성립하려면 우리가 가정한 분포(gaussian, uniform등) 가 성립해야 하므로 network의 output인 p(x_t)가 실제 x_0가 주어졌을 때 x_t의 분포와 같게 맞추어주는 term입니다.

Fig18. Auxiliary loss term

앞선 loss가 일반적인 diffusion model의 loss라면 해당 term은 3-2파트의 mask diffusion reverse process를 support하기 위한 loss term입니다. reverse process에서 x_t를 받아서 x_0를 예측하는걸 좀더 잘하게 해주기 위한 역할을 합니다.

 

5. Architecture

Fig19. Architecture of VQ-diffusion

전반적인 architecture는 다음과 같습니다 code book vector를 위한 VQ-VAE를 pretrain하고 pretrain된 text encoder를 통해 인코딩된 text를 반영하여 codebook index에 mask diffusion을 적용시켜 줍니다.

 

6. Quantitative results

Fig20. Quantitative results

Quantitative results는 다음과 같습니다. 성능 평가는 inception v3 기반 FID score를 통해 진행하였습니다.

너무 당연한 말이지만 SOTA를 찍었습니다.

 

7. Qualitative results

Fig21-1. Qualitative results

Fig21-2. Qualitative results

Qualitative results는 다음과 같은데요 비교적 쉬운 CUB-200의 경우 그럭저럭 괜찮은 거 같은 결과를 보였지만 MSCOCO데이터 셋에는 아직 어색한점이 좀 보입니다. 그 외에도 inpainting, text-to-image synthesis, image editing 등 다양한 task에 적용한 결과를 확인할 수 있습니다.

 

번외

Fig22. Throughput

추가적으로 한가지 눈여겨 볼 점은 autogregressive한 transformer base의 method의 경우 n * m의 token을 생성할 때

n * m time의 step이 필요했다면 discrete diffusion base의 해당 논문의 경우 한번에 reverse process시에 n * m의 token을 전부 생성하게 되어 throughput면에서 이득을 봤다고 합니다.

물론 diffusion model의 특성상 reverse process의 step수가 많아질수록 당연히 throughput이 안좋아지긴 해도 어느정도 FID score를 보장하는 선에서 조금 더 빠른 성능을 냈다는거는 나름 재밌는 점 같습니다.

차후에 reverse process의 step수를 더 줄일 수 있는 논문들이 많이 나온다면 더 빠른 throughput을 낼 수 있지 않을까 기대가 되는 점입니다.

 

 

7. Summary

Fig23. Summary

지금 까지 나온 본논문의 VQ-diffusion의 흐름은 다음과 같습니다.

 

Text-to-Image generation task에서 transformer를 사용하는 것이 성능이 잘 나온다

-> 하지만 그냥 image를 그대로 transformer에 적용하기에는 많은 computational resource가 필요하다

 

-> 하여 기존의 DALL-E와 같은 논문에서는 VQ-VAE와 같은 방식을 통해 latent space를 discrete한 codebook vector의 index로 표현함으로써 이를 해결하였다

 

-> 그런데 DALL-E와 같은 auto-regressive하게 latent space를 생성하는 경우 image generation task에서 image생성을 위한 code vector를 고르는 과정중에 좌상단에 해당하는 값들만 참고할 수 있다는 문제점과 한번 잘못 예측된 에러를 고칠 수 없어 계속 축적된다는 문제점이 있다.

 

-> 이를 해결하기 위해 auto-regressive한 방법이 아닌 discrete space상에서의 diffusion을 적용하였다.

 

-> 이런 discrete space상의 diffusion도 reverse process중에 어느 부분이 forward diffusion process로 바뀐 부분인지 모델이 알기 힘들다는 문제가 있고 한번 잘못된 category의 codebook vector로 변경된 경우 전체적인 semantic가 바뀔 수 있다는 점이 있었다.

 

-> 이를 해결하기 위해 mask-and-replace diffusion strtategy를 discrete diffusion에 적용하여 일정확률로 token을 마스킹함으로써 상기 문제점들을 해결하였다.

 

-> 또한 일반적인 diffusion loss에 Auxiliary loss term을 추가하였는데 이는 x_t step으로부터 x_(t-1) 스탭을 예측할 때 그 사이의 noise를 예측하는 것이 아닌 x_0를 예측하고 x_0로부터 x_(t-1)까지 가는데 필요한 noise를 예측하는 것을 surpport해주기 위함이다.

 

-> 결과적으로 Text-to-Image generation task에서 강점을 보이는 transformer 구조를 사용하여 image generation을 진행하되 mask diffusion을 통해 auto-regressive한 모델의 단점을 극복하였다.

 

VQ-diffusion = VQ-VAE + Discrete diffusion + Mask-and-replace diffusion strategy

 

 

 

 

논문을 읽은건 2022년 10월쯤이였는데 쓰다보니 길어지게 되어 리뷰 마무리를 이제와서 하게 됐네요.

추가로 중간에 설명이 생략이 된 부분이 있을 수 있으니 덧글로 질문 주시면 여유가 되는 내에서 답변해드리도록 하겠습니다.

 

그 외에 diffusion model 및 computer vision 관련 논의도 환영합니다!

 

 

 

################

2023.3.17 수정