Enseignements de la mise à l'échelle des tokenizers visuels pour la reconstruction et la génération

Learnings from Scaling Visual Tokenizers for Reconstruction and Generation

January 16, 2025

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

cs.AI

Résumé

La tokenisation visuelle via l'auto-codage renforce les modèles génératifs d'images et de vidéos de pointe en comprimant les pixels dans un espace latent. Bien que l'extension des générateurs basés sur les Transformers ait été au cœur des récentes avancées, le composant de tokenisation lui-même est rarement étendu, laissant des questions ouvertes sur la manière dont les choix de conception de l'auto-codeur influencent à la fois son objectif de reconstruction et les performances génératives en aval. Notre travail vise à mener une exploration de l'extension des auto-codeurs pour combler cette lacune. Pour faciliter cette exploration, nous remplaçons l'épine dorsale convolutive typique par une architecture améliorée de Vision Transformer pour la Tokenisation (ViTok). Nous formons ViTok sur des ensembles de données d'images et de vidéos à grande échelle dépassant largement ImageNet-1K, éliminant les contraintes de données sur l'extension du tokeniseur. Nous étudions d'abord comment l'extension de l'étranglement de l'auto-codeur affecte à la fois la reconstruction et la génération - et constatons que bien qu'elle soit fortement corrélée à la reconstruction, sa relation avec la génération est plus complexe. Nous explorons ensuite l'effet de l'extension séparée de l'encodeur et du décodeur des auto-codeurs sur les performances de reconstruction et de génération. De manière cruciale, nous constatons que l'extension de l'encodeur apporte des gains minimes pour la reconstruction ou la génération, tandis que l'extension du décodeur améliore la reconstruction mais les avantages pour la génération sont mitigés. En nous appuyant sur notre exploration, nous concevons ViTok comme un auto-codeur léger qui atteint des performances compétitives avec des auto-codeurs de pointe sur ImageNet-1K et les tâches de reconstruction COCO (256p et 512p) tout en surpassant les auto-codeurs existants sur la reconstruction vidéo 16 images de 128p pour UCF-101, le tout avec 2 à 5 fois moins d'opérations en virgule flottante (FLOPs). Lorsqu'il est intégré aux Transformers de Diffusion, ViTok démontre des performances compétitives en matière de génération d'images pour ImageNet-1K et établit de nouveaux benchmarks de pointe pour la génération vidéo conditionnelle par classe sur 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.