Visual Autoregressive Modeling: Scalable Image Generation via Next-Scale Prediction [review]
VAR paper
Jul 16, 2024
읽은 이유
visual modal - image 영역에서 autoregressive model을 사용해서 diffusion보다 높은 성능(어떤 성능인지는 아직 모름)을 달성했다고 들어서 관심이 생김. 사실이라면 꽤 중요하고, vision 생성에서 diffusion이외의 방식들이 주목을 받을 수 있는 결과라고 생각함.
요약
기존의 vision autoregressive model이 갖는 문제를 multi-scale의 token map을 만들어 예측하는 next-scale이라는 방법으로 해결하면서 DiT보다 좋은 성능을 냈다.
그리고 AR의 특성, scaling law도 가지는걸 보였고 이건 매우 중요한 사실이다.
autoregressive하게 해서 diffusion을 이긴 첫 논문(내가 알기론)이기도 하고, 아직 uncondition으로 ImageNet 256에서만 실험이 된 만큼 발전과 활용될 여지가 많다.
Abstract
Visual AutoRegressive modeling(VAR)을 제안한다. 특히 coarse-to-fine의 형태로, next token prediction 태스크에 새로운 학습 방식-구조를 추가해 next-scale prediction을 수행해서 진행함.
처음으로 GPT-style의 AR model로 diffusion transformer를 이미지 생성 분야에서 이겼다. ImageNet 256*256 벤치마크에서 FID, IS를 이겼고 inference 속도가 20배 빠르다고 함.
DiT를 이미지 퀄리티, inference 스피드, data efficiency, scalability 측면에서 이겼다.
1) VAR model을 scaling up하는건 LLM에서 관찰된 clear power-law를 비슷하게 보인다.(모델의 사이즈가 클수록 성능이 리니어하게 증가한다. 약 0.998)
2) downstream 태스크인 image in-painting, out-painting, editing도 제로샷으로 잘한다.
→ 이건 AR model의 형태를 차용한 것 만으로 LLM의 두가지 중요한 특성 : scaling-law와 zero-shot : 을 띄게 된다는 사실을 의미한다.
Introduction의 직관
sequence의 다음 토큰을 예측한다는 self-supervised 학습 방식으로 인해 auto-regressive 모델들의 scalability(이건 scaling law와 함께 모델의 크기를 키움에 따라 성능이 예측됨으로 인해 작게 테스트 → 크게 적용이라는 방식으로 많은 리소스를 아꼈다.)와 generalizability 성능이 매우 좋았고, 이게 발전을 이끌었다. 이건 매우 많은 unlabeled data로 학습한다는 AGI의 essence에 가깝다.
VQGAN과 DALL-E가 이미지에서도 AR model을 시도했음. 근데 성능은 디퓨전보다 안좋았다.
AR modeling은 데이터의 순서가 있어야한다. 우리는 이미지 순서를 어떻게 할지 고려했다. 근데 단순히 한 이미지의 순서뿐만 아니라, 사람이 이미지를 떠올리는 방식-대략 이런 모양 → 디테일-까지 반영된 데이터를 원했다. 그래서 multi-scale, coarse-to-fine 구조를 적용함(본 논문에서 가장 중요한 개념)
그래서 일단 이미지를 multi-scale token map으로 인코딩하고 시작한다. 1*1에서부터 resolution을 증가시켜나감. 학습과 inference 자체는 LM의 AR과 똑같다. 직접적을, GPT-2와 비슷한 구조의 transformer architecture를 차용했다.
FID가 1.73이고, IS가 350.2다. 체리피킹이 아니라 당당하게 가장 최신의 DiT와도 비교했는데 이겼다.
한 이미지안에서의 토큰 순서는 ViT에서 한것처럼 raster-scan manner(왼쪽 위에서 오른쪽 아래로)로 정렬했다.
여기까지 + figure를 보고 든 생각
결국 여기서도 AR model에다가 DDPM, score matching method처럼 abstract → detail의 방법이 적용되었다. 이게 얼마나 중요한걸까? SDE에서 설명한게 여기도 적용되는걸까
AR과 coase-to-fine 이 두가지가 가장 중요한 적용사항일텐데 전자가 스케일링과 강건도, 후자가 퀄리티에 영향을 미친 요소들이 아닐까 추측.
동일한 해상도의 image 학습/생성을 위해 기존의 visual AR보다 많은 토큰을 예측해야하는데, 그래도 충분히 빠른가?
Diffusion model이 모델 크기를 키웠을 경우 성능이 그대로 유지되지 않는다는 사실을 처음 알았다. N년 전과 다르게 하드웨어가 발달하면서 scalability도 꽤 중요한 요소일 수 있다고 생각하는데, 가치있는 논문으로 생각됨
Method
visual tokenizer를 통해 2d image를 1D token sequence로 바꾼 후 학습한다. Decoder-only LM의 학습 next token prediction과 동일하다. 이전 토큰 시퀀스로 조건부확률을 계산해 다음 토큰을 예측한다.
Tokenization
1) discrete하게 토큰화 2) ordering 필요
VQGAN의 autoencoder를 써서
vector quantizer에서와 같이 image를 encoder 태우고, C차원의 벡터 V개를 가진 learnable coderbook으로 구성된 quantizer를 사용한다(VQGAN을 봐야 이해가능한 부분). image를 인코딩해서 나온 continuouse한 vector를 유클리디안 거리상으로 가장 가까운 code로 매핑해서 사용한다.
우선 이 quantizer부터 학습하고 시작한다. loss는 VQGAN과 같다. 이미지를 VQ한 다음 다시 reconstruction한 뒤 reconstruction loss(L2 거리)와 perceptual loss(LPIPS 사용.. 이거 믿을만한가?), discriminative loss를 사용해서 학습한다.
그래도 2D긴 한데 1D로 나열해서 학습한다. to maximize likelihood.
근데 그냥 이렇게 auto-regressive 적용하면 단점이 있다.
- 이미지를 encoder 태워 feature map을 얻은 뒤 얘를 quatize하는데, 이 때 feature map 자체는 bidirectional한 정보를 가지고 압축된다. 그래서 얘를 토큰화하면 각 토큰에도 약간의 bidirectional한 정보가 담긴다. 근데 여기서는 next token으로 unidirectioinal하게 학습하는데, 기본 가정 자체와 조금 모순됨. 현재 가정에서의 학습이 잘되려면 t번째 token이 1 ~ t-1 토큰에 의해서만 결정될 수 있는 형태여야 하는데 그게 아니다.
- zero-shot으로 bidirectional한 태스크는 못한다. 제일 왼쪽 위의 이미지를 아래 이미지에 기반해서 예측하지 못하는 형태. 이건 어쩔 수 없고
- 1D sequence로 flattening하기 때문에 spatial한 정보가 손실된다.
- n 토큰 예측이라고하면 autoregressive step을 n^2번 해야하고 computational cost는 O(n^6)이다. 왜 6제곱이지?
전부 appendix에서 자세히 설명함. 전부 중요한 사항들이긴 하다. 아래에서 우리는 이걸 좀 해결했다 라는 식으로 풀어낸다.
Next-scale prediction
next-scale prediction을 한다. 인코더를 태워 얻은 피처맵을 K개의 multi-scale token map(r1, r2, …, rk)으로 양자화한다. 마지막 K번째의 scale이 기존 이미지의 h,w와 같다.
이건 스케일을 조건(prefix)으로 다음 스케일 예측.
k번째의 autoregressive step에서, 모든 분포는 rk의 prefix와 k번째라는 조건에 의존해 parallel하게 생성된다.(?? 이렇게해서 어떻게 잘되는지 이해 안됨) 즉, 한번에 예측하는건 다음 토큰 한개가 아니라 다음 해상도의 모든 토큰이다. 어떻게 이걸 실제 구현하는지는 안나와있어서 repo를 봐야한다.
이렇게 하기 위해서 block-wise causal attention mask를 사용해서 같은 해상도의 이전 토큰을 참고하지 못하게 마스킹 했다.
따라서 discussion
- K번째 scale에서는 k-1까지의 scale만 참고하도록 해서 unidirectional한 특성이 생겨났다.
- 본 연구에서는 flattening이 일어나지 않기 때문에 spatial한 특성이 유지된다.
- computational cost가 O(n^4)로 감소한다. 같은 해상도 안에서 parallel하게 생성하니까.
그리고 이걸 위해서 당연히 multi-scale quantization autoencoder가 필요하다. 일단 feature map을 얻은 뒤 각 해상도에 맞추기 위해 interpolation해서 피처를 만든뒤 양자화 한다.
자세한 알고리즘.
Al 1. 보면 각 해상도 k마다 양자화 → recon해서 복구한 값을 기존 피처맵에서 뺀다. → RQ-VAE에서 사용한 residual quantization 개념을 적용한건데, 각 해상도의 양자화 벡터가 갖는 정보는 이전 해상도의 정보를 제외한 정보(residual의 개념)이다. 더 잘된다고 한다. 참 신기한 세상이다
Al 2. 예측한 token을 기반으로 reconstruction할 때도, 모든 해상도의 토큰이 사용된다. 각 해상도의 값은 residual이니까 이렇게 해야지!
그리고 나머지에서는 scaling 효과가 얼마나 잘 적용되는지를 자세하게 분석해놓았다. scaling law가 적용된다는 것을 꽤 중요하게 생각하는걸로 보인다. 실제로 practical하게 매우 중요한 요소이기도 하고.
text-based generation 연구 하고 있다고 함.
Share article
Subscribe to our newsletter