Inzichten uit het opschalen van visuele tokenizers voor reconstructie en generatie

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

Samenvatting

Visuele tokenisatie via auto-encoding versterkt toonaangevende beeld- en videogeneratiemodellen door pixels samen te drukken in een latente ruimte. Hoewel het schalen van op Transformer gebaseerde generatoren centraal heeft gestaan in recente vooruitgang, wordt het tokenizer-component zelf zelden geschaald, waardoor er nog vragen openstaan over hoe ontwerpkeuzes van de auto-encoder zowel de reconstructiedoelstelling als de prestaties van downstream generatie beïnvloeden. Ons werk heeft tot doel een verkenning van schalen in auto-encoders uit te voeren om deze leemte op te vullen. Om deze verkenning te vergemakkelijken, vervangen we de typische convolutionele ruggengraat door een verbeterde Vision Transformer-architectuur voor Tokenisatie (ViTok). We trainen ViTok op grootschalige beeld- en videodatasets die ver uitstijgen boven ImageNet-1K, waardoor de gegevensbeperkingen voor het schalen van de tokenizer worden weggenomen. We bestuderen eerst hoe het schalen van de bottleneck van de auto-encoder zowel de reconstructie als de generatie beïnvloedt - en ontdekken dat hoewel het sterk gecorreleerd is met reconstructie, de relatie met generatie complexer is. Vervolgens onderzochten we het effect van afzonderlijk schalen van de encoder en decoder van de auto-encoders op de prestaties van reconstructie en generatie. Cruciaal is dat we ontdekken dat het schalen van de encoder minimale voordelen oplevert voor zowel reconstructie als generatie, terwijl het schalen van de decoder de reconstructie verbetert, maar de voordelen voor generatie gemengd zijn. Voortbouwend op onze verkenning ontwerpen we ViTok als een lichtgewicht auto-encoder die concurrerende prestaties levert met toonaangevende auto-encoders op ImageNet-1K en COCO-reconstructietaken (256p en 512p), terwijl het bestaande auto-encoders overtreft op 16-frame 128p video-reconstructie voor UCF-101, allemaal met 2-5x minder FLOPs. Wanneer geïntegreerd met Diffusion Transformers, toont ViTok concurrerende prestaties op beeldgeneratie voor ImageNet-1K en stelt nieuwe toonaangevende benchmarks voor klasse-geconditioneerde videogeneratie op 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.