Уроки масштабирования визуальных токенизаторов для реконструкции и генерации

Learnings from Scaling Visual Tokenizers for Reconstruction and Generation

January 16, 2025

Авторы: Philippe Hansen-Estruch, David Yan, Ching-Yao Chung, Orr Zohar, Jialiang Wang, Tingbo Hou, Tao Xu, Sriram Vishwanath, Peter Vajda, Xinlei Chen

cs.AI

Аннотация

Визуальная токенизация с использованием автоэнкодера усиливает современные генеративные модели изображений и видео путем сжатия пикселей в скрытое пространство. Хотя масштабирование генераторов на основе трансформеров было ключевым в последних достижениях, компонент токенизатора сам по себе редко масштабируется, что оставляет открытыми вопросы о том, как выбор дизайна автоэнкодера влияет как на его цель восстановления, так и на последующую генеративную производительность. Наша работа направлена на исследование масштабирования в автоэнкодерах, чтобы заполнить этот пробел. Для облегчения этого исследования мы заменяем типичную сверточную основу на улучшенную архитектуру Vision Transformer для токенизации (ViTok). Мы обучаем ViTok на крупномасштабных наборах данных изображений и видео, значительно превышающих ImageNet-1K, устраняя ограничения данных на масштабирование токенизатора. Сначала мы изучаем, как масштабирование узкого места автоэнкодера влияет как на восстановление, так и на генерацию, и обнаруживаем, что хотя оно сильно коррелирует с восстановлением, его отношение к генерации более сложно. Затем мы исследовали эффект отдельного масштабирования кодера и декодера автоэнкодеров на производительность восстановления и генерации. Критически важно, что мы обнаружили, что масштабирование кодера приносит минимальные выгоды как для восстановления, так и для генерации, в то время как масштабирование декодера улучшает восстановление, но выгоды для генерации неоднозначны. Основываясь на нашем исследовании, мы разработали ViTok как легкий автоэнкодер, который достигает конкурентоспособной производительности с передовыми автоэнкодерами на задачах восстановления ImageNet-1K и COCO (256p и 512p), превосходя существующие автоэнкодеры на восстановлении видео из 16 кадров 128p для UCF-101, все это с 2-5 раз меньшим количеством операций с плавающей запятой (FLOPs). При интеграции с Диффузионными Трансформерами, ViTok демонстрирует конкурентоспособную производительность в генерации изображений для ImageNet-1K и устанавливает новые передовые стандарты для классовой генерации видео на 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.