Apprendimenti dalla scalabilità dei tokenizzatori visivi per la ricostruzione e la generazione

Learnings from Scaling Visual Tokenizers for Reconstruction and Generation

January 16, 2025

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

cs.AI

Abstract

La tokenizzazione visuale tramite auto-codifica potenzia i modelli generativi di immagini e video all'avanguardia comprimendo i pixel in uno spazio latente. Sebbene l'incremento di scala dei generatori basati su Transformer sia stato centrale negli avanzamenti recenti, il componente tokenizer stesso raramente viene scalato, lasciando aperte domande su come le scelte progettuali dell'auto-codificatore influenzino sia l'obiettivo di ricostruzione che le prestazioni generative successive. Il nostro lavoro mira a condurre un'esplorazione dello scaling negli auto-codificatori per colmare questa lacuna. Per agevolare questa esplorazione, sostituiamo la tipica struttura convoluzionale di base con un'architettura migliorata del Vision Transformer per la Tokenizzazione (ViTok). Alleniamo ViTok su set di dati di immagini e video su larga scala che superano di gran lunga ImageNet-1K, eliminando vincoli di dati sullo scaling del tokenizer. Iniziamo studiando come lo scaling del collo di bottiglia dell'auto-codificatore influenzi sia la ricostruzione che la generazione, e scopriamo che sebbene sia altamente correlato alla ricostruzione, la sua relazione con la generazione è più complessa. Successivamente esploriamo l'effetto dello scaling separato dell'encoder e del decoder degli auto-codificatori sulle prestazioni di ricostruzione e generazione. In modo cruciale, scopriamo che lo scaling dell'encoder porta a guadagni minimi sia per la ricostruzione che per la generazione, mentre lo scaling del decoder potenzia la ricostruzione ma i benefici per la generazione sono misti. Basandoci sulla nostra esplorazione, progettiamo ViTok come un auto-codificatore leggero che raggiunge prestazioni competitive con auto-codificatori all'avanguardia su ImageNet-1K e compiti di ricostruzione COCO (256p e 512p), superando gli auto-codificatori esistenti nella ricostruzione video a 16 frame 128p per UCF-101, il tutto con 2-5 volte meno FLOPs. Integrato con i Diffusion Transformers, ViTok dimostra prestazioni competitive nella generazione di immagini per ImageNet-1K e stabilisce nuovi benchmark all'avanguardia per la generazione video condizionata alla classe su 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.