[論文レビュー] An Image is Worth 32 Tokens for Reconstruction and Generation

Qihang Yu, Mark Weber+4 more|arXiv (Cornell University)|Jun 11, 2024

Anatomy and Medical Technology被引用数 4

ひとこと要約

この論文は、局所的再帰性を活用して、256×256の画像をたった32個の離散トークンで表現する1Dトランスフォーマー基盤の画像トークナイザー、TiTokを紹介する。2D VQベースの手法と比較して、トークン数を顕著に削減している。ImageNet-256で1.97 gFIDを達成し、SOTA性能を発揮するとともに、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潜在シーケンス表現を提案すること。
• よりコンパクトで効率的かつ意味的に豊かな画像表現を可能にし、高解像度画像生成を支援すること。
• 1Dトークナイゼーションが、特に低トークン数の状況下で、再構築忠実度と生成速度の両面で2D対比手法を上回ることを示すこと。
• アブレーションスタディおよびMaskGIT やDiT-XL/2といった最先端モデルとの比較を通じて、1Dトークナイゼーションの有効性を検証すること。

提案手法

• TiTokは、フラット化された画像パッチと学習可能な1D潜在トークンを連結して処理するビジョントランスフォーマー（ViT）エンコーダーを用いる。
• ベクトル量子化（VQ）により、符号化された特徴量が離散トークンにマッピングされ、離散潜在コードのコンパクトな1Dシーケンスが形成される。
• マスクされた1Dトークンシーケンスから元の画像を再構築するためのViTデコーダーを採用し、再構築損失を用いたエンドツーエンド学習が可能になる。
• 限られたトークン数での品質向上を図るため、プロキシコードを用いた2段階学習戦略を採用する。
• 生成時における自己回帰的生成を模倣するために、マスキングスケジュールを適用し、線形およびアークコサインスケジュールに対して、アブレーションによりロバストネスの向上が示された。
• 2Dグリッド制約を解除することで、各潜在トークンが非局所的で意味的に関連する画像領域に注目できるようになり、表現の効率性が向上する。

実験結果

リサーチクエスチョン

• RQ1 1D潜在シーケンス表現は、2Dグリッドベースのトークナイザーと比較して、競争力ある画像再構築・生成性能を達成できるか？
• RQ2 1Dトークナイゼーションは、忠実度を損なわず、画像冗長性を効果的に活用してトークン数を削減できるか？
• RQ3 モデルサイズやトークン数の増加に伴う1Dトークナイゼーションの性能スケーリング特性、特に32や64といった低トークン数での性能はいかがなものか？
• RQ4 2D対比手法と比較して、1Dトークナイゼーションは生成速度とスループットをどの程度向上させるか？
• RQ5 2段階学習におけるプロキシコードは、1DトークナイゼーションでSOTA性能を達成するために不可欠であるか？

主な発見

• TiTokは、256×256の画像をたった32個の離散トークンにまで圧縮し、2D VQGAN手法と比較して64倍の削減を実現した。
• ImageNet-256では、1.97 gFIDを達成し、MaskGITベースライン（4.21 gFID）を4.21 gFID上回った。
• ImageNet-512では、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設計の優位性を実証した。
• プロキシコードを用いた2段階学習により、TiTokは1.70 rFIDを達成し、MaskGIT-VQGAN（2.28 rFID）を上回った。これは、訓練レシピの有効性を示している。