[논문 리뷰] Vector Quantized Diffusion Model for Text-to-Image Synthesis 논문 리뷰

Link

• 깃허브 링크 : https://github.com/cientgu/VQ-Diffusion
• 논문 링크 : https://arxiv.org/abs/2111.14822

 

 

Abstract

Diffusion + VQ-VAE + Masking

Diffusion으로 벡터 양자화에 의한 오차의 누적을 피하고 이미지 해상도에 따라 선형적으로 비용이 증가하고 선형 회기 방법(AR)

과 달리 엄청 빠른 속도가 나오면서 더 나은 품질을 보여 줍니다.

 

Introduction

DALL-E와 같은 AR 방식의 단점

1. 왼쪽 위에서 오른쪽 아래 텐서 순으로 예측이 되기 때문에 편향이 생길 수가 있습니다.

2. 추론 단계가 이전에 샘플링된 토큰을 기반으로 실행되기 때문에 앞선 토큰의 오차가 계속해서 누적이 됩니다.

따라서 Diffusion 모델을 사용하고 네트워크 수렴을 위하여 mask-and-replace 방식을 채용하였습니다.

 

Background: Learning Discrete Latent Space of images Via VQ-VAE

 

VQ-VAE는 인코더 E, 디코더 G, 코드북 Z(i ∈ {1, 2, ..., K})로 구성됩니다.

 

이미지 x가 주어진다면 공간 feature zij를 가까운 코드북 항목 zk네 매핑하는 quantizer Q(·)와 컬렉션 zq를 얻습니다.

 

식은 x̃ = G(zq)이고 따라서 이미지 합성은 잠재 분포에서 이미지 토큰을 샘플링하는 것과 같습니다. 모델은 다음과 같이 

훈련이 될 수 있습니다.

 

첫째 항부터 순서대로 전체 과정 최적화를 진행하고, 코드북 최적화, 인코더 최적화순으로 진행을 합니다.

코드북 매핑은 역전파가 불가능하기 때문에 인코더 아웃풋과 디코더 인풋을 그냥 복사해서 연결을 하였더니 잘 되었다고 합니다.

(VQ-Diffusion의) 연구진들은 둘째 항을 EMA로 대체하는게 더 작동이 잘 되었다고 합니다.

 

Vector Quantized Diffusion Model

코드북에 매핑된 이미지 토큰 x와 BPE 인코딩을 통해 얻은 텍스트 토큰 y를 통해서 conditional transition distribution

q(x|y)를 최적화 하는 것입니다.

 

AR 방식 모델들은 소개에서 말했듯이 이와 같은 문제가 있어서 masked language modeling (MLM)에서 영감을 받은

Diffusion 과정과 훈련 방법을 소개합니다.

 

Discrete diffusion process

 

이산 확산 과정은 노이지를 추가하는 것이 아니라 예를 들어 토큰 일부를 무작위로 교체하는 방식으로 진행이 됩니다.

또한 역확산을 진행하는 모델은 U-Net이 아니라 트랜스포머를 사용합니다.

 

자세히 들어가면, 이미지 x0의 한 이미지 토큰 x0i에 대해 K x K 마르코프 transition matrix Qt를 활용하여 다음 단계에서

어떤 토큰으로 바뀔지에 대하여 확률을 정의가 가능합니다.

(v(x)는 길이가 K고 해당 토큰의 코드북 인덱스만 1인 one-hot 열벡터. 윗첨자 i는 앞으로 생략.)

또한, 마르코프 체인의 특성 덕분에 중간단계를 배제하게 되고 0에서 t단계에서의 확률을 바로 도출하는 것이 가능합니다.

Qt를 보면 그대로 유지될 확률이 조금 더 높고 나머지 인덱스로 대체될 확률입니다.

 

다만, Diffusion 과정에서 데이터 손상에 균일한 확률을 사용한다는 것을 급격한 의미 변화, 추정 난이도가 급격하게 높아지는

불안정 요소를 의미합니다.

 

Mask-and-replace diffusion strategy

위 문제를 해결하기 위해서 [MASK]토큰을 추가해줍니다.

이미지 토큰은 γt의 확률로 마스크 토큰으로 진행이 되고, 한 번 마스크 토큰이 되어버리면 다음 단계부터는 변하지 않게 됩니다.

Qt는 (K+1)x(K+1) 행렬이 됩니다.

 

Mask-and-replace의 이점

• 손상된 토큰이 식별이 가능하며 역확산에 용의합니다. 다만, 무작위로 대체된 토큰의 경우 해단 토큰이 대체되었는지 아닌지 알 방법은 없습니다.)
• 무작위 토큰 대체는 네트워크가 마스크에만 집중하지 않고 전체 context를 이해하도록 장려합니다.
• 또한 닫힌 형태 계산이 가능하기 때문에 계산 비용도 감소가 됩니다.

 

 

Learning the reverse process

 

Reparameterization trick on discrete stage

xt-1을 직접  얘측하기 보다 다른 대리 변수를 근사화하는 것이 더 좋다.

따라서 모델이 xt-1이 아닌 pθ(x̃0|xt, y)를 추정하도록 하고 xt-1을 구한다.

 

Model architecture

 

 

Fast inference strategy

재개매개변수화로 인해 모든 추룬 단계를 거치지 않고 일정한 보폭으로 단계를 건너뛸 수 있다.

 

Expertiments

사용 파라미터

256x256 이미지를 32X32 토큰으로 (32x32개),

K = 2886,

텍스트 인코드 CLIP,

Timesteps = 100 정도로 작게

최종 단계에서 마스크 비율은 90%