Lições da Escala de Tokenizadores Visuais para Reconstrução e Geração

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

Resumo

A tokenização visual via autoencoders capacita modelos gerativos de imagem e vídeo de última geração comprimindo pixels em um espaço latente. Embora a escalabilidade de geradores baseados em Transformers tenha sido central para avanços recentes, o componente tokenizador em si raramente é escalado, deixando questões em aberto sobre como as escolhas de design do autoencoder influenciam tanto seu objetivo de reconstrução quanto o desempenho gerativo subsequente. Nosso trabalho tem como objetivo realizar uma exploração da escalabilidade em autoencoders para preencher essa lacuna. Para facilitar essa exploração, substituímos a espinha dorsal convolucional típica por uma arquitetura aprimorada de Vision Transformer para Tokenização (ViTok). Treinamos o ViTok em conjuntos de dados de imagem e vídeo em larga escala que excedem em muito o ImageNet-1K, removendo as restrições de dados na escalabilidade do tokenizador. Primeiramente, estudamos como a escalabilidade do gargalo do autoencoder afeta tanto a reconstrução quanto a geração - e descobrimos que, embora esteja altamente correlacionada com a reconstrução, sua relação com a geração é mais complexa. Em seguida, exploramos o efeito de escalar separadamente o codificador e o decodificador dos autoencoders no desempenho de reconstrução e geração. Crucialmente, descobrimos que escalar o codificador resulta em ganhos mínimos tanto para a reconstrução quanto para a geração, enquanto escalar o decodificador impulsiona a reconstrução, mas os benefícios para a geração são mistos. Com base em nossa exploração, projetamos o ViTok como um autoencoder leve que alcança desempenho competitivo com autoencoders de última geração em tarefas de reconstrução do ImageNet-1K e COCO (256p e 512p), superando os autoencoders existentes na reconstrução de vídeo de 16 quadros em 128p para UCF-101, tudo com 2-5 vezes menos FLOPs. Quando integrado com os Transformadores de Difusão, o ViTok demonstra desempenho competitivo na geração de imagens para o ImageNet-1K e estabelece novos benchmarks de última geração para a geração de vídeo condicional de classe no 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.