[논문리뷰] Infinity: Scaling Bitwise AutoRegressive Modeling for High-Resolution Image Synthesis

arXiv 2024. [Paper] [Page] [Github]
Jian Han, Jinlai Liu, Yi Jiang, Bin Yan, Yuqi Zhang, Zehuan Yuan, Bingyue Peng, Xiaobing Liu
ByteDance
5 Dec 2024

Introduction

최근, Visual AutoRegressive modeling (VAR)은 이미지에 대한 autoregressive (AR) 학습을 coarse-to-fine “next-scale prediction”으로 재정의했다. 이는 DiT에 비해 더 우수한 일반화 및 scaling 능력을 보여주면서도 더 적은 step이 필요하다. VAR은 LLM의 강력한 scaling 속성을 활용하고 동시에 이전 scale step을 정제하여 diffusion model의 장점도 활용할 수 있다. 그러나 AR 또는 VAR 모델에 사용되는 index-wise discrete tokenizer는 제한된 vocabulary 크기로 인해 상당한 quantization 오차에 직면하여 특히 고해상도 이미지에서 세밀한 디테일을 재구성하는 데 어려움이 있다. 생성 단계에서 index-wise token은 모호한 supervision으로 인해 시각적 디테일 손실과 로컬한 왜곡이 발생한다. 게다가 LLM에 내재된 teacher-forcing 학습의 학습-테스트 불일치는 시각적 디테일의 누적 오차를 증폭시킨다.

본 논문은 프로세스 전반에 걸쳐 index-wise token을 bitwise token으로 대체하는 bitwise modeling이라는 새로운 접근 방식을 제안하였다. Bitwise modeling 프레임워크는 다음의 세 가지 기본 모듈로 구성된다.

bitwise visual tokenizer
bitwise infinite-vocabulary classifier
bitwise self-correction

Bitwise modeling을 통해 tokenizer의 vocabulary 크기를 $2^{64}$로 확장하여 이전의 모든 AR 모델의 vocabulary를 크게 능가한다. 이 확장을 통해 continuous VAE를 크게 능가하는 재구성 품질이 가능해졌다. 또한, 학습 중에 일부 bit를 무작위로 뒤집어서 예측 실수를 시뮬레이션하고 residual feature를 re-quantize하는 bitwise self-correction을 통합하였다. 본 논문의 방법인 Infinity는 VAR 모델링의 scaling 및 속도 이점을 유지하면서도 diffusion model과 비슷한 디테일 재구성 및 생성 품질을 달성하였다.

Infinity는 AR 모델에 대한 새로운 SOTA를 세웠으며, 여러 까다로운 text-to-image 벤치마크에서 SDXL, PixArt-Sigma, DALL-E 3, Stable-Diffusion 3를 포함한 주요 diffusion model을 능가하였다. 특히, Infinity는 Stable-Diffusion 3와의 인간 평가에서 66%의 승률을 달성하고 동일한 모델 크기에서 inference 시간을 2.6배 줄였다. 또한, Infinity는 이미지 tokenizer의 vocabulary 크기와 transformer 크기를 scaling하여 이미지 생성 능력에 대한 강력한 scaling law를 보여주며, scaling함에 따라 고품질 이미지 생성의 콘텐츠와 디테일이 상당히 향상되었다.

Method

1. Visual AutoRegressive Modeling

Infinity는 visual tokenizer와 이미지 합성을 위한 transformer를 통합하였다. 학습 단계에서 샘플은 텍스트 프롬프트 $t$와 실제 이미지 $I$로 구성된다. Visual tokenizer는 먼저 이미지 $I$를 feature map $F \in \mathbb{R}^{h \times w \times d}$로 인코딩한 다음 (stride $s$), feature map $F$를 $K$개의 multi-scale residual map $(R_1, \ldots, R_K)$으로 quantize한다. $R_k$의 해상도는 $h_k \times w_k$이고 $k$가 커짐에 따라 점차 커진다. 이 residual 시퀀스를 기반으로 다음과 같이 continuous한 feature $F$를 점차 근사할 수 있다.

\[\begin{equation} F_k = \sum_{i=1}^k \textrm{up} (R_i, (h, w)) \end{equation}\]

($\textrm{up}(\cdot)$은 bilinear upsampling)

그 후, transformer는 이전 예측과 텍스트 입력을 조건으로 다음 scale의 residual $R$을 예측하는 법을 AR 방식으로 학습한다.

\[\begin{equation} p (R_1, \cdots, R_K) = \prod_{k=1}^K p (R_k \,\vert\, R_1, \ldots, R_{k-1}, \Psi (t)) \end{equation}\]

여기서 $\Psi (t)$는 Flan-T5 모델의 텍스트 임베딩이며, 추가로 cross-attention 메커니즘을 통해 예측을 더욱 가이드한다. $\Psi (t) \in \mathbb{R}^{L \times C}$는 첫 번째 scale의 입력인 $\langle \textrm{SOS} \rangle \in \mathbb{R}^{1 \times 1 \times h}$로 projection된다. 여기서 $h$는 transformer의 hidden dimension이다. Transformer는 첫 번째 scale의 $\langle \textrm{SOS} \rangle$에 따라 $R_1$을 예측해야 한다. $k$번째 scale에서 입력과 출력 레이블 $R_k$의 공간적 크기를 맞추기 위해 마지막 scale $k−1$의 다운샘플링된 feature $\tilde{F}_{k-1}$을 입력으로 사용하여 병렬로 $R_k$를 예측한다.

\[\begin{equation} \tilde{F}_{k-1} = \textrm{down}(F_{k-1}, (h_k, w_k)) \end{equation}\]

($\textrm{down}(\cdot)$은 bilinear downsampling)

$\tilde{F}_{k-1}$과 $R_k$의 공간적 크기는 $(h_k, w_k)$이며, 이전 index-wise 표현에서 예측의 모양은 $(h_k, w_k, V_d)$이다. $V_d$는 visual tokenizer의 vocabulary 크기이다. 코드 임베딩 차원이 $d$인 binary quantization의 경우 $V_d = 2^d$이며, $d$가 크면 필요한 계산 리소스가 감당할 수 없을 정도로 커진다.

Transformer는 반복되는 block의 스택으로 구성되며, 각 block에는 RoPE2d, self-attention layer, cross-attention layer, FFN layer가 포함된다. 텍스트 임베딩 $\Psi (t)$는 각 cross-attention layer에서 이미지 합성을 위한 guidance를 제공한다. 학습 단계에서는 block별 causal attention mask를 활용하여 transformer가 $F_k$를 예측할 때 $(\langle \textrm{SOS} \rangle, \tilde{F}_1, \ldots, \tilde{F}_{k-1})$에만 attention할 수 있도록 한다. Inference 단계에서는 KV-Caching을 수행하여 속도를 높이며 마스킹이 필요하지 않다.

2. Visual Tokenizer

Vocabulary 크기를 늘리면 재구성 및 생성 품질을 개선된다. 그러나 기존 tokenizer에서 vocabulary를 직접 확대하면 메모리 소비와 계산 부담이 상당히 증가한다. 이 문제를 해결하고 discrete tokenizer의 잠재력을 최대한 활용하기 위해 본 논문에서는 새로운 bitwise multi-scale residual quantizer를 제안하였다. 이는 메모리 사용량을 크게 줄여 매우 큰 어휘(ex. $2^{64}$)의 학습을 가능하게 한다.

Bitwise Multi-scale Residual Quantizer

저자들은 VAR의 원래 vector quantizer를 차원에 독립적인 bitwise quantizer로 대체하였다. 본 논문에서는 LFQ와 BSQ를 고려하였다. Multi-scale quantizer의 $K$개의 scale이 주어지면, $k$번째 scale에서 continuous한 residual 벡터 $z_k \in \mathbb{R}^d$가 다음과 같이 바이너리 출력 $q_k$로 quantize된다.

\[\begin{equation} q_k = \mathcal{Q} (z_k) = \begin{cases}\textrm{sign} (z_k) & \textrm{if LFQ} \\ \frac{1}{\sqrt{d}} \textrm{sign} (\frac{z_k}{\vert z_k \vert}) & \textrm{if BSQ}\end{cases} \end{equation}\]

Codebook 활용을 장려하기 위해 엔트로피 페널티

\[\begin{equation} \mathcal{L}_\textrm{entropy} = \mathbb{E}[H(q(z))] - H(\mathbb{E}[q(z)]) \end{equation}\]

를 사용한다. LFQ를 사용할 때 $q(z)$의 분포를 얻으려면 입력 $z$와 전체 codebook 간의 유사도를 계산해야 한다. 그러나 이는 $O(2^d)$의 감당할 수 없는 공간 및 시간 복잡도로 이어진다. 반면 BSQ의 입력과 출력은 모두 단위 벡터이므로 BSQ는 엔트로피 페널티에 대한 근사를 통해 계산 복잡도를 $O(d)$로 줄인다. 따라서 저자들은 BSQ를 채택하였다.

3. Infinite-Vocabulary Classifier

Visual tokenizer는 residual feature를 quantize하여 discrete한 레이블을 얻는다. 결과적으로, transformer는 다음 scale의 residual feature의 레이블 $y_k \in [0, V_d)^{h_k \times w_k}$를 예측하며, cross-entropy loss를 통해 학습된다. 이전 방법들은 $V_d$개의 클래스에 대한 classifier를 사용하여 레이블을 직접 예측하였다. 그러나 막대한 계산 비용과 모호한 학습의 두 가지 단점이 있다.

앞서 설명한 대로, visual tokenizer로 vocabulary 크기 $V_d$가 매우 큰 bitwise VQ-VAE를 활용한다. 예를 들어, $V_d = 2^{32}$이고 $h = 2048$인 경우 기존 classifier는 8.8조 개의 파라미터로 구성된 가중치 행렬 $W \in \mathbb{R}^{h \times V_d}$가 필요하며, 이는 계산 리소스의 한계를 초과한다.

또한 VQ-VAE는 quantization 중에 sign function을 활용하며, 다음과 같이 인덱스 레이블 $y_k (m, n)$을 얻는다.

\[\begin{equation} y_k (m, n) = \sum_{p=0}^{d-1} \unicode{x1D7D9}_{R_k (m,n,p) > 0} \cdot 2^p \end{equation}\]

Quantization 방법으로 인해, 거의 0에 가까운 feature에 대한 약간의 noise는 레이블에 상당한 변화를 초래한다. 결과적으로, 기존의 index-wise classifier는 최적화하기 어렵다.

본 논문은 계산 및 최적화 문제를 해결하기 위해 Infinite-Vocabulary Classifier (IVC)를 제안하였다. $V_d$개의 클래스를 사용하는 기존 classifier를 사용하는 대신 $d = \log_2 (V_d)$개의 binary classifier를 병렬로 사용하여 다음 scale의 residual $R_k (m, n, p)$가 양수인지 음수인지 예측한다.

IVC는 기존 classifier에 비해 메모리와 파라미터 측면에서 훨씬 효율적이다. 예를 들어, $V_d = 2^{32}$이고 $h = 2048$일 때 13만 개의 파라미터만 필요하다. 게다가 일부 차원에서 모델을 혼란스럽게 하는 거의 0에 가까운 값이 존재하더라도 다른 차원에서의 supervision은 여전히 명확하다. 따라서 기존 index-wise classifier에 비해 IVC를 최적화하기가 더 쉽다.

4. Bitwise Self-Correction

Teacher-forcing 학습의 취약점

VAR은 LLM에서 teacher-forcing 학습을 물려받았다. 그러나 비전에서의 next-scale prediction은 언어에서의 next-token prediction과 상당히 다르다. 구체적으로, 모든 scale에서 residual $R_k$를 얻을 때까지 전체 이미지를 디코딩할 수 없다. Teacher-forcing 학습은 시각적 생성에 심각한 학습-테스트 불일치를 초래한다. 특히 teacher-forcing 학습은 transformer가 실수를 인식하고 수정할 수 있는 능력 없이 각 scale의 feature만 정제하게 한다. 이전 scale에서 발생한 실수는 이후 scale에 전파되고 증폭되어 결국 생성된 이미지를 엉망으로 만든다.

본 논문에서는 이 문제를 해결하기 위해 Bitwise Self-Correction (BSC)을 제안하였다. 구체적으로, $R_k$의 비트들을 $[0, p]$에서 균일하게 샘플링한 확률로 무작위로 뒤집어서 $R_k^\textrm{flip}$을 얻는다. 이는 $k$번째 scale의 예측에서 발생한 오차의 다양한 강도를 모방한다.

$R_k^\textrm{flip}$에는 오차가 있지만 $R_k$에는 없다. $k$번째 scale에서 예측으로 $R_k$를 $R_k^\textrm{flip}$으로 대체한 후 transformer 입력 $\tilde{F}_k$를 다시 계산한다. 또한 새로운 $R_{k+1}$을 얻기 위해 re-quantization을 수행한다. 즉, transformer의 입력과 레이블을 수정한다. 추가적인 계산 비용이 필요하지 않고 원래의 병렬 학습 특성이 손상되지 않는다.

각 scale은 동일한 random flipping과 re-quantization 프로세스를 거친다. Transformer는 부분적으로 랜덤하게 flip된 feature를 입력으로 사용하여 예측 오차를 고려한다. Re-quantize된 비트 레이블을 통해 transformer는 이전 예측에서 발생한 오차를 자동으로 수정할 수 있다. 이런 방식으로 teacher-forcing 학습으로 인해 발생한 학습-테스트 불일치를 해결하고 Infinity가 자체 수정 능력을 갖도록 한다.

5. Dynamic Aspect Ratios and Position Encoding

Infinity는 다양한 종횡비로 사실적인 이미지를 생성할 수 있는데, 이는 정사각형 이미지만 생성할 수 있는 VAR과 상당히 다르다. 다양한 종횡비 이미지를 생성하는 데 있어 주요 장애물은 두 가지 측면에 있다.

첫째, 다양한 종횡비에 따라 $R_k$의 높이 $h_k$와 너비 $w_k$를 정의해야 한다. 본 논문에서는 각 종횡비에 대해 $\{(h_1^r, w_1^r), \ldots, (h_K^r, w_K^r)\}$로 scale schedule이라고도 하는 scale 목록을 미리 정의한다. 각 튜플 $(h_k^r, w_k^r)$의 종횡비가 $r$과 거의 같도록 한다. 또한 동일한 scale $k$에서 다른 종횡비의 경우 $h_k^r \times w_k^r$의 면적을 거의 같게 유지하여 학습 시퀀스 길이가 거의 같도록 한다.

둘째, 다양한 scale과 종횡비의 feature를 처리하기 위해 해상도를 고려하는 위치 인코딩 방법을 신중하게 설계해야 한다. 본 논문에서는 각 scale의 feature에 RoPE2d를 적용하여 이미지의 고유한 2D 구조를 보존하였다. 또한 학습 가능한 scale 임베딩을 활용하여 서로 다른 scale의 feature 간 혼동을 피한다. Feature에 element-wise하게 적용된 학습 가능한 absolute positional encoding과 비교할 때, scale에 적용된 학습 가능한 임베딩은 파라미터가 적게 필요하고 다양한 시퀀스 길이에 적응할 수 있으며 최적화하기가 더 쉽다.

Experiments

데이터셋: LAION, COYO, OpenImages
구현 디테일
- Pre-training
  - 해상도: 512
  - iteration: 11만
  - batch size: 4096
  - learning rate: $6 \times 10^{-5}$
- Fine-tuning
  - 해상도: 1024
  - iteration: 6만
  - batch size: 2048
  - learning rate: $2 \times 10^{-5}$