Lecciones de la escalabilidad de tokenizadores visuales para reconstrucción y generación

Learnings from Scaling Visual Tokenizers for Reconstruction and Generation

January 16, 2025

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

cs.AI

Resumen

La tokenización visual a través del autoencodificación potencia los modelos generativos de imágenes y videos de última generación al comprimir píxeles en un espacio latente. Aunque escalar generadores basados en Transformadores ha sido fundamental para los avances recientes, el componente tokenizador en sí rara vez se escala, dejando preguntas abiertas sobre cómo las elecciones de diseño del autoencoder influyen tanto en su objetivo de reconstrucción como en el rendimiento generativo posterior. Nuestro trabajo tiene como objetivo llevar a cabo una exploración de la escala en los autoencoders para llenar este vacío. Para facilitar esta exploración, reemplazamos la típica estructura convolucional con una arquitectura mejorada de Vision Transformer para la Tokenización (ViTok). Entrenamos ViTok en conjuntos de datos de imágenes y videos a gran escala que superan con creces a ImageNet-1K, eliminando las restricciones de datos en la escala del tokenizador. Primero estudiamos cómo afecta la escala del cuello de botella del autoencoder tanto a la reconstrucción como a la generación, y encontramos que si bien está altamente correlacionado con la reconstrucción, su relación con la generación es más compleja. A continuación, exploramos el efecto de escalar por separado el codificador y el decodificador de los autoencoders en el rendimiento de reconstrucción y generación. Crucialmente, descubrimos que escalar el codificador aporta ganancias mínimas tanto en la reconstrucción como en la generación, mientras que escalar el decodificador impulsa la reconstrucción, pero los beneficios para la generación son mixtos. Basándonos en nuestra exploración, diseñamos ViTok como un autoencoder ligero que logra un rendimiento competitivo con los autoencoders de última generación en tareas de reconstrucción de ImageNet-1K y COCO (256p y 512p), superando a los autoencoders existentes en la reconstrucción de videos de 16 fotogramas a 128p para UCF-101, todo con 2-5 veces menos FLOPs. Al integrarse con Transformadores de Difusión, ViTok demuestra un rendimiento competitivo en la generación de imágenes para ImageNet-1K y establece nuevos puntos de referencia de última generación para la generación de videos condicionados por clase en 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.