Taming Transformers for High Resolution Image Synthesis (VQGAN)

https://arxiv.org/abs/2012.09841

Taming Transformers for High-Resolution Image Synthesis

 

## Background

 

본 논문에 들어가기 앞서 Background 지식을 조금 알아보겠다.

### Inductive Bias

Inductvie Bias란 머신러닝의 학습 시 직면하지 않았던 상황에 대해 정확한 예측을 하기 위해 사용하는 추가적인 가정을 의미한다. 머신러닝에서는 Target을 예측하기 위해 학습할 수 있는 모델 구축을 목표로 하고 있는데 이를 이루기 위해 제한된 훈련 데이터가 주어진다. 여기서 제한된 훈련 데이터를 예측하는 것 뿐만 아닌 일반화를 할 수 있게 만든느 것이 Inductive Bias다.

Inductive Bias의 가정은 관계성의 유무에 따라 나눠진다. 여기선 관계성이 있는 경우 즉, Relational Inductive Bias에 대해 얘기하겠다. 그리고 이를 FCN과 CNN에 따라서 설명한다. Relational Inductive Bias란 입력값이 출력값에 미치는 관계를 말한다.

FCN, CNN

FCN의 Inductive Bias는 Weak하다. FCN은 weight 들을 모두 공유하기 때문에 이는 모든 입력의 요소가 어떤 출력 요소던지 영향을 미칠 수 있다. 따라서 이는 Weak한 Relational Inductive Bias임을 의미한다.

반면에 CNN은 일반적으로 입력값의 이미지 크기보다 작은 하나의 filter를 통해 sliding window 방식으로 convolution 연산을 한다. 필터가 이미지의 일부를 대상으로 연산을 수행하기 때문에 이는 CNN의 entity 간의 관계가 서로 가까운 값들 간에 존재한다 말할 수 있고, 이는 어떠한 특성을 가지는 값들이 모여 있는지에 따라 중요하게 판단할 수 있는 구조가 된다.

따라서 Entity 간의 관계가 서로 거리가 가까운 요소들이 존재한다는 Locality의 특성을 가지게 된다. 그리고 CNN은 Paramter Sharing의 특성에 따라 Fully Connected 보다 훨씬 적은 수의 파라미터를 활용하여 지역적인 영역에서 공간적인 특성을 잘 뽑을 수 있다. 이는 효율적인 학습 방법이라고 할 수 있다.

이러한 Inductive Bias는Relational inductive Biases, deep learning, and graph networks 논문에 자세하게 설명되어있다.

 

### Transformer

2022.07.12 - [논문 정리] - Attention is all you need

Attention is all you need

 

자연어에 대한 Transformer는 위 게시물 참고

 

CNN은 Locality의 특성을 강하게 갖는다. 그러면 이미지의 특성을 Global하게 뽑을 수 있는 모델은 무엇일까. 바로 Transformer이다. 

Transformer는 본래 NLP task에서 Machine Translation에 많이 쓰인 모델이었다. Vit를 시작으로 vision task에서도 많이 사용되고 있다.

Vit 에서는 이미지가 patch 단위로 Transformer Encoder에 들어간다. 여기서 들어간 patch 단위의 이미지들은 Transformer Encoder의 Self-Attention과 그 전에 활용되는 positional encoding에 의해서 이미지의 global한 특성을 잘 뽑을 수 있게 된다.

Vit

또한 Transformer는 Auto-Regressive 하게 예측이 가능하다. 이는 다음 단어를 예측하는 것과 같이 Embedding 된 input sequence 값을 Decoder에 넣어서 auto-regressive 하게 target sequence를 예측하는 것과 같다.

하지만 이러한 transformer는 단점이 존재한다. CNN과 다르게 transformer는 inductive bias를 학습하지 않아 local structure를 보는 것이 떨어진다. 또한, 모든 이미지에 대한 관계를 학습 하므로 많은 computation이 소요된다.

## 중간 정리 (CNN vs Transformer)

CNN과 Transformer의 장단점을 비교해서 적어보자

모델	장점	단점
CNN	Strong local Correlation	Image의 Global한 이해 부족
	학습의 효율성	표현력의 한계
Transformer	Global Correlation 파악 용이	많은 Computation
	Weak Inductive Bias ( 풍부한 표현력)	High Resolution 생성에 관한 어려움

CNN은 inductive Bias의 영향으로 강한 local correlation을 갖고 있어 locality 파악에 용이하다. 또한, convolution filter를 이용하여 학습을 하기 때문에 parameter sharing을 하게 된다. 이는 적은 paramter로 연산을 비교적 적게하여 지역적인 영역에서 공간적인 정보를 뽑아내기 때문에 효율적인 학습 방법이라는 장점이 있다.

그러나 이러한 장점으로 인한 단점이 생기는데 local structure를 중심적으로 파악하기 때문에 이미지의 global한 이해가 떨어진다. 또한, locality를 강조하여 학습을 하기 때문에 표현력 측면에서도 한계가 있고, 특히 high resolution 이미지를 표현하는 것에 대한 한계가 존재한다.

Transformer는 이미지의 전체를 self-attention으로 학습을 하기 때문에 이미지의 global한 관계를 잘 파악할 수 있다. 또한 이러한 특성과 Inductive Bias의 관계가 약하기 때문에 풍부한 표현력을 가지고 있다.

그러나 이미지 전체를 학습하기 때문에 많은 computation이 소요되고, 이렇기 때문에 Encoder와 Decoder 를 모두 사용하는 high resolution 이미지를 생성할 때에 어려움이 있다.

따라서 VQGAN은 이러한 CNN과 Transformer의 두 장점을 살리면 더욱 powerful한 모델이 될 수 있지 않을까라는 질문으로 논문을 시작하게 되고 각각의 모델의 단점을 극복하기 위한 방법을 제시한다.

 

## VAE (Variational Auto Encoder)

VQGAN의 기반이 되는 VQ-VAE에 대해 설명하기에 앞서 간단하게 VAE에 대해 설명한다. VAE란 autoencoder 기반의 구조를 활용한 데이터 생성 모델이다. AutoEncoder는 Image를 Encoding 하여 Embedding Vector 즉 latent vector를 만든 후 이를 Decoder의 input으로 넣어 이미지를 reconstruction 한다. 

 

Auto Encoder

VAE는 이러한 AutoEncoder 방식을 사용하여 input 데이터의 분포 자체를 학습하여 데이터를 생성한다. 

VAE

먼저 Encoder 부분에선 input의 정규분포를 따르는 latent variable을 생성하도록 학습한 뒤 이 분포의 \( \mu \) 와 \( \sigma \)를 학습한다. 그리고 gaussian noise와 latent variable의 \( \sigma \)를 먼저 곱하고 이 후 이 분포에 \( \mu \)를 곱하여서 decoder의 input으로 전달한다. 이러한 과정을 통해 이는 backpropogation이 가능하게 되고 이를 Reparameteriztion Trick 이라고 한다.

## VQ - VAE (Vector Quantization Variational Auto Encoder)

본 논문의 VQGAN과 가장 닮아있는 VQ-VAE에 대해 설명하겠다.

 

VQ-VAE

VQ-VAE는 기존 VAE의 문제점인 Decoder가 너무 강해 latent vector들이 무시되는 현상이 발생하는 Posterior Collapse 문제를 해결한 논문이다. 
VQ-VAE는 이러한 Posterior Collapse 문제를 해결하기 위해 Vector Quantization 개념을 도입하였다.
Vector Quantization은 간단히 말하자면 하나의 딕셔너리 형태로 각 카테고리를 매핑하는 과정이다.
그리고 이 방법을 사용하기 위해 codebook 개념을 도입하였다.
codebook은 embedding space에 있는 embedding vector들의 집합으로 NLP에서의 lookup table과 유사하다. 
따라서 Vector Quantization 방법을 사용한 VQ-VAE는 다음과 같이 동작한다.


먼저, input image를 encoder에 통과시켜 \( z_{e}(x) \)를 생성하고 이를 codebook의 각각의 원소와 비교한다. 그리고 \( z_{e}(x) \)의 원소와 codebook의 vector 들 간의 각각의 거리를 구한다. 그리고 이렇게 구한 거리 중 가장 가까운 값의 index를 반환하고 이 index와 mapping 되는 codebook의 vector들로 quantized vector \( z_{q}(x) \)를 구성한다. 이  \( z_{q}(x) \)  Decoder에 넣어 이미지를 생성한다.
그러면 이러한 VQ-VAE는 codebook을 어떻게 update 시킬까?
먼저 codebook을 update하기에는 한가지 어려움이 있다. 바로 argmin 연산은 미분이 불가능한 함수인 것이다. 
즉, codebook과 encoding vector 간의 최소 거리 인덱스를 반환하는 것은 argmin 연산이므로 이는 backpropogation이 안된다는 것이다.
따라서 이를 해결하기 위해 Gradient를 copy하는 방법을 이용한다. 즉, decoder에서 나온 gradient를 바로 encoder에 넣어주는 방식이다.

이를 적용한 VQ-VAE의 loss 식은 다음과 같다.

먼저 Reconstruction loss는 실제 이미지와 생성된 이미지 간의 차이를 구하는 loss로 생성된 이미지가 실제 이미지와의 차이가 없도록 만드는 것을 목적으로 한다.

다음으로 VQ loss는 codebook만 update 하는 loss로 codebook vector가 encoder의 출력과 비슷하게 만들도록 하는 목적을 가진다. 여기서 \( sg \)는 stop gradient 라는 표기로 encoder \( z_{e}(x) \)를 update하지 않는다. 

마지막으로 commitment loss이다. commitment loss 는 Encoder만 update 하는 loss로 Encoder의 출력이 codebook vector와 가까운 값을 출력하는 것이 목적인 loss이다.

## VQGAN

지금까지 설명한 Background를 바탕으로 VQGAN에 설명하겠다.

VQGAN의 Architecture는 다음과 같다.

VQGAN

VQGAN은 VQ-VAE의 기본구조를 따르며 앞서 설명한 CNN과 Transformer의 장점을 결합한 모델이다.

VQGAN은 CNN으로 Locality 를 잘 반영하는 codebook을 학습하고, Transformer의 풍부한 표현력으로 Image Synthesis를 이룬다.

그리고 VQGAN은 2-stage 모델로 첫번째 stage 에선 codebook을 학습하여 Transformer에 사용하기 위한 이미지의 구성요소를 학습하는 것이고, 2번째 stage에선 이러한 codebook을 바탕으로 구성된 구성요소들로 transformer를 이용하여 이미지를 구성하는 것이다.

 

## 1-stage

먼저 첫번째 stage는 codebook 학습하는 단계로 이미지 구성 요소를 학습하는 과정이다.

VQGAN 1-stage

 

VQGAN의 첫번째 stage는 VQ-VAE와 매우 유사하다. 그러나 VAE와 달리 adversarial learning을 사용하여 학습을 진행한다. 

VQ-VAE 처럼 Encoder에서 나온 vector 값과 codebook 간의 유클리디안 distance를 비교한 후 distance가 가장 작은 vector 들의 값으로 quantized vector \(z_q\)를 구성한다.

이렇게 구성한 \( z_q \)를 decoder에 넣어 reconstruction image를 생성합니다. 그리고 이를 discriminator에 넣어 patch 단위로 real 인지 fake 인지 판단한다.

다음에 loss에 대해 살펴보겠다. 이에 대해 설명하기 전에 notation을 정의하면 다음과 같다.

notation

VQ loss

먼저 VQ loss에 대해 살펴보겠다. 이는 그 전 VQ-VAE와 거의 비슷한 것을 알 수 있다.

VQ loss에서 달라진 것은 reconstruciont loss이다. 이는 단순 L2 loss 가 아닌 perceptual loss로 VGG 16 기반으로 real image와 fake image를 넣어서 모델 중간중간의 feature map 간의 loss를 구해서 update하는 것이다. 

그리고 나머지는 VQ-VAE의 loss와 동일하다.

여기서 hyperparameter \( \beta \)는 commitment loss의 중요도를 결정하는 역할을 한다. 만약 \( \beta \)가 1보다 작으면 encoder가 codebook 보다 더 빠르게 update 된다는 의미이다. VQ-VAE에선 codebook이 천천히 update 되고 encoder sample을 더 빠르게 update 하기를 원하여서 \( \beta = 0.25 \)로 설정하였다.

GAN loss

다음으로 Discriminator를 학습하는 loss는 Vanilla GAN loss와 동일하다. 

그래서 이  \( L_{VQ} \)와 \( L_{GAN} \) 을 합쳐 최종적인 loss를 설정한다.

최종적인 loss

여기서 \( \lambda\)는adaptive weight로 두 loss 간의 균형을 맞추는 값입니다. \( \lambda \)에 대한 식은 단순한 hyperparameter가 아니라 다음과 같이 정의된다.

adaptive weight

여기서 \( G_{L} \)은 decoder의 마지막 layer를 말하고 \( \sigma = 10^{-6}\)이다. 즉 GAN loss의 decoder의 마지막 layer의 loss 값으로 decoder의 마지막 layer의 reconstruction loss를 나눠주는 것이다.

 

이러한 과정을 통해서 CNN Encoder Decoder를 이용하여 효율적인 학습을 이룰 수 있게 되고 이미지 구성 요소를 담은 codebook을 생성할 수 있게된다.

 

## 2-stage

다음으로는 2번째 stage에 대해 설명한다.

Transformer는 앞서 언급했듯이 풍부한 표현력을 갖고 있다. 따라서 본 논문에서는 이러한 특성을 이용하여 Transformer를 Image Synthesis 하는데 사용한다.

Transformer는 codebook의 index를 전에 나왔던 값을 기반으로 auto-regressive 하게 예측한다.

2-stage

위 그림처럼 i 번째 보다 작은 값들을 이용하여 i 번째의 index를 예측하는 것이다. 따라서 이렇게 예측한 i 번째 index 값을 codebook 과 mapping 하여 생성하는 이미지 다음 patch vector를 구성하는 방식이다. 여기서 Transformer는 첫번째 stage 에서 학습한 \( z_{q} \)를 label 값으로 삼아 NLL 로 학습을 한다.

이러한 Transformer는 unconditional한 상태 뿐만 아니라 conditional한 상태에서도 동일하게 진행이 된다.

condition

여기서 condition은 depth map이나 semantic segmentation map, keypoint , image class 등이 될 수 있다. 

이러한 condition은 input 과 마찬가지로 1-stage에서 학습을 하여 각각의 codebook Z와 \( Z_c \)를 생성한다. 그 후 이를 바탕으로  codebook \( Z_c \)를 조건으로 주어 다음의 patch 를 예측하게 된다.

 

## Sliding Attention Window

위와 같은 방식으로 이미지를 구성하는데 그러면 어떻게 computation 을 줄여 high resolution 이미지를 생성하는 것인지 알아보자.

이는 본 논문에선 Sliding Attention Window라는 기법을 이용한다.

Sliding attention window

아이디어는 간단하다.

위 그림처럼 검은색 크기의 임의의 window size를 설정한 후 window size 안에 있는 인접한 patch 들과만 attention을 진행해주는 것이다.

이로 인해 모든 관계를 학습 하는 것이 아닌 주변에 있는 patch 들만을 이용하여 attention을 진행하여 computation을 줄일  수 있다.

하지만 이러한 방법은 landscape 이미지에선 잘 동작이 되지만 공간적인 정보가 많이 변하는 이미지에 대해선 잘 동작이 안되는 한계점이 있다.

## Experiments

실험 파트에서는 3가지의 결과를 볼 것이다. 첫째로, Transformer의 중요성을 나타내는 실험, 두번째로 codebook의 중요성을 나타내는 실험, 마지막으로 질적인 결과에 대한 실험으로 마무리 한다.

## Transformer vs Pixel Snail

  

 

 

이 실험은 Transformer 자리 즉, 두 번째 stage를 PixelCNN 계열의 Pixel snail로 대체하여 Transformer와 Pixel Snail의 차이를 보여주는 실험으로 Transformer의 중요성을 보여준다. 각각 동일한 데이터셋에 대해 실험하였고, 동일한 step과 동일한 time으로 각각 학습하였을 때, NLL을 비교하였다.

이 표를 해석하기 위해 예를 들어서 설명한다.

t 라는 정해진 시간 동안 Transformer는 k step을 학습하고, Pixel Snail은 2k step을 학습한다.

이 가정을 통해 설명하면 초록색 box는 t라는 시간에 2k step을 학습한 Pixel -Snail NLL, 파란색 box는 t 라는 시간에 k step을 학습한 Transformer NLL, 빨간색 box 는 2t 라는 시간에 2k step을 학습한 Transformer의 NLL

따라서 같은 step, 같은 time으로 훈련을 하였을 때를 비교하면 Transformer가 더욱 성능이 좋은 것을 알 수 있다.

 

 

 

## Effect of Downsampling Factor

 

 

이는 codebook의 중요성을 downsampling factor라는 것을 통해 나타내는 실험이다.

Downsampling factor는 \( f = H/h \)로 이미지 사이즈를 codebook 사이즈로 나눠준 값이다.

이러한 downsampling factor에 대한 변화로 codebook의 사이즈를 변화시킬 수 있고, 이에 대한 차이를 보여준다.

만약에 f 가 커지면 codebook의 사이즈는 작아지고 이는 codebook에 있는 k 번째 code가 담당하는 영역이 작아지는 것이다. 이는 즉, receptvie field가 작아지는 것이다. 

만약에 f 가 작아지면 codebook의 사이즈는 커지고 이는 codebook에 있는 k 번째 code가 담당하는 영역이 커지는 것이다. 이는 즉, receptvie field가 커지는 것이다. 

따라서 이에 따른 결과를 분석하면 다음과 같다.

Effect of Downsampling factor

왼쪽의 표는 downsampling factor의 변화에 따른 값을 나타낸다. y 축은 NLL, x 축은 reconstruction error 이다. 

따라서 위 결과와 같이 downsampling factor가 작아지면 reconstruction는 잘하고, image generation은 잘 못한다.

반대로 downsampling factor가 커지면 reconstruction은 잘 못하고, image generation은 잘한다.

따라서 이런 trade-off 관계를 고려하여 본 논문에서는 downsampling factor f를 16으로 설정하였다.

 

## Qualitative Result

1. condition을 주었을 때 생성 결과

 

 2. condition을 주고 high-resolution 생성 결과

 

 

 3. Unconditional High Resolution Image Generation Results

 

 

## Contributions

1. 기존의 모델은 Transformer의 많은 computation으로 인해 low-resolution image generation 에서만 적용을 하였는데, VQGAN은 Transformer를 High Resolution 생성에 적용할 수 있는 방법론을 제시

2. Adversarial learning을 통해서 이미지의 중요한 local structure를 담은 codebook을 생성하는 방식을 제안