[논문 리뷰] VQ-GAN: Taming Transformers for High-Resolution Image Synthesis (2021)

 

이번에 리뷰할 논문은 Taming Transformers for High-Resolution Image Synthesis 입니다.

https://arxiv.org/abs/2012.09841

Taming Transformers for High-Resolution Image Synthesis

 

 

 

 

 

 

Introduction

Transformer는 장거리 상호작용(long-range interaction) 학습에 특화되어 있으며, NLP 뿐만 아니라 오디오, 비전 분야에서도 SOTA 성능을 보여주고 있습니다. 하지만 고해상도 이미지와 같이 Sequence 길이가 매우 긴 데이터에서는 계산 비용이 Sequence 길이의 제곱으로 증가(quadratic)하며, 실제 적용이 매우 비효율적이라는 단점이 존재합니다.

 

CNN(Convolutional Neural Network)은 local적으로 상호작용을 모델링하기에 효율적이고, convolutional kernel로 인한 inductive bias 덕분에 local level의 이미지 구조를 잘 학습하는데 유리합니다. 하지만 장거리 상호작용을 잘 모델링하지 못하고, 공간 불변성이라는 편향이 있기 때문에 전체적, 의미론적 구조 이해에는 한계가 존재합니다.

 

반면, Transformer는 표현력은 높지만, 지역 구조에 대한 prior이 없기 때문에, low-level의 구조 학습에 비효율적이며, 계산 비용이 크다는 한계가 존재합니다.

 

---

 

본 논문에서는 이런 CNN과 Transformer의 장점을 결합합니다. 

1. CNN으로 low-level의 지역 구조를 학습하여, 맥락 정보를 담은 Codebook을 만듭니다. (이미지를 discrete token sequence로 변환)

 

2. Transformer는 그 Codebook index sequence 상에서, 전역적 조합을 모델링 합니다. (이미지 내 구성 요소들 간의 장거리 관계를 효과적으로 학습)

 

3. VQ-VAE 기반 구조 + Adversarial loss 를 사용하여 local structure을 잘 보존하도록 합니다. 이를 통해 Transforemr가 low-level 통계까지 모델링할 필요가 없도록 설계합니다.

 

---

 

본 논문의 main contribution은 다음과 같습니다.

• Transformer로 메가픽셀급 고해상도 이미지를 생성할 수 있는 최초의 구조 제안
• 조건부 생성(Conditional synthesis) 가능
  • 객체 클래스, 세그멘테이션 등 non-spatial / spatial 조건 모두 지원
• 기존 CNN 기반 Codebook 방식보다 일관되게 더 좋은 성능

Proposed approach enables transformers to synthesize high-resolution images like this one, which contains 1280x460 pixels.

 

 

 

 

 

 

Approach

본 논문에서 목표로 하고자 하는 것은 Transformer 모델의 뛰어난 학습 능력을 활용하여, 메가픽셀 수준의 고해상도 이미지 생성에 적용하는 것입니다. 기존 연구들에서도 Transformer를 이미지 생성에 활용하여 64x64 크기까지는 좋은 성능을 보여주었지만, Sequence 길이에 따른 계산 비용이 제곱으로 quadratic하게 증가하기 때문에 해상도를 단순히 높이는 방식으로 확장하는 것은 불가능했습니다. 

 

고해상도 이미지를 생성하기 위해서는 이미지의 전역적 구성(global composition)을 이해할 수 있는 모델이 필요합니다. 즉, 지역적으로 자연스러운 패턴과 전역적으로 일관된 구조를 동시에 생성할 수 있어야 하죠.

 

이를 위해, 저자는 이미지를 픽셀 단위로 표현하는 대신, perceptually rich image constituents들의 codebook으로 표현합니다. 즉, 효율적인 codebook 학습 방법을 통해 구성 요소들의 description length를 줄이고, 그들의 global interrelations는 Transformer로 모델링합니다. 

 

 

이러한 접근을 통해 unconditional 이미지 생성 뿐만 아니라 conditional 생성에도 적용이 가능해집니다.

 

 

 

 

 

 

Learning an Effective Codebook of Image Constituents

우선 Transformer의 높은 표현력을 이미지 생성에 활용하기 위해서는 이미지 구성 요소들을 Sequence 형태로 표현해야 합니다. 다만, 픽셀 단위로 sequence를 구성하는 것은 너무 복잡하기 때문에, 학습된 표현으로 이루어진 discrete codebook을 사용합니다. 

 

만약 어떤 이미지 $ x \in \mathbb{R}^{H \times W \times 3} $ 이 있다고 하면, Codebook 항목들의 공간적 집합은

$ z_{q} \in \mathbb{R}^{h \times w \times n_{z}} $ 으로 표현됩니다. 

• $ H \times W $ : 원본 이미지 해상도
• $ h \times w $ : 인코더를 거쳐 down sample된 latent 해상도
  • 예) $ H = 256, W = 256 $ 이면 $h = 16, w = 16 $ (downsample factor 16)
• $ n_{z} $ : 코드북 벡터의 차원 (예: 256)

동일하게, 이 표현은 $ h \times w $ 길이의 sequence (codebook 인덱스들의 sequence) 로도 나타낼 수 있습니다. 

 

---

 

이제 Codebook을 학습해야 하는데, 다음 두 가지의 구조를 학습하게 됩니다.

• Encoder $E$ 와 Decoder $G$ 로 이루어진 합성곱 기반 모델
• 학습된 discrete codebook $ Z = \{ z_{k} \}^{K}_{k=1} \subset \mathbb{R}^{n_{z}} $

구체적으로, 주어진 이미지 $x$를 다음과 같이 근사합니다.

 

$$ \hat{x} = G(z_{q}) $$

 

여기서 original image에 대해 latent vector로 변환하기 위해 다음과 같이 Encoding 과정을 거치며,

 

$$ \hat{z} = E(x) \in \mathbb{R}^{h \times w \times n_{z}} $$

 

양자화(Quantization)을 통해, 각 위치 $ (i, j) $의 latent vector를 가장 가까운 codebook 항목으로 치환하죠.

 

$$ z_q = q(\hat{z}) := \underset{z_k \in Z}{\arg\min} \, \| \hat{z}_{ij} - z_k \| $$

 

마지막으로 reconstruction을 수행하여 이미지를 생성합니다.

 

$$ \hat{x} = G(z_{q}) = G(q(E(x))) $$

 

또한 위의 양자화 연산은 비미분 가능(non-differentiable) 이기 때문에, Straight-through gradient estimator (STE)를 사용해 gradient를 전달합니다. 이를 통해 모델과 codebook을 end-to-end로 학습하는 것이 가능해집니다.

 

학습을 위한 최종적인 손실 함수는 다음과 같습니다.

 

$$ \mathcal{L}_{\mathrm{VQ}}(E, G, Z) = \| x - \hat{x} \|_2^2 
+ \| \mathrm{sg}[E(x)] - z_q \|_2^2 
+ \| \mathrm{sg}[z_q] - E(x) \|_2^2 $$

 

위 식은 총 3가지의 항으로 구성되어있으며, 

• 1 term : reconstruction loss
• 2 term : update codebook
• 3 term : commitment loss

의 역할을 가집니다. 위 손실함수는 VQ-VAE 논문에서 처음 제안되었으며, 자세한 내용은 다음 글을 참고하시기 바랍니다.

2025.03.24 - [논문 리뷰/Image generation] - [논문 리뷰] VQ-VAE: Neural Discrete Representation Learning (2018)

[논문 리뷰] VQ-VAE: Neural Discrete Representation Learning (2018)

 

이미지를 latent image constituents의 분포로 표현하기 위해 transformer를 사용하려면, 압축률을 최대한 높이면서도 풍부한 codebook을 학습해야 합니다. 하지만 VQ-VAE의 loss와 같이 이미지를 압축해서 discrete token으로 표현하는 코드북을 학습할 때 단순히 $ l_{2} loss $를 사용하면, 이미지가 흐릿하게 재구성되는 문제가 발생했습니다.

 

이를 위해 본 논문에서는 VQ-VQE의 변형인 VQ-GAN을 제안하고, 판별기(discriminator)와 지각 손실(perceptual loss)을 활용하여 높은 압축률에서도 지각적으로 좋은 품질을 유지할 수 있도록 합니다. 직관적으로 말하자면, GAN을 이용해 더 "진짜 같은" 복원 이미지를 만들도록 학습한다는 것이죠.

 

우선 VQ-GAN에 대해 들어가기 전 GAN 에 대해 간단히 설명 후 넘어가도록 하겠습니다.

 

GAN의 경우 판별기와 생성기가 존재하며, 서로 적대적으로 학습이 진행됩니다.

 

판별기의 경우 진짜 이미지 $x$는 1로, 가짜 이미지 $\hat{x}$는 0으로 분류를 하도록 학습되며, 다음과 같은 likelihood를 최대화 하도록 학습됩니다.

 

$$ \log D(x) + \log (1 - D(\hat{x})) $$

 

생성자의 경우 처음에는 $ \log (1-D(\hat{x})) $을 최소화하는 것으로 제안되었지만, 요즘은 $ -\log D(\hat{x}) $ 를 최소화하는 것이 더 학습이 잘 된다고 합니다. 즉, 생성자는 가짜 이미지 $\hat{x}$를 1로 예측하도록 생성자를 업데이트 합니다. 

 

---

 

여기까지 간단하게 GAN loss function에 대해 알아보았으니, VQ-GAN의 손실함수에 대해 알아보도록 하겠습니다. 앞서 설명하였던 GAN과 크게 다르지 않으며, 단순히 진짜 이미지와 Decoder로부터 생성된 가짜 이미지를 판별할 수 있도록 하는 판별자 손실함수를 적용하면 됩니다.

 

$$ \mathcal{L}_{\mathrm{GAN}}(\{E, G, Z\}, D) = \log D(x) + \log (1 - D(\hat{x})) $$

 

즉, $D(x)$ : 진짜 이미지를 진짜라고 예측하고, $D(\hat{x})$ : 가짜 이미지를 가짜라고 예측하도록 최적화 합니다. 여기서 $\hat{x} = G(g(E(x)))$는 인코딩 및 양자화된 latent vector를 디코더로 복원한 이미지를 의미합니다.

 

---

 

다시 정리하자면, VQ-VAE의 손실함수를 통해 코드북과 모델에 대해 학습을 진행하지만, 단순히 $l_{2} loss$만으로는 고품질의 이미지를 생성하는데 한계가 있기 때문에, 생성된 이미지와 진짜 이미지를 판별하는 판별기를 추가하여 학습시키는 GAN 손실함수를 추가하여 고화질의 이미지도 생성할 수 있도록 합니다.

 

$$ Q^{*} = \{ E^{*}, G^{*}, Z^{*} \} $$

 

목적은 위 식과 같이 인코더, 디코더, 코드북을 최적화하는 것이며, 

 

$$ Q^{*} = \underset{E, G, Z}{\arg\min} \; \underset{D}{\max} \; \mathbb{E}_{x \sim p(x)} \left[ \mathcal{L}_{\mathrm{VQ}}(E, G, Z) + \lambda \, \mathcal{L}_{\mathrm{GAN}}(\{E, G, Z\}, D) \right] $$

 

최종적인 손실함수는 위와 같이 VQ loss와 GAN loss를 가중합한 형태로 구성됩니다. VQ loss의 경우 재구성 품질을 최적화하는 데 사용되며, GAN loss는 시각적으로 진짜 같은지 판별할 수 있도록 최적화됩니다.

 

하지만 이런 GAN 기반 손실함수에는 한가지 문제점이 존재하는데, GAN loss는 gradient가 너무 크면, 너무 강하게 작동되어 학습이 불안정해질 수 있다는 점입니다. 

 

$$ \lambda = \frac{\left\| \nabla_G \mathcal{L}_{\mathrm{rec}} \right\|}{\left\| \nabla_G \mathcal{L}_{\mathrm{GAN}} \right\| + \delta} $$

 

그래서 위와 같이 가중치 $ \lambda $를 구하여, 현재 GAN loss가 지나치게 커지면 $ \lambda $를 줄이고, GAN loss가 너무 약하면 $ \lambda $를 늘려서 영향력을 키워 adaptive balace를 맞추도록 합니다.

 

 

 

 

 

 

Learning the Composition of Image with Transformers

앞서 우리는 인코더 $E$와 디코더 $G$를 준비하였습니다. 이제 이미지를 인코딩된 코드북 인덱스 sequence로 표현할 수 있죠. 우선 이미지 $x$의 양자화된 인코딩은 다음과 같습니다. 

 

$$ z_{q} = q(E(x)) \in \mathbb{R}^{h \times w \times n_{z}} $$

 

또한 이는 코드북 인덱스 sequence $ s \in \{ 0,..., \| Z \| - 1 \}^{h \times w} $ 와 동등합니다. 이때, $Z$ 는 codebook의 entry 개수가 됩니다.

 

즉, 각 위치 $(i,j)$에 대해 

 

$$ s_{ij} = k \ \ \  such \ that \ \ \  (z_{q})_{ij} = z_{k} $$

 

가 됩니다.

 

시퀀스 $s$이 인덱스를 다시 대응하는 코드북 항목으로 매핑하면, 

 

$$ z_{q} = z_{s_{ij}} $$

 

가 되며, 이를 디코더 $G$에 통과시키면 복원된 이미지 $ \hat{x} = G(z_{q}) $ 가 얻어집니다.

 

---

 

Autoregressive Generation

다음으로, 시퀀스 $s$에 순서를 지정하면, 이미지 생성을 자기회귀(next-index prediction) 문제로 정의할 수 있습니다. 즉, 이전 인덱스 $ s_{<i} $가 주어졌을 때 다음 인덱스 $ s_{i} $의 분포를 예측합니다.

 

$$ p(s) = \prod_{i} p(s_i \mid s_{<i}) $$

 

이젠, 전체 데이터 표현의 log-likelihood를 최대화하는 것이 가능합니다.

 

$$ \mathcal{L}_{\mathrm{Transformer}} = \mathbb{E}_{x \sim p(x)} \left[ -\log p(s) \right] $$

 

 

Conditioned Synthesis

이미지 생성 작업에서는 사용자가 추가 정보를 통해 생성 과정에 대한 제어를 요구하는 경우가 많습니다. 예를 들어, 조건이 되는 정보 $c$는 이미지의 클래스 라벨 이거나 다른 이미지 그 자체가 될 수 있죠.

 

즉, 우리는 조건 $c$가 주어졌을 때 시퀀스 $s$의 확률을 학습해야 합니다. 

 

$$ p(s \mid c) = \prod_{i} p(s_i \mid s_{<i}, c) $$

 

위의 경우, 조건 정보 $c$가 클래스 정보와 같이 discrete한 정보가 들어올 경우 단순하게 condition을 주어 conditioned probability를 구하는 것이 가능합니다.

 

하지만 조건 정보 $c$가 공간적 정보(spatial extent)를 가진다면, 해당 정보를 또 다른 VQ-GAN을 통해 index-based representation $r \in \{0, \dots, |Z_c| - 1\}^{h_c \times w_c}$ 으로 인코딩합니다. 

 

$$ p(s_i \mid s_{<i}, r) $$

 

Transformer의 자기회귀 구조때문에 단순히 $r$ 시퀀스를 $s$ 앞에 붙인 후 위와 같이 조건부 확률을 학습하면 됩니다.

 

 

 

 

 

 

Generating High-Resolution Images

Transformer의 attention mechanism은 input sequence $s$의 길이, 즉, $h \times w$에 제한을 둡니다. 이를 해결하기 위해서는 VQ-GAN의 down-sampling 블록 수 $m$을 조정하여 입력 이미지 크기 $H \times W$를 다음과 같이 줄일 수 있습니다.

 

$$ h = \frac{H}{2^{m}}, \quad w = \frac{W}{2^{m}} $$

 

하지만 $m$이 너무 커지면, reconstruction 품질이 심하게 저하되는 문제가 발생합니다. 따라서 본 논문에서는 메가 픽셀급 이미지 생성을 위해서 학습 시 이미지를 patch 단위로 잘라 학습하고, 생성 시에는 Sliding window 방식으로 순차적으로 생성합니다.

 

Sliding attention window