再構築と生成のためのビジュアルトークナイザのスケーリングに関する学び

Learnings from Scaling Visual Tokenizers for Reconstruction and Generation

January 16, 2025

著者: Philippe Hansen-Estruch, David Yan, Ching-Yao Chung, Orr Zohar, Jialiang Wang, Tingbo Hou, Tao Xu, Sriram Vishwanath, Peter Vajda, Xinlei Chen

cs.AI

要旨

自動符号化を介した視覚的トークン化は、ピクセルを潜在空間に圧縮することで、最先端の画像およびビデオ生成モデルを強化します。最近の進歩においてTransformerベースのジェネレータのスケーリングが中心的である一方、トークナイザー部分自体はほとんどスケーリングされておらず、自動符号化の設計選択が再構成の目的と下流の生成パフォーマンスの両方にどのように影響するかという未解決の問題が残されています。本研究は、この空白を埋めるために自動符号化のスケーリングに関する探索を行うことを目的としています。この探索を容易にするために、通常の畳み込みバックボーンを強化したVision Transformerアーキテクチャをトークナイゼーション用に使用したViTokを導入します。ViTokをImageNet-1Kをはるかに超える大規模な画像およびビデオデータセットで訓練し、トークナイザーのスケーリングに関するデータ制約を除去します。まず、自動符号化のボトルネックのスケーリングが再構成と生成の両方にどのように影響するかを調査し、再構成と強く相関する一方で生成との関係はより複雑であることがわかりました。次に、自動符号化のエンコーダーとデコーダーを別々にスケーリングした場合の再構成と生成パフォーマンスへの影響を探りました。重要なのは、エンコーダーのスケーリングは再構成または生成のいずれに対してもわずかな利益しかもたらさず、一方でデコーダーのスケーリングは再構成を向上させますが、生成に対する利点は混在しています。私たちの探索を基に、ViTokを軽量な自動符号化器として設計し、ImageNet-1KおよびCOCOの再構成タスク（256pおよび512p）で最先端の自動符号化器と競争力のあるパフォーマンスを達成し、UCF-101の16フレーム128pビデオ再構成では既存の自動符号化器を2〜5倍のFLOPsで上回りました。Diffusion Transformersと統合すると、ViTokはImageNet-1Kで画像生成において競争力のあるパフォーマンスを発揮し、UCF-101でのクラス条件付きビデオ生成の新たな最先端ベンチマークを設定しました。

English

Visual tokenization via auto-encoding empowers state-of-the-art image and video generative models by compressing pixels into a latent space. Although scaling Transformer-based generators has been central to recent advances, the tokenizer component itself is rarely scaled, leaving open questions about how auto-encoder design choices influence both its objective of reconstruction and downstream generative performance. Our work aims to conduct an exploration of scaling in auto-encoders to fill in this blank. To facilitate this exploration, we replace the typical convolutional backbone with an enhanced Vision Transformer architecture for Tokenization (ViTok). We train ViTok on large-scale image and video datasets far exceeding ImageNet-1K, removing data constraints on tokenizer scaling. We first study how scaling the auto-encoder bottleneck affects both reconstruction and generation -- and find that while it is highly correlated with reconstruction, its relationship with generation is more complex. We next explored the effect of separately scaling the auto-encoders' encoder and decoder on reconstruction and generation performance. Crucially, we find that scaling the encoder yields minimal gains for either reconstruction or generation, while scaling the decoder boosts reconstruction but the benefits for generation are mixed. Building on our exploration, we design ViTok as a lightweight auto-encoder that achieves competitive performance with state-of-the-art auto-encoders on ImageNet-1K and COCO reconstruction tasks (256p and 512p) while outperforming existing auto-encoders on 16-frame 128p video reconstruction for UCF-101, all with 2-5x fewer FLOPs. When integrated with Diffusion Transformers, ViTok demonstrates competitive performance on image generation for ImageNet-1K and sets new state-of-the-art benchmarks for class-conditional video generation on UCF-101.