diffusion과 비슷한 quality에다가 더 효율적인 1024 image 생성 Auto Regressive model.
- (A100 기준) 4.5 - 7.7배 높은 throughput
- 3.1 - 5.9배 낮은 latency
기존의 AR 모델들은 discrete tokenizer를 이용한다는 단점 때문에 recon 퀄리티가 diffusion보다 좋지 않았다. 그리고 1024px image를 학습할 때 드는 비용도 높았다. sequence가 제곱으로 길어지기 때문에
hybrid tokenizer 제시.
auto-encoder의 continuous한 latent를 두가지 컴포넌트로 분리한다.
- 큰그림을 그리는 discrete tokens
- discrete token만으로 충분히 표현되지 않은 residual을 채워주는 continuous tokens
discrete token은 scalable-resolution AR model로 모델링하고, continuous 부분은 37M의 작은 residual diffusion model로 학습한다.
AR과 Diffusion이라는 이미지 생성의 두 축 - diffusion이 생성 퀄리티는 더 좋지만 inference시 iterative하게 약 8.5배 리소스가 많이 든다.
This paper addresses the following question:
Can we develop an autoregressive model that matches the visual generation quality of diffusion models while still being significantly faster?
현재까지 AR의 두가지 단점
- Discrete tokenizer가 reconstruction 퀄리티를 심각하게 낮춘다. diffusion에서 쓰는 continuous tokenizer에 비교해 - Consequently, AR models have a lower generation upper bound and struggle to accurately model fine image details
- high-resolution image 생성의 어려움
AR 모델 중에서는 VAR이 가장 좋다고 본다
Hybrid models
최근들어 꽤 발전하고 있는걸로 보임
HART
다른 방법들은 latent space자체는 discrete이나 continuous 중 한가지로만 진행하는 것과 달리 HART는 discrete과 continuou token 모두를 통합된 transformer로 모델링한다.
- Hybrid tokenizer
VQ codebook을 이용해 image generation을 할 수 있는데, tokenizer의 recon quality가 이미지 생성의 상한선을 정한다. finite vocab codebook의 제한으로 인해 discrete tokenizer는 세밀함이 필요한 이미지를 잘 못 생성하는 경향이 있다.
우선 CNN base의 visual encoder를 이용해서 이미지를 latent space에서의 continuous token으로 바꾼다. 그리고 이 토큰을 discrete token으로 quantization한다(VAR 처럼 multiple scale로). 이 과정에서 변환된 discrete 토큰을 전부 사용해도 기존의 continuous token과 완전히 일치할 수는 없고, 손실이 발생한다.
우리는 이 차이를 residual token이라고 부른다 - residual diffusion에 의해 모델링된다.
Hybrid tokenizer 학습
- visual encoder, quantizer (codebook), visual decoder (from a pretrained discrete VAR tokenizer) 학습
- encoder, quantizer는 freeze하고 decoder만 학습
각 학습에서 50% 확률로 continuous token으로 바로 recon하도록 학습하거나(기존의 auto encoder와 동일), VQ quantization을 통해서 나온 discrete token으로 학습하도록 한다(quantizer를 학습하는 개념).
결론적으로 두 방법 전부의 성능이 좋다. continuous token의 recon 성능은 SDXL과 비슷, discre token recon 성능도 SOTA와 비슷
Decoder의 입장에서 보면,, continuous와 discrete latent가 비슷하게 형성되게 된다(둘다 같은 이미지로 recon 해야하니까) → continuous latent의 모델링도 어렵지 않게 할 수 있게 된다.
Discussion : reconstruction FID가 낮다고 해서 generation FID가 좋다는걸 의미하진 않는다(우리 얘기?). continuous와 discrete token을 따로 모델링했을 때도 recon 성능은 비슷했지만 generation quality는 크게 차이가 났다.
Modeling
HART는 최종적으로는 Continuous image token을 모델링한다고 볼 수 있는데 대신 이 토큰이 2가지로 분리되어서 모델링된다 - discrete, residual
Scalable-resolution autoregressive transformer
discrete token modeling 방법 자체는 VAR을 1) text-to-image + 2) high resolution으로 확장시킨 방법.
- text token을 visual token에 concat해서 학습한다(time-wise concat 같음) → cross-attention보다 파라미터가 25% 적다.
- low resolution checkpoint 위에 학습하고, 모든 absolute PE를 relative PE로 바꿨다. position index들은 text token에서 시작해서 visual token까지 연결되도록 했다.
Residual diffusion
MAR에서처럼, 큰 DiT가 이 residual을 위해 필요하진 않다고 판단했다 → 37M의 작은 MLP로 residual diffusion modeling.
MAR에서 diffusion의 output은 continuous token 자체이지만 HART에서는 residual만 예측한다 - 효율성의 측면에서 중요. 특히 inference 때 동일한 퀄리티를 위해 MAR의 30-50에 비해 8 step만 중요하다 → 4-6배 빠르다.
- input of MLP
- discrete token들
- last layer hidden states
- output
- discrete 토큰마다의 residual(continuous)
discrete token + continuous residual → continuous token ← 이게 Decoder로 들어가서 이미지 생성
Setup
AR model은 600M, 1B, 2B 실험 - diffusion MLP는 37M
- class condition으로 adaLN써서 학습한 다음 그 위에 text condition으로 학습, 할 때는 adaLN 전부 제거.
- text encoder는 Qwen2-1.5B 사용, SANA 처럼 프롬프트 바꿔서 사용
성능
text alignment도 SD-XL 2.6B보다 좋다(CLIP이 맞다면)
Hybrid tokenizer의 성능
continuous token의 성능은 생성 모델의 상한선을 결정하고
discrete token의 성능은 학습 과정에서 diffusion model이 가질 부담을 보여줌.
class condition 생성 VAR, MAR과 비교
X는 discrete token만 사용, 체크는 residual까지 적용. FID가 좋아진다.
Share article
Subscribe to our newsletter
