Visual Autoregressive Modeling: Scalable Image Generation via Next-Scale Prediction 리뷰

Tian, Keyu, Yi Jiang, Zehuan Yuan, Bingyue Peng, and Liwei Wang. "Visual Autoregressive Modeling: Scalable Image Generation via Next-Scale Prediction." arXiv preprint arXiv:2404.02905 (2024).

 

Paperswithcode에 상위 논문으로 플로팅 되어 있길래 하는 리뷰NLP에서 사용하는 autoregressive modeling을 비전에 적용한 첫 모델? 인듯 하다.

 

Abstract

우리는 autoregressive learning을 이용한 이미지 생성에 새로운 패러다임을 제시한다. 다음 토큰을 예측하는 전형적인 방법에서 벗어나 단계별로 해상도/스케일을 예측한다. GPT스타일의 AR모델인 VAR(Visual AutoRegressive)은 이미지 생성에서 diffusion transformers를 능가한다. Diffusion Transformer(DiT)에 대해 이미지의 품질, 추론 속도, 데이터 효율성, 크기 측면에서 모두 VAR이 더 뛰어난 성능을 보인다.또한 zero-shot으로 image in-painting, out-painting, editing이 가능한 것을 확인했다.

 

(b)의 AR의 경우 ViT에서 사용되는 이미지 토큰들을 이용해 다음 토큰을 예측하여 이미지를 생성하는데 VAR은 (c)와 같이 저수준에서 온전한 이미지까지 화소를 높이며(coarse-to-fine 라고 표현) 이미지를 생성한다고 한다. 때문에 VAR은 multi-scale VQVAE를 사용한다.

 

 

Method

1. Preliminary : autoregressive modeling via next-token prediction

각각의 토큰 x 가 \(V\) 크기의 vocabulary에서 \(x_t\in[V]\) 일 때, \(x=(x_1, x_2,...,x_T)\)라고 가정하면 autoregressive model은 현재 토큰 \(x_t\)를 예측하기 위해 \((x_1,x_2,...,x_{t-1})\)에만 의존한다. 이러한 단방향 토큰에 의존한 추측은 sequence x에 대해 조건부 T 확률의 곱으로 표현할 수 있다 (1).

이미지는 2D이기 때문에 next-token prediction의 형태로 모델링 하기 위해서는 다음과 같은 작업이 필요하다.

1) 이미지를 discrete한 토큰의 형태로 나누어야 하고, 2) 단방향 토큰 모델링을 위한 1D형태를 정의해야 한다.

 

이미지를 discrete token으로 변형하기 위해서는 최근 주로 사용되는 방법인 quantized autoencoder을 사용하였다.

이는 이미지 feature map \(f\in R^{h\times w\times c}\)를 \(q\in[V]^{h\times w}\)의 형태로 변환시킨다.

\(im\) : raw image

\(\xi(\cdot)\) : encoder

Q(·) : quantizer

quantizer은 전형적으로 \(V\) 벡터를 가지는 학습가능한 코드북 \(Z\in R^{V\times C}\)을 포함한다. 이는 각각의 특성맵 \(f^{(i,j)}\) 를 유클리디안 거리의 가장 가까운 인덱스 \(q^{(i,j)}\)로 맵핑시킨다.

 

\(q\in [V]^{h\times w}\)에 속한 이미지 토큰들은 2차원 격자로 정렬된다. 이전의 AR모델들은 2d 격자의 q를 1차원으로 펼친 후 다음 토큰을 예측하기 위해 likelihood를 최대화 하는 방식으로 모델을 훈련하였다.

 

하지만 다음과 같은 문제점들이 발생

1) Mathematical premise violation : 인코더는 전형적으로 모든 i,j에 대해 inter-dependent feature vectors \(f^{(i,j)}\)을 제공한다. 이는 양방향의 관계성을 띄기 때문에 단방향성의 추측을 사용하는 AR모델에 적합하지 않다.

 

2) Structural degradation : flattening은 이미지의 공간적 특성 정보를 저해시킨다. 예를 들어 한 토큰과 이웃한 토큰들은 서로 연관 관계를 갖지만 linear sequence에서는 이러한 공간적 특성이 상쇄된다.

 

3) Inefficiency : self-attention 트랜스포머를 이용한 autoregressive 스텝은 \(O(n^2)\)의 시간과 \(O(n^6)\)의 컴퓨팅 비용이 발생한다.

 

2. Visual autoregressive modeling via next-scale prediction

우리는 위에서 설명한 next-token prediction이 아닌 next-scale prediction 전략으로 재정의 한다. 우리는 하나의 a 토큰이 아닌 entire token을 이용해 autoregressive unit을 정의한다. 우리는 특징맵 \(f \in R^{h\times w\times C}\)을 \(h_k\times w_k\)크기까지 점차 화소를 올리는 K multi-scale의 토큰맵 \((r_1,r_2,...r_k)\)으로 양자화 한다. 이는 \(r_k\)일 때 원본 화소인 \(h\times w\)의 특징맵과 일치한다.

(1)과 수식은 똑같지만 다른점은, 기존에는 이미지 토큰맵들을 feature map으로 만들어 다음 토큰을 예측하는 방법을 사용했다면 (6)은 하나의 이미지를 특징맵으로 양자화 시켜 이미지의 화소를 키워가며 다음 scale을 예측하는 방법이다.

Fig4 에서 VAR의 훈련 과정이 묘사되어 있는데, VAR은 두개의 훈련 과정으로 나누어 진다. multi-scale quantized autoencoder은 K token maps \(R=(r_1,r_2,...,r_K)\)으로 이미지를 인코딩하고 reconstruction loss를 계산한다(좌측).

VAR transformer은 likelihood를 최대화 시키거나 cross-entropy loss를 최소화 시키도록 훈련되는데, VQVAE로부터 ground truth인 \(R=(r_1,r_2,...,r_k)\)인 토큰맵을 제공받는다.

 

우리는 VQGAN과 같은 구조를 사용하였지만 multi-scale로 양자화 시키기 위해 multi-scale quantization layer을 추가하였다. 이는 K개의 cnn layer이며 (0.03M)의 추가적인 파라미터가 요구된다.

 

Results

자세한 결과들은 논문을 참고하길 바란다...

사실 아직 자세히 모르는 분야라서 잘 이해가 가질 않는다..