Erkenntnisse aus dem Skalieren von visuellen Tokenizern für Rekonstruktion und Generierung

Learnings from Scaling Visual Tokenizers for Reconstruction and Generation

January 16, 2025

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

cs.AI

Zusammenfassung

Die visuelle Tokenisierung mittels Auto-Encoding stärkt modernste Bild- und Videogenerierungsmodelle, indem Pixel in einen latenten Raum komprimiert werden. Obwohl das Skalieren von Transformer-basierten Generatoren entscheidend für jüngste Fortschritte war, wird die Tokenizer-Komponente selbst selten skaliert, was Fragen offen lässt, wie Designentscheidungen des Auto-Encoders sowohl die Rekonstruktionsziele als auch die nachgelagerte generative Leistung beeinflussen. Unsere Arbeit zielt darauf ab, eine Untersuchung des Scalings in Auto-Encodern durchzuführen, um diese Lücke zu schließen. Um diese Untersuchung zu erleichtern, ersetzen wir das typische konvolutionale Grundgerüst durch eine verbesserte Vision-Transformer-Architektur für Tokenisierung (ViTok). Wir trainieren ViTok auf umfangreichen Bild- und Videodatensätzen, die weit über ImageNet-1K hinausgehen, und beseitigen Datenbeschränkungen beim Skalieren des Tokenizers. Zunächst untersuchen wir, wie sich das Skalieren des Engpasses des Auto-Encoders sowohl auf die Rekonstruktion als auch auf die Generierung auswirkt - und stellen fest, dass es zwar stark mit der Rekonstruktion korreliert ist, aber seine Beziehung zur Generierung komplexer ist. Anschließend untersuchten wir die Auswirkung des separaten Scalings des Encoders und Decoders der Auto-Encoder auf die Rekonstruktions- und Generierungsleistung. Entscheidend ist, dass wir feststellen, dass das Skalieren des Encoders nur minimale Gewinne für die Rekonstruktion oder Generierung bringt, während das Skalieren des Decoders die Rekonstruktion verbessert, aber die Vorteile für die Generierung gemischt sind. Aufbauend auf unserer Untersuchung entwerfen wir ViTok als einen leichten Auto-Encoder, der eine wettbewerbsfähige Leistung mit modernsten Auto-Encodern bei ImageNet-1K und COCO-Rekonstruktionsaufgaben (256p und 512p) erzielt, während er bestehende Auto-Encoder bei der 16-Bild-128p-Video-Rekonstruktion für UCF-101 übertrifft, und das mit 2-5x weniger FLOPs. Wenn ViTok mit Diffusion-Transformern integriert wird, zeigt es eine wettbewerbsfähige Leistung bei der Bildgenerierung für ImageNet-1K und setzt neue Bestmarken für die klassenkonditionierte Videogenerierung auf 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.