VibeToken: Scaling 1D Image Tokenizers and Autoregressive Models for Dynamic Resolution Generations

以下は、論文「VibeToken: Scaling 1D Image Tokenizers and Autoregressive Models for Dynamic Resolution Generations」の詳細な技術的サマリーです。

1. 問題提起

自己回帰（AR）画像生成モデル（例：LlamaGen）は競争力のある性能を示していますが、拡散モデルに比べて生産環境での採用に向けた大きな障壁に直面しています。主なボトルネックは、解像度とアスペクト比に関するスケーラビリティです：

• トークンの爆発的増加: 従来の 2D トークナイザー（例：VQGAN）は、画像解像度に比例して線形的に増加するトークン系列を生成します。1024×1024 の画像の場合、数千のトークン（例：4,096 トークン）が発生します。
• 二次的な計算コスト: AR モデルは自己注意（self-attention）に依存しており、これは系列長 $T$ に対して $O(T^2)$ の複雑さを持ちます。解像度が高くなるにつれて、推論時の FLOPs が二次的に増加し、高解像度生成が計算的に非現実的になります。
• 硬直性: ほとんどの AR モデルは固定解像度（例：256×256）で訓練されており、再訓練や複雑な超解像パイプラインなしでは、任意のアスペクト比やより高い解像度への汎化が困難です。
• 拡散モデルの支配: 拡散モデルは任意の解像度を自然に処理できるため、柔軟で生産グレードのユースケースにおいて AR モデルを凌駕しています。

2. 手法

著者らは、VibeToken（新しいトークナイザー）とVibeToken-Gen（AR 生成器）という 2 部構成のソリューションを提案します。

A. VibeToken: 解像度非依存の 1D トークナイザー

VibeToken は、任意の解像度の画像を固定された、ユーザーが制御可能なトークン数（例：32–256）に符号化するよう設計された 1D Transformer ベースの画像トークナイザーであり、トークン数と空間解像度を切り離します。

主なアーキテクチャ上の革新点は以下の通りです：

1. 動的グリッド位置符号化: 固定された絶対位置符号化の代わりに、VibeToken は入力パッチグリッド（$T_H \times T_W$）に一致するように微分的にサイズ変更（バイリニア/バイキュービック補間を介して）される学習済みグリッドの符号化を使用します。これにより、品質の劣化なしに任意のアスペクト比を処理できます。
2. 適応型パッチ埋め込み: FlexiViT に着想を得て、モデルは訓練および推論中に可変のパッチサイズ（$k$）をサポートします。重みサイズ変更演算子がベース投影重みを現在のパッチサイズにマッピングし、モデルが再パラメータ化なしに異なる解像度に適応できるようにします。
3. 適応型デコーダー解像度: デコーダーは固定された高解像度（パッチサイズの 4 倍）でピクセルを予測し、学習可能な 2D 畳み込みダウンサンプリング層を使用してターゲット解像度にマッピングします。これにより、個別のアップスケーラーモジュールなしでネイティブな超解像が可能になります。
4. 動的長さトークナイゼーション: モデルは潜在長さ $L \in [32, 256]$ の均一サンプリングで訓練されます。エンコーダーは正確に $L$ 個のトークンを生成し、デコーダーは正確に $L$ 個を消費するため、パディングや固定長の制約が不要になります。
5. 量子化: 圧縮と再構成品質のバランスを取るため、8 つのコードブックを用いたマルチベクトル量子化（MVQ）を使用し、語彙サイズを 32,768 にします。

B. VibeToken-Gen: 効率的な AR 生成器

VibeToken-Gen は、VibeToken が生成する短く解像度非依存のトークン系列上で動作する、クラス条件付き AR モデル（LlamaGen アーキテクチャに基づく）です。

• 一定の計算予算: トークン系列長 $L$ は出力解像度に関係なく固定（例：64 トークン）であるため、推論時の FLOPs は一定に保たれます。
• 条件付け: 幾何学的な伸縮アーティファクトを防ぐため、クラスラベルに加えてターゲット解像度 $(H, W)$ とアスペクト比に基づいて条件付けを行います。
• 訓練: ImageNet-1k 上で、固定解像度のベースラインと同等のトークン予算を使用して、256×256 から 512×512 までの混合解像度とアスペクト比で訓練されます。

3. 主な貢献

1. 解像度非依存トークナイゼーション: 任意の解像度とアスペクト比をネイティブにサポートし、かつ短く動的なトークン予算（32–256 トークン）を維持する、最初の 1D Transformer トークナイザー。
2. 計算効率: AR 推論コストを画像解像度から切り離します。VibeToken-Gen は任意の解像度に対して一定の179 GFLOPsを維持するのに対し、LlamaGen などのベースラインは 1024×1024 の画像で約 11 TFLOPsにスケーリングします（63.4 倍の効率向上）。
3. ネイティブ超解像: 個別のアップスケーリングモジュールなしに、低解像度入力（または任意のターゲット）から高解像度画像を生成することをアーキテクチャがサポートします。
4. 最先端の性能: 2D トークナイザーや拡散モデルと比較しても、大幅に少ないトークン数を使用しながら、競争力のある再構成（rFID）および生成（gFID）品質を達成します。

4. 実験結果

本論文は、ImageNet-1k および FFHQ データセット上で、VibeToken と VibeToken-Gen を、256×256 から 1024×1024 までのさまざまな解像度とアスペクト比で評価しました。

• 再構成品質:
  • VibeToken-LL は 256×256 で0.40 rFID、1024×1024 で2.40 rFIDを達成します。
  • 解像度汎化において他の 1D トークナイザー（例：TiTok、One-D-Piece）を上回り、圧縮効率において 2D トークナイザー（例：IBQ、MAGViT）と同等かそれ以上の性能を発揮します。
• 生成品質:
  • VibeToken-Gen (XXL) は、わずか64 トークンを使用して 1024×1024 画像を合成し、gFID 3.94を達成します。
  • 比較:
    • 拡散モデル (NiT) 対比: VibeToken-Gen は 1024×1024 において2.35 倍高速（0.46 秒対 1.08 秒）であり、より高い品質（3.94 gFID 対 5.87 gFID）を達成します。
    • AR ベースライン (LlamaGen) 対比: LlamaGen は 1024×1024 で約 11 TFLOPs を必要とするのに対し、VibeToken-Gen は一定の 179 GFLOPs を使用します。
• 効率性:
  • トークン数は 1024×1024 の画像において、2D トークナイザーの約 4,096 から64 トークンに削減されました。
  • 高解像度生成における推論速度は 1 秒未満（0.46 秒）であり、同様の解像度のベースライン AR モデルの 32 秒以上と比較して大幅に高速です。

5. 意義

この研究は、柔軟性と効率性の面で、自己回帰モデルと拡散モデルの間のギャップを埋めます。

• 生産への実用性: AR モデルが一定の計算量で任意の解像度を処理できるようにすることで、VibeToken は、入力/出力の寸法が変化する実世界のアプリケーション（例：モバイルアプリ、多様な表示形式）において AR 生成を実用的なものにします。
• スケーラビリティ: AR モデルが二次的な計算の爆発なしに高解像度にスケーリングできることを実証し、高解像度生成 AI への道筋として拡散モデルが唯一の選択肢であるという考えに挑戦します。
• 将来の方向性: このアプローチは、生成コンテンツの空間的複雑さに関係なく、固定されたユーザー制御可能な計算予算の下で動作できる統合マルチモーダルモデル（テキスト、画像、動画）への扉を開きます。

要約すると、VibeTokenは、解像度適応の負担を生成モデルから柔軟で解像度非依存のトークナイザーへ移行させることで、AR 生成における「解像度スケーリング」の問題を解決し、現在の最先端の代替手段よりも高速で、安価で、柔軟なシステムを実現します。