[논문 리뷰] An Image is Worth 32 Tokens for Reconstruction and Generation

Qihang Yu, Mark Weber외 4명|arXiv (Cornell University)|2024. 06. 11.

Anatomy and Medical Technology인용 수 4

한 줄 요약

이 논문은 지역적 부재성의 특성을 활용하여 $256\times256$ 이미지를 단지 32개의 이산 토큰으로 표현하는 1D 트랜스포머 기반 이미지 토크나이저인 TiTok을 소개한다. 이는 2D VQ 기반 방법에 비해 토큰 수를 크게 줄였으며, ImageNet-256에서 1.97 gFID의 최고 성능을 기록하고, 512×512 해상도에서 최대 74배 빠른 생성 속도를 달성한다. 이는 압축된 1D 잠재 시퀀스가 효율적이고 고해상도의 정밀한 이미지 재구성 및 생성을 가능하게 함을 보여준다.

ABSTRACT

Recent advancements in generative models have highlighted the crucial role of image tokenization in the efficient synthesis of high-resolution images. Tokenization, which transforms images into latent representations, reduces computational demands compared to directly processing pixels and enhances the effectiveness and efficiency of the generation process. Prior methods, such as VQGAN, typically utilize 2D latent grids with fixed downsampling factors. However, these 2D tokenizations face challenges in managing the inherent redundancies present in images, where adjacent regions frequently display similarities. To overcome this issue, we introduce Transformer-based 1-Dimensional Tokenizer (TiTok), an innovative approach that tokenizes images into 1D latent sequences. TiTok provides a more compact latent representation, yielding substantially more efficient and effective representations than conventional techniques. For example, a 256 x 256 x 3 image can be reduced to just 32 discrete tokens, a significant reduction from the 256 or 1024 tokens obtained by prior methods. Despite its compact nature, TiTok achieves competitive performance to state-of-the-art approaches. Specifically, using the same generator framework, TiTok attains 1.97 gFID, outperforming MaskGIT baseline significantly by 4.21 at ImageNet 256 x 256 benchmark. The advantages of TiTok become even more significant when it comes to higher resolution. At ImageNet 512 x 512 benchmark, TiTok not only outperforms state-of-the-art diffusion model DiT-XL/2 (gFID 2.74 vs. 3.04), but also reduces the image tokens by 64x, leading to 410x faster generation process. Our best-performing variant can significantly surpasses DiT-XL/2 (gFID 2.13 vs. 3.04) while still generating high-quality samples 74x faster.

연구 동기 및 목표

• 이미지 토크나이제이션에서 2D 구조적 제약 조건이 고품질 이미지 표현을 위해 필수적인가를 조사하는 것.
• 이미지의 중복성을 효율적으로 처리하지 못하는 2D 토크나이저의 비효율성을 해결하기 위해 1D 잠재 시퀀스 표현 방식을 제안하는 것.
• 더 컴act하고 효율적이며 의미적으로 풍부한 이미지 표현을 가능하게 하여 고해상도 이미지 생성을 지원하는 것.
• 1D 토크나이제이션이 낮은 토큰 수에서 재구성 정밀도와 생성 속도 면에서 2D 대비 뛰어난 성능을 보일 수 있음을 입증하는 것.
• MaskGIT 및 DiT-XL/2와 같은 최신 모델과의 비교 및 추론 분석을 통해 1D 토크나이제이션의 효과성을 검증하는 것.

제안 방법

• TiTok는 이미지 패치를 평탄화하여 학습 가능한 1D 잠재 토큰과 연결한 후, 비전 트랜스포머(ViT) 인코더를 사용해 처리한다.
• 벡터 양자화(VQ) 기법을 통해 인코딩된 특징을 이산 토큰으로 매핑하여 압축된 1D 이산 잠재 코드 시퀀스를 형성한다.
• 마스킹된 1D 토큰 시퀀스에서 원본 이미지를 재구성하기 위해 ViT 디코더를 사용하며, 재구성 손실을 통해 엔드 투 엔드 학습을 가능하게 한다.
• 제한된 토큰 수를 사용할 경우 토큰 품질을 향상시키기 위해 프oxy 코드를 활용한 이중 단계 학습 전략을 적용한다.
• 생성 과정에서 자동회귀적 생성을 시뮬레이션하기 위해 마스킹 스케줄을 적용하며, 추론 분석을 통해 선형 및 아코스 스케줄에 대한 강건성이 향상됨을 확인하였다.
• 2D 격자 제약 조건을 해체함으로써 각 잠재 토큰이 국소적이지 않은 의미적으로 관련된 이미지 영역에 접근할 수 있게 되어 표현 효율성이 향상된다.

실험 결과

연구 질문

• RQ1 2D 격자 기반 토크나이저에 비해 1D 잠재 시퀀스 표현 방식이 경쟁적인 이미지 재구성 및 생성 성능을 달성할 수 있는가?
• RQ2 1D 토크나이제이션은 정밀도를 희생시키지 않고 이미지의 중복성을 효과적으로 활용하여 토큰 수를 줄일 수 있는가?
• RQ3 특히 32개 또는 64개와 같은 낮은 토큰 수에서 1D 토크나이제이션의 성능이 모델 크기 및 토큰 수에 따라 어떻게 스케일링되는가?
• RQ4 2D 대비 1D 토크나이제이션은 생성 속도와 처리량 면에서 얼마나 향상되는가?
• RQ5 2단계 학습 전략과 프록시 코드가 1D 토크나이제이션에서 최고 성능을 달성하기 위해 필수적인가?

주요 결과

• TiTok는 $256\times256$ 이미지를 단지 32개의 이산 토큰으로 압축하여, 2D VQGAN 방법 대비 64배의 토큰 수 감소를 이룬다.
• ImageNet-256에서 TiTok는 1.97 gFID를 기록하며, MaskGIT 기준선 대비 4.21 gFID 향상된 성능을 보였다.
• ImageNet-512에서 TiTok는 2.13 gFID를 달성하여 DiT-XL/2(3.04 gFID)를 초월했으며, 샘플 생성 속도는 최대 74배 빠르게 구현하였다.
• 디코더 미세조정 및 더 큰 코드북을 사용할 경우, TiTok-L-32는 재구성에서 2.21 rFID 성능을 기록하여 기준선 대비 5.8배 향상되었다.
• TiTok-B-64의 1D 버전은 2D 대비 10.43 rFID 향상(5.15 vs. 15.58)을 기록하여, 압축된 잠재 공간에서 1D 아키텍처의 우수성을 입증하였다.
• 프록시 코드를 활용한 이중 단계 학습 전략 덕분에 TiTok는 1.70 rFID 성능을 기록하며, MaskGIT-VQGAN(2.28 rFID)을 능가하여 학습 레시피의 효과성을 입증하였다.